diff --git a/tools/run_confidence_validation25_reference_single.sh b/tools/run_confidence_validation25_reference_single.sh new file mode 100644 index 0000000000000000000000000000000000000000..8eac98bab8dbbb3a10dfc7859f179de224717e0b --- /dev/null +++ b/tools/run_confidence_validation25_reference_single.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "$#" -ne 5 ]]; then + echo "Usage: $0 EVAL_NAME CONFIDENCE_CHECKPOINT REPLACE_STEPS THRESHOLD_MODE THRESHOLD" >&2 + echo "THRESHOLD_MODE must be global or per_step." >&2 + exit 2 +fi + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +eval_name="$1" +confidence_checkpoint="$2" +replace_steps="$3" +threshold_mode="$4" +threshold="$5" + +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_july20260720/bin/python}" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-${repo_root}/datasets/${dataset_name}}" +validation_root="${dataset_root}/validation" +eval_root="${dataset_root}/evaluations/${eval_name}" +log_dir="${eval_root}/logs" +base_model="${PREDICTOR_BASE_MODEL:-${repo_root}/outputs/hf_snapshot_repaired_9b49404}" +action_ckpt="${PREDICTOR_ACTION_CKPT:-/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub/models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/diffusion_pytorch_model.safetensors}" +weights="${repo_root}/checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" + +if [[ ! "${eval_name}" =~ ^[A-Za-z0-9._-]+$ ]]; then + echo "Invalid eval name: ${eval_name}" >&2 + exit 2 +fi +if [[ "${replace_steps}" != "1,2" && "${replace_steps}" != "1,2,3" ]]; then + echo "replace_steps must be 1,2 or 1,2,3" >&2 + exit 2 +fi +if [[ "${threshold_mode}" != "global" && "${threshold_mode}" != "per_step" ]]; then + echo "threshold mode must be global or per_step" >&2 + exit 2 +fi + +[[ -x "${python_bin}" ]] +[[ -d "${base_model}" ]] +[[ -f "${action_ckpt}" ]] +[[ -f "${weights}" ]] +[[ -f "${confidence_checkpoint}" ]] +[[ -f "${validation_root}/cases.jsonl" ]] +[[ $(find "${validation_root}/input_images" -type f -name '*.jpg' | wc -l) -eq 25 ]] +[[ $(find "${validation_root}/full" -type f -name '*.mp4' | wc -l) -eq 100 ]] + +mkdir -p "${log_dir}" +export PYTHONPATH="${repo_root}" +export HY_PROFILE_TIMING=1 + +confidence_args=(--confidence_checkpoint "${confidence_checkpoint}") +if [[ "${threshold_mode}" == "global" ]]; then + confidence_args+=(--confidence_threshold "${threshold}") +else + confidence_args+=(--confidence_thresholds "${threshold}") +fi + +run_workers() { + local phase="$1" + shift + local -a pids=() + local worker_id pid failed=0 + for worker_id in $(seq 0 7); do + CUDA_VISIBLE_DEVICES="${worker_id}" "${python_bin}" -u "$@" \ + --worker_id "${worker_id}" --num_workers 8 \ + >"${log_dir}/${phase}_worker_${worker_id}.log" 2>&1 & + pid="$!" + pids+=("${pid}") + echo "${pid}" >"${log_dir}/${phase}_worker_${worker_id}.pid" + done + for pid in "${pids[@]}"; do + if ! wait "${pid}"; then failed=1; fi + done + if [[ "${failed}" -ne 0 ]]; then + echo "${phase} worker failure; inspect ${log_dir}/${phase}_worker_*.log" >&2 + return 1 + fi +} + +if [[ ! -f "${eval_root}/generation.completed" ]] \ + || [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' 2>/dev/null | wc -l) -ne 100 ]]; then + date -u +%FT%TZ >"${eval_root}/generation.started" + run_workers generate \ + "${repo_root}/tools/generate_predictor_v4_train_case_eval.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --base_model "${base_model}" \ + --action_ckpt "${action_ckpt}" \ + --weights "${weights}" \ + --source_blocks 1,52 \ + --predictor_variant v4 \ + --previous_condition_mode same_step \ + --replace_steps "${replace_steps}" \ + --modes predictor \ + "${confidence_args[@]}" + [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' | wc -l) -eq 100 ]] + date -u +%FT%TZ >"${eval_root}/generation.completed" +else + echo "[skip] generation already complete: ${eval_name}" +fi + +metric_records=0 +if [[ -f "${eval_root}/metrics.json" ]]; then + metric_records=$("${python_bin}" -c \ + 'import json,sys; print(len(json.load(open(sys.argv[1]))["records"]))' \ + "${eval_root}/metrics.json") +fi +if [[ ! -f "${eval_root}/metrics.completed" || "${metric_records}" -ne 100 ]]; then + date -u +%FT%TZ >"${eval_root}/metrics.started" + run_workers metrics \ + "${repo_root}/tools/compute_predictor_v4_train_case_metrics.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --num_cases 25 \ + --reference_subdir full \ + --modes predictor \ + --metric_batch_size 4 + "${python_bin}" "${repo_root}/tools/merge_predictor_v4_train_case_metrics.py" \ + --output_root "${eval_root}" \ + --num_workers 8 \ + --num_cases 25 \ + --modes predictor \ + --replace_steps "${replace_steps}" \ + --reference_label "Full-DiT validation videos" \ + >"${log_dir}/metrics_merge.log" 2>&1 + [[ $("${python_bin}" -c \ + 'import json,sys; print(len(json.load(open(sys.argv[1]))["records"]))' \ + "${eval_root}/metrics.json") -eq 100 ]] + date -u +%FT%TZ >"${eval_root}/metrics.completed" +else + echo "[skip] metrics already complete: ${eval_name}" +fi + +date -u +%FT%TZ >"${eval_root}/sync.completed" +if [[ "${PREDICTOR_SKIP_SYNC:-1}" != "1" ]]; then + bucket="${BUCKET:-s3-us-west2-default}" + durable_eval_uri="s3://${bucket}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/evaluations/${eval_name}" + aws s3 sync "${eval_root}" "${durable_eval_uri}" --only-show-errors +fi + +echo "Validation evaluation complete: ${eval_root}" diff --git a/tools/run_offline_ffff_targets_validation25_evaluation.sh b/tools/run_offline_ffff_targets_validation25_evaluation.sh new file mode 100644 index 0000000000000000000000000000000000000000..fe40277576e21a9423e3d6b4f909b60c031696aa --- /dev/null +++ b/tools/run_offline_ffff_targets_validation25_evaluation.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +vbench_repo="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +weights="${repo_root}/checkpoints/predictor_v4_detached-chunk-timestep-rollout-offline-ffff-targets-fppp_blocks1-52_fulltrain100_seed20260721_fp32master_8gpu_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1_swanlab/predictor_step_02000.safetensors" +eval_prefix="eval_validation25_detached-chunk-timestep-rollout-offline-ffff-targets-fppp-fulltrain100" +eval_suffix="blocks1-52_ckptstep2000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1" +fppf_name="${eval_prefix}_schedule-fppf_${eval_suffix}" +fppp_name="${eval_prefix}_schedule-fppp_${eval_suffix}" +run_name="validation25_offline_ffff_targets_step2000" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" + +export PYTHONPATH="${repo_root}" + +bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${fppf_name}" "${weights}" "1,52" v4 "1,2" +bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${fppp_name}" "${weights}" "1,52" v4 "1,2,3" + +export PYTHONPATH="${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/vbench" +export TORCH_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/torch" +export HF_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/huggingface" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +group_fppf="offline_ffff_targets_step2000_fppf" +group_fppp="offline_ffff_targets_step2000_fppp" +declare -a groups=( + "${group_fppf}|${dataset_root}/evaluations/${fppf_name}/predictor" + "${group_fppp}|${dataset_root}/evaluations/${fppp_name}/predictor" +) + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne 200 ]]; then + echo "Expected 200 staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 200 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json +import sys + +path, dimension = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == 200, (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + --group "${group_fppf}" \ + --group "${group_fppp}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +echo "[validation] all standard and VBench metrics complete: ${run_root}" diff --git a/tools/run_oldenv_validation_reruns_sequential.sh b/tools/run_oldenv_validation_reruns_sequential.sh new file mode 100644 index 0000000000000000000000000000000000000000..2e6347dfd08794f121b14f2cccdc0101566f5b24 --- /dev/null +++ b/tools/run_oldenv_validation_reruns_sequential.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_july20260720/bin/python}" +base_model="${PREDICTOR_BASE_MODEL:-${repo_root}/outputs/hf_snapshot_repaired_9b49404}" +run_tag="${OLDENV_REVAL_TAG:-oldenv_reval_20260821}" +run_root="${repo_root}/logs/${run_tag}" +master_log="${run_root}/sequential.log" +status_file="${run_root}/status.tsv" + +mkdir -p "${run_root}" +exec > >(tee -a "${master_log}") 2>&1 + +export PREDICTOR_PYTHON_BIN="${python_bin}" +export PREDICTOR_BASE_MODEL="${base_model}" +export PREDICTOR_SKIP_SYNC="${PREDICTOR_SKIP_SYNC:-0}" +export PYTHONPATH="${repo_root}" + +on_error() { + local exit_code="$?" + printf 'failed\t%s\tline=%s\texit=%s\n' "$(date -u +%FT%TZ)" "${BASH_LINENO[0]}" "${exit_code}" >>"${status_file}" + exit "${exit_code}" +} +trap on_error ERR + +record_status() { + printf '%s\t%s\t%s\n' "$1" "$(date -u +%FT%TZ)" "$2" >>"${status_file}" +} + +run_confidence_eval() { + local short_name="$1" checkpoint="$2" replace_steps="$3" threshold_mode="$4" threshold="$5" + local eval_name="eval_validation25_${short_name}_${run_tag}" + record_status start "${eval_name}" + bash "${repo_root}/tools/run_confidence_validation25_reference_single.sh" \ + "${eval_name}" "${checkpoint}" "${replace_steps}" "${threshold_mode}" "${threshold}" + record_status complete "${eval_name}" +} + +echo "[preflight] fixed environment" +"${python_bin}" - <<'PY' +import importlib.util +import torch +assert torch.__version__ == "2.6.0+cu124", torch.__version__ +assert importlib.util.find_spec("flash_attn") is None, "flash_attn must be absent" +print(f"torch={torch.__version__} cuda={torch.version.cuda} flash_attn=False") +PY +[[ -d "${base_model}" ]] +[[ $(find "${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50/validation/full" -type f -name '*.mp4' | wc -l) -eq 100 ]] +record_status preflight "torch=2.6.0+cu124; flash_attn=False" + +small_c1="${repo_root}/confidence_validation/checkpoints/confidence_c1.pt" +wide_c1="${repo_root}/confidence_validation_wide4096x2_c1/checkpoints/confidence_c1.pt" +wide_c2="${repo_root}/confidence_validation_wide4096x2_c2/checkpoints/confidence_c2.pt" + +# Keep the trained thresholds unchanged. The fixed-environment run records the +# actual fallback count, so distribution drift remains visible in the result. +run_confidence_eval stage1-c1-threshold-r10 "${small_c1}" 1,2,3 global -3.523588657 +run_confidence_eval stage1-c1-threshold-r25 "${small_c1}" 1,2,3 global -3.915276289 +run_confidence_eval stage1-c1-threshold-r33 "${small_c1}" 1,2,3 global -4.329442501 +run_confidence_eval stage1-c1-threshold-r50 "${small_c1}" 1,2,3 global -4.798926830 +run_confidence_eval stage1-wide4096x2-c1-threshold-r33 "${wide_c1}" 1,2,3 global -4.308430990 +run_confidence_eval stage1-wide4096x2-c2-hybrid-p3full-p12r10 "${wide_c2}" 1,2 per_step 1:-4.136674309,2:-3.798280764 +run_confidence_eval stage1-wide4096x2-c2-hybrid-p3full-p12r25 "${wide_c2}" 1,2 per_step 1:-4.400576830,2:-4.049407601 + +record_status complete "all validation reruns" +date -u +%FT%TZ >"${run_root}/all.completed" + +if [[ "${PREDICTOR_SKIP_SYNC}" != "1" ]]; then + bucket="${BUCKET:-s3-us-west2-default}" + aws s3 sync "${run_root}" \ + "s3://${bucket}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/logs/${run_tag}" \ + --only-show-errors +fi + +echo "[complete] all fixed-environment validation reruns" diff --git a/tools/run_predictor_confidence_experiment.sh b/tools/run_predictor_confidence_experiment.sh new file mode 100644 index 0000000000000000000000000000000000000000..a1c2d9d20e176c215261030c3e05d5f06ecfc3fe --- /dev/null +++ b/tools/run_predictor_confidence_experiment.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT" +PYTHON_BIN="${PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_runtime/bin/python}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +NPROC="${NPROC:-8}" +TRAIN_ROOT="${TRAIN_ROOT:-$ROOT/datasets/hyworldplay_predictor_prefeature_v4}" +VAL_ROOT="${VAL_ROOT:-$ROOT/datasets/hyworldplay_predictor_confidence_validation25_prefeature}" +FEATURE_ROOT="${FEATURE_ROOT:-$ROOT/confidence_validation/features}" +OUTPUT_ROOT="${OUTPUT_ROOT:-$ROOT/confidence_validation}" +PREDICTOR="${PREDICTOR:-$ROOT/checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors}" +BASE_MODEL="${BASE_MODEL:-$ROOT/outputs/model_cache_repaired/HunyuanVideo-1.5-9b49404-v2}" +TEACHER="${TEACHER:-/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub/models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/diffusion_pytorch_model.safetensors}" +export PYTHONPATH="$ROOT" + +if [[ ! -f "$VAL_ROOT/train_eval_manifest.jsonl" ]]; then + "$PYTHON_BIN" tools/prepare_confidence_validation_prefeature.py --output_root "$VAL_ROOT" + mkdir -p "$VAL_ROOT/logs" + pids=() + for worker in $(seq 0 $((NPROC - 1))); do + CUDA_VISIBLE_DEVICES="${worker}" "$PYTHON_BIN" -u tools/build_predictor_prefeature_dataset.py \ + --dataset_root "$VAL_ROOT" --base_model "$BASE_MODEL" --action_ckpt "$TEACHER" \ + --worker_id "$worker" --num_workers "$NPROC" \ + >"$VAL_ROOT/logs/worker_${worker}.log" 2>&1 & + pids+=("$!") + done + failed=0 + for worker in "${!pids[@]}"; do + wait "${pids[$worker]}" || { echo "validation trajectory worker $worker failed" >&2; failed=1; } + done + (( failed == 0 )) || exit 1 + "$PYTHON_BIN" tools/merge_predictor_prefeature_manifests.py \ + --dataset_root "$VAL_ROOT" --expected_cases 25 --expected_actions 4 --expected_chunks 8 +fi + +extract() { + local split="$1" manifest="$2" output="$FEATURE_ROOT/$1" + mkdir -p "$output" + if [[ "$(find "$output" -maxdepth 1 -name 'features_rank_*.pt' | wc -l)" -ne "$NPROC" ]]; then + rm -f "$output"/features_rank_*.pt + CUDA_VISIBLE_DEVICES="$GPU_IDS" "$PYTHON_BIN" -m torch.distributed.run \ + --standalone --nproc_per_node="$NPROC" tools/extract_predictor_confidence_features.py \ + --manifest "$manifest" --predictor_checkpoint "$PREDICTOR" \ + --teacher_checkpoint "$TEACHER" --output_dir "$output" --batch_size 8 + fi +} + +extract train "$TRAIN_ROOT/train_eval_manifest.jsonl" +extract validation "$VAL_ROOT/train_eval_manifest.jsonl" +CUDA_VISIBLE_DEVICES="$GPU_IDS" "$PYTHON_BIN" -m torch.distributed.run \ + --standalone --nproc_per_node="$NPROC" tools/train_evaluate_predictor_confidence.py \ + --train_features "$FEATURE_ROOT/train" --validation_features "$FEATURE_ROOT/validation" \ + --output_dir "$OUTPUT_ROOT" --epochs 40 --batch_size 256 --seed 0 diff --git a/tools/run_predictor_v4_atc_full_experiment.sh b/tools/run_predictor_v4_atc_full_experiment.sh new file mode 100644 index 0000000000000000000000000000000000000000..6d2119fde6639bfda5cd5c6806631be3c7e74f91 --- /dev/null +++ b/tools/run_predictor_v4_atc_full_experiment.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +max_steps="${ATC_MAX_STEPS:-2000}" +variants=(full anchor_only raw_prev no_gate no_distance) + +for variant in "${variants[@]}"; do + bash "${repo_root}/tools/launch_predictor_v4_atc_stage1_8gpu.sh" \ + "${variant}" "${max_steps}" +done + +for variant in "${variants[@]}"; do + checkpoint_dir="${repo_root}/checkpoints/predictor_v4_atc-${variant}_blocks1-52_fulltrain100_fp32master_8gpu_flr1e-4_blr1e-5_steps${max_steps}" + weights="${checkpoint_dir}/predictor_step_$(printf '%05d' "${max_steps}").safetensors" + for schedule in fppf fppp; do + replace_steps="1,2" + [[ "${schedule}" == "fppp" ]] && replace_steps="1,2,3" + eval_name="eval_validation25_atc-${variant}_schedule-${schedule}_blocks1-52_step${max_steps}_oldenv" + bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${eval_name}" "${weights}" 1,52 v4 "${replace_steps}" \ + same_step concat atc "${variant}" + done +done diff --git a/tools/run_predictor_v4_atc_global_sequential_8gpu.sh b/tools/run_predictor_v4_atc_global_sequential_8gpu.sh new file mode 100644 index 0000000000000000000000000000000000000000..96b34eaf87617db309966883073b6a431cb230c3 --- /dev/null +++ b/tools/run_predictor_v4_atc_global_sequential_8gpu.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Train the global-transport ATC variants one after another on 8 GPUs: +# per-rank batch 8, no gradient accumulation (global batch 64), 2000 steps each. +set -euo pipefail +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +max_steps="${ATC_MAX_STEPS:-2000}" +scopes=("${@:-}") +[[ -n "${scopes[0]}" ]] || scopes=(chunk last_frame) +for scope in "${scopes[@]}"; do + echo "[$(date -Is)] start scope=${scope}" + PREDICTOR_RDZV_PORT=29501 bash "${repo_root}/tools/launch_predictor_v4_atc_global_stage1.sh" \ + "${scope}" 0,1,2,3,4,5,6,7 8 1 "${max_steps}" + echo "[$(date -Is)] done scope=${scope}" +done diff --git a/tools/run_predictor_v4_detached_timestep_rollout_offline_targets_overfit8.sh b/tools/run_predictor_v4_detached_timestep_rollout_offline_targets_overfit8.sh new file mode 100644 index 0000000000000000000000000000000000000000..7ab1815eb95a1337b3ff046efed5205d33ba5502 --- /dev/null +++ b/tools/run_predictor_v4_detached_timestep_rollout_offline_targets_overfit8.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +config_tag="flr1e-4_blr1e-5_schedsteps1000_swarmup8_bwarmup8" +run_name="detached-timestep-rollout-offline-ffff-targets_blocks1-52_overfit8_seed20260721_fp32master_8gpu_${config_tag}" +checkpoint_dir="${repo_root}/checkpoints/predictor_v4_${run_name}" +eval_prefix="detached-timestep-rollout-offline-ffff-targets-overfit8" +summary_name="summary_${eval_prefix}_ckptstep1000_${config_tag}" +summary_root="/mnt/local_nvme/zoubin/cz/hyworldplay_${summary_name}" +final_weights="${checkpoint_dir}/predictor_step_01000.safetensors" +mkdir -p "${checkpoint_dir}" "${summary_root}" + +export BUCKET="${BUCKET:-s3-us-west2-default}" +export FS_ID="${FS_ID:-fs-05c60be7352499179}" +export MOUNT_POINT="${MOUNT_POINT:-/mnt/s3files/s3-us-west2-default}" +export PYTHONPATH="${repo_root}" + +exec 9>"${summary_root}/pipeline.lock" +if ! flock -n 9; then + echo "[pipeline] detached offline-target rollout pipeline is already active" + exit 0 +fi + +run_with_retries() { + local stage="$1" + shift + local attempt=1 + while true; do + echo "[pipeline] stage=${stage} attempt=${attempt} start=$(date -u +%FT%TZ)" + if "$@"; then + echo "[pipeline] stage=${stage} complete=$(date -u +%FT%TZ)" + return 0 + fi + if [[ "${attempt}" -ge 5 ]]; then + echo "[pipeline] stage=${stage} failed after ${attempt} attempts" >&2 + return 1 + fi + attempt=$((attempt + 1)) + sleep 30 + done +} + +train_model() { + if [[ -f "${final_weights}" ]]; then + echo "[pipeline] step 1000 already exists" + return 0 + fi + local resume_path="" + if [[ -f "${checkpoint_dir}/training_latest.pt" ]]; then + resume_path="${checkpoint_dir}/training_latest.pt" + fi + bash "${repo_root}/tools/launch_predictor_v4_detached_timestep_rollout_offline_targets_overfit8_training.sh" \ + "${checkpoint_dir}" "predictor-v4-${run_name//_/-}" "${resume_path}" 1000 + [[ -f "${final_weights}" ]] +} + +evaluate_model() { + local schedule="$1" + local replace_steps="$2" + local eval_root="/mnt/local_nvme/zoubin/cz/hyworldplay_eval_v4_${eval_prefix}_schedule-${schedule}_blocks1-52_ckptstep1000_${config_tag}" + local durable_name="eval_${eval_prefix}_schedule-${schedule}_blocks1-52_ckptstep1000_${config_tag}" + if [[ -f "${eval_root}/metrics.completed" && -f "${eval_root}/sync.completed" ]]; then + echo "[pipeline] ${schedule} step-1000 evaluation already complete" + return 0 + fi + bash "${repo_root}/tools/launch_predictor_v4_rollout_overfit8_eval.sh" \ + "${eval_root}" "${final_weights}" "${durable_name}" "${replace_steps}" + [[ -f "${eval_root}/metrics.json" && -f "${eval_root}/sync.completed" ]] +} + +compare_to() { + local baseline_label="$1" + local baseline_fppf="$2" + local baseline_fppp="$3" + local output_name="$4" + local candidate_prefix="/mnt/local_nvme/zoubin/cz/hyworldplay_eval_v4_${eval_prefix}" + "${python_bin}" "${repo_root}/tools/compare_predictor_overfit8_metrics.py" \ + --baseline_label "${baseline_label}" \ + --baseline_fppf "${baseline_fppf}" \ + --baseline_fppp "${baseline_fppp}" \ + --candidate_label "detached-timestep-rollout-offline-ffff-targets-step1000" \ + --candidate_fppf "${candidate_prefix}_schedule-fppf_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + --candidate_fppp "${candidate_prefix}_schedule-fppp_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + --output "${summary_root}/${output_name}.json" \ + >"${summary_root}/${output_name}.log" 2>&1 +} + +run_with_retries dataset \ + bash "${repo_root}/tools/ensure_predictor_prefeature_nvme.sh" +run_with_retries train_step1000 train_model +run_with_retries eval_step1000_fppf evaluate_model fppf "1,2" +run_with_retries eval_step1000_fppp evaluate_model fppp "1,2,3" + +run_with_retries compare_stage1 compare_to \ + "stage1-step2000" \ + "${repo_root}/outputs/predictor_v4/eval_stage1_overfit8_schedule-fppf_blocks1-52_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/metrics.json" \ + "${repo_root}/outputs/predictor_v4/eval_stage1_overfit8_schedule-fppp_blocks1-52_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/metrics.json" \ + comparison_vs_stage1 +run_with_retries compare_grouped_offline compare_to \ + "grouped-teacher-forcing-offline-prefeature-step1000" \ + "${repo_root}/outputs/predictor_v4/eval_grouped-teacher-forcing-offline-prefeature-overfit8_schedule-fppf_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + "${repo_root}/outputs/predictor_v4/eval_grouped-teacher-forcing-offline-prefeature-overfit8_schedule-fppp_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + comparison_vs_grouped_offline +run_with_retries compare_detached_online compare_to \ + "detached-timestep-only-online-same-state-step1000" \ + "${repo_root}/outputs/predictor_v4/eval_timestep-rollout-detached-timestepfp32-overfit8-pure_schedule-fppf_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + "${repo_root}/outputs/predictor_v4/eval_timestep-rollout-detached-timestepfp32-overfit8-pure_schedule-fppp_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + comparison_vs_detached_online_same_state + +date -u +%FT%TZ >"${summary_root}/pipeline.completed" +aws s3 sync "${summary_root}" \ + "s3://${BUCKET}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/outputs/predictor_v4/${summary_name}" \ + --only-show-errors +echo "[pipeline] training, final evaluations, and three comparisons complete" diff --git a/tools/run_predictor_v4_grouped_teacher_forcing_overfit8.sh b/tools/run_predictor_v4_grouped_teacher_forcing_overfit8.sh new file mode 100644 index 0000000000000000000000000000000000000000..3b5923a00b8bd76886bf2f226449ceb5a27eb74a --- /dev/null +++ b/tools/run_predictor_v4_grouped_teacher_forcing_overfit8.sh @@ -0,0 +1,99 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +config_tag="flr1e-4_blr1e-5_schedsteps1000_swarmup8_bwarmup8" +run_name="grouped-teacher-forcing-offline-prefeature_blocks1-52_overfit8_seed20260721_fp32master_8gpu_${config_tag}" +checkpoint_dir="${repo_root}/checkpoints/predictor_v4_${run_name}" +eval_prefix="grouped-teacher-forcing-offline-prefeature-overfit8" +summary_name="summary_grouped-teacher-forcing-offline-prefeature-vs-detached-timestepfp32-overfit8_ckptstep1000_${config_tag}" +summary_root="/mnt/local_nvme/zoubin/cz/hyworldplay_${summary_name}" +final_weights="${checkpoint_dir}/predictor_step_01000.safetensors" +mkdir -p "${checkpoint_dir}" "${summary_root}" + +export BUCKET="${BUCKET:-s3-us-west2-default}" +export FS_ID="${FS_ID:-fs-05c60be7352499179}" +export MOUNT_POINT="${MOUNT_POINT:-/mnt/s3files/s3-us-west2-default}" +export PYTHONPATH="${repo_root}" + +exec 9>"${summary_root}/pipeline.lock" +if ! flock -n 9; then + echo "[pipeline] grouped teacher-forcing pipeline is already active" + exit 0 +fi + +run_with_retries() { + local stage="$1" + shift + local attempt=1 + while true; do + echo "[pipeline] stage=${stage} attempt=${attempt} start=$(date -u +%FT%TZ)" + if "$@"; then + echo "[pipeline] stage=${stage} complete=$(date -u +%FT%TZ)" + return 0 + fi + if [[ "${attempt}" -ge 5 ]]; then + echo "[pipeline] stage=${stage} failed after ${attempt} attempts" >&2 + return 1 + fi + attempt=$((attempt + 1)) + sleep 30 + done +} + +train_grouped() { + if [[ -f "${final_weights}" ]]; then + echo "[pipeline] grouped teacher-forcing step 1000 already exists" + return 0 + fi + local resume_path="" + if [[ -f "${checkpoint_dir}/training_latest.pt" ]]; then + resume_path="${checkpoint_dir}/training_latest.pt" + fi + bash "${repo_root}/tools/launch_predictor_v4_grouped_teacher_forcing_overfit8_training.sh" \ + "${checkpoint_dir}" \ + "predictor-v4-${run_name//_/-}" \ + "${resume_path}" 1000 + [[ -f "${final_weights}" ]] +} + +evaluate_grouped() { + local schedule="$1" + local replace_steps="$2" + local eval_root="/mnt/local_nvme/zoubin/cz/hyworldplay_eval_v4_${eval_prefix}_schedule-${schedule}_blocks1-52_ckptstep1000_${config_tag}" + local durable_name="eval_${eval_prefix}_schedule-${schedule}_blocks1-52_ckptstep1000_${config_tag}" + if [[ -f "${eval_root}/metrics.completed" && -f "${eval_root}/sync.completed" ]]; then + echo "[pipeline] ${schedule} step-1000 evaluation already complete" + return 0 + fi + bash "${repo_root}/tools/launch_predictor_v4_rollout_overfit8_eval.sh" \ + "${eval_root}" "${final_weights}" "${durable_name}" "${replace_steps}" + [[ -f "${eval_root}/metrics.json" && -f "${eval_root}/sync.completed" ]] +} + +run_with_retries dataset \ + bash "${repo_root}/tools/ensure_predictor_prefeature_nvme.sh" +run_with_retries train_step1000 train_grouped + +# Only the final step-1000 checkpoint is evaluated. +run_with_retries eval_step1000_fppf evaluate_grouped fppf "1,2" +run_with_retries eval_step1000_fppp evaluate_grouped fppp "1,2,3" + +baseline_prefix="${repo_root}/outputs/predictor_v4/eval_timestep-rollout-detached-timestepfp32-overfit8-pure" +baseline_suffix="blocks1-52_ckptstep1000_${config_tag}/metrics.json" +candidate_prefix="/mnt/local_nvme/zoubin/cz/hyworldplay_eval_v4_${eval_prefix}" +"${python_bin}" "${repo_root}/tools/compare_predictor_overfit8_metrics.py" \ + --baseline_label "detached-timestep-only-step1000" \ + --baseline_fppf "${baseline_prefix}_schedule-fppf_${baseline_suffix}" \ + --baseline_fppp "${baseline_prefix}_schedule-fppp_${baseline_suffix}" \ + --candidate_label "grouped-teacher-forcing-offline-prefeature-step1000" \ + --candidate_fppf "${candidate_prefix}_schedule-fppf_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + --candidate_fppp "${candidate_prefix}_schedule-fppp_blocks1-52_ckptstep1000_${config_tag}/metrics.json" \ + --output "${summary_root}/comparison.json" \ + >"${summary_root}/comparison.log" 2>&1 +date -u +%FT%TZ >"${summary_root}/pipeline.completed" +aws s3 sync "${summary_root}" \ + "s3://${BUCKET}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/outputs/predictor_v4/${summary_name}" \ + --only-show-errors +echo "[pipeline] grouped teacher forcing, final evaluations, and comparison complete" diff --git a/tools/run_predictor_v4_timestep_rollout_detached_overfit8_experiments.sh b/tools/run_predictor_v4_timestep_rollout_detached_overfit8_experiments.sh new file mode 100644 index 0000000000000000000000000000000000000000..8d7385cc70dd875df584c331018a797d948c3c0f --- /dev/null +++ b/tools/run_predictor_v4_timestep_rollout_detached_overfit8_experiments.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +export TIMESTEP_ROLLOUT_SUITE_NAME="timestep_rollout_detached_timestepfp32_blocks1-52_overfit8_seed20260721_fp32master_8gpu" +export TIMESTEP_ROLLOUT_EVAL_SUITE_NAME="timestep-rollout-detached-timestepfp32-overfit8" +export TIMESTEP_ROLLOUT_DETACH_BETWEEN_TIMESTEPS="true" +export TIMESTEP_ROLLOUT_VARIANTS="pure:0.0" +export TIMESTEP_ROLLOUT_MAX_STEPS="1000" + +exec bash "${repo_root}/tools/run_predictor_v4_timestep_rollout_overfit8_experiments.sh" diff --git a/tools/run_predictor_v4_timestep_rollout_overfit8_experiments.sh b/tools/run_predictor_v4_timestep_rollout_overfit8_experiments.sh new file mode 100644 index 0000000000000000000000000000000000000000..5cba0453ce7f9f1e0ecaa593d9abf07af2b46342 --- /dev/null +++ b/tools/run_predictor_v4_timestep_rollout_overfit8_experiments.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +suite_name="${TIMESTEP_ROLLOUT_SUITE_NAME:-timestep_rollout_timestepfp32_blocks1-52_overfit8_seed20260721_fp32master_8gpu}" +eval_suite_name="${TIMESTEP_ROLLOUT_EVAL_SUITE_NAME:-timestep-rollout-timestepfp32-overfit8}" +detach_between_timesteps="${TIMESTEP_ROLLOUT_DETACH_BETWEEN_TIMESTEPS:-false}" +variant_specs="${TIMESTEP_ROLLOUT_VARIANTS:-replay0:0.0 replay0p2:0.2 replay0p4:0.4}" +max_steps="${TIMESTEP_ROLLOUT_MAX_STEPS:-200}" +stop_after_steps="${TIMESTEP_ROLLOUT_STOP_AFTER_STEPS:-${max_steps}}" +checkpoint_parent="${repo_root}/checkpoints/predictor_v4_${suite_name}" +config_tag="flr1e-4_blr1e-5_schedsteps${max_steps}_swarmup8_bwarmup8" +checkpoint_parent="${checkpoint_parent}_${config_tag}" +eval_parent="/mnt/local_nvme/zoubin/cz/hyworldplay_eval_v4_${eval_suite_name}_blocks1-52_ckptstep${stop_after_steps}_${config_tag}" +pipeline_root="/mnt/local_nvme/zoubin/cz/hyworldplay_summary_v4_${eval_suite_name}_ckptstep${stop_after_steps}_${config_tag}" +mkdir -p "${checkpoint_parent}" "${eval_parent}" "${pipeline_root}" + +export BUCKET="${BUCKET:-s3-us-west2-default}" +export FS_ID="${FS_ID:-fs-05c60be7352499179}" +export MOUNT_POINT="${MOUNT_POINT:-/mnt/s3files/s3-us-west2-default}" +export PYTHONPATH="${repo_root}" + +exec 9>"${pipeline_root}/pipeline.lock" +if ! flock -n 9; then + echo "[pipeline] another timestep-rollout sweep is already active" + exit 0 +fi + +run_with_retries() { + local stage="$1" + shift + local attempt=1 + while true; do + echo "[pipeline] stage=${stage} attempt=${attempt} start=$(date -u +%FT%TZ)" + if "$@"; then + echo "[pipeline] stage=${stage} complete=$(date -u +%FT%TZ)" + return 0 + fi + if [[ "${attempt}" -ge 5 ]]; then + echo "[pipeline] stage=${stage} failed after ${attempt} attempts" >&2 + return 1 + fi + attempt=$((attempt + 1)) + sleep 30 + done +} + +train_variant() { + local label="$1" + local replay_weight="$2" + local output_dir="${checkpoint_parent}/${label}" + local final_step + printf -v final_step "%05d" "${stop_after_steps}" + local final_weights="${output_dir}/predictor_step_${final_step}.safetensors" + if [[ -f "${final_weights}" ]]; then + echo "[pipeline] training ${label} already complete" + return 0 + fi + local resume_path="" + if [[ -f "${output_dir}/training_latest.pt" ]]; then + resume_path="${output_dir}/training_latest.pt" + fi + bash "${repo_root}/tools/launch_predictor_v4_timestep_rollout_overfit8_training.sh" \ + "${output_dir}" "${replay_weight}" \ + "predictor-v4-${suite_name}-${label}-${config_tag//_/-}" \ + "${resume_path}" "${detach_between_timesteps}" "${max_steps}" \ + "${stop_after_steps}" + [[ -f "${final_weights}" ]] +} + +evaluate_variant() { + local label="$1" + local schedule="$2" + local replace_steps="$3" + local final_step + printf -v final_step "%05d" "${stop_after_steps}" + local weights="${checkpoint_parent}/${label}/predictor_step_${final_step}.safetensors" + local eval_root="${eval_parent}/${label}_${schedule}" + local durable_name="eval_${eval_suite_name}_${label}_schedule-${schedule}_blocks1-52_ckptstep${stop_after_steps}_${config_tag}" + if [[ -f "${eval_root}/metrics.completed" && -f "${eval_root}/sync.completed" ]]; then + echo "[pipeline] evaluation ${label}/${schedule} already complete" + return 0 + fi + bash "${repo_root}/tools/launch_predictor_v4_rollout_overfit8_eval.sh" \ + "${eval_root}" "${weights}" "${durable_name}" "${replace_steps}" + [[ -f "${eval_root}/metrics.json" && -f "${eval_root}/sync.completed" ]] +} + +run_with_retries sidecar \ + bash "${repo_root}/tools/ensure_predictor_timestep_rollout_overfit8_sidecar.sh" + +for spec in ${variant_specs}; do + label="${spec%%:*}" + replay_weight="${spec##*:}" + run_with_retries "train_${label}" train_variant "${label}" "${replay_weight}" +done + +for spec in ${variant_specs}; do + label="${spec%%:*}" + run_with_retries "eval_${label}_fppp" evaluate_variant "${label}" fppp "1,2,3" + run_with_retries "eval_${label}_fppf" evaluate_variant "${label}" fppf "1,2" +done + +summary_path="${pipeline_root}/summary.json" +"${python_bin}" "${repo_root}/tools/summarize_predictor_timestep_rollout_overfit8.py" \ + --eval_parent "${eval_parent}" --output "${summary_path}" \ + --experiment "v4 ${suite_name}; scheduler_max_steps=${max_steps}; stop_after_steps=${stop_after_steps}; detach_between_timesteps=${detach_between_timesteps}" \ + --variants "${variant_specs}" \ + >"${pipeline_root}/summary.log" 2>&1 +date -u +%FT%TZ >"${pipeline_root}/pipeline.completed" +aws s3 sync "${pipeline_root}" \ + "s3://${BUCKET}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/outputs/predictor_v4/summary_${eval_suite_name}_ckptstep${stop_after_steps}_${config_tag}" \ + --only-show-errors +echo "[pipeline] all training and evaluation stages complete" diff --git a/tools/run_predictor_v4_validation5_rollout_comparison.sh b/tools/run_predictor_v4_validation5_rollout_comparison.sh new file mode 100644 index 0000000000000000000000000000000000000000..17399d77874c3a74ab12ac8318ec95f755732263 --- /dev/null +++ b/tools/run_predictor_v4_validation5_rollout_comparison.sh @@ -0,0 +1,176 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +validation_root="${dataset_root}/validation" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +case_ids="0,8,15,19,21" +subset_tag="validation5_seed20260722_cases0-8-15-19-21" +config_tag="flr1e-4_blr1e-5_schedsteps1000_swarmup8_bwarmup8" +summary_root="${dataset_root}/evaluations/summary_${subset_tag}_three-stage2-variants" +pipeline_log_dir="${summary_root}/logs" +mkdir -p "${pipeline_log_dir}" + +export BUCKET="${BUCKET:-s3-us-west2-default}" +export FS_ID="${FS_ID:-fs-05c60be7352499179}" +export MOUNT_POINT="${MOUNT_POINT:-/mnt/s3files/s3-us-west2-default}" +export PYTHONPATH="${repo_root}" +export HY_PROFILE_TIMING=1 + +durable_dataset_uri="s3://${BUCKET}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}" +durable_eval_uri="${durable_dataset_uri}/evaluations" + +exec 9>"${summary_root}/pipeline.lock" +if ! flock -n 9; then + echo "[pipeline] validation5 comparison is already active" + exit 0 +fi + +if [[ ! -f "${validation_root}/cases.jsonl" ]] \ + || [[ $(find "${validation_root}/input_images" -type f -name '*.jpg' 2>/dev/null | wc -l) -ne 25 ]] \ + || [[ $(find "${validation_root}/full" -type f -name '*.mp4' 2>/dev/null | wc -l) -ne 100 ]]; then + echo "[dataset] restoring validation cache from ${durable_dataset_uri}" + mkdir -p "${dataset_root}" + aws s3 sync "${durable_dataset_uri}" "${dataset_root}" --only-show-errors +fi + +[[ -f "${validation_root}/cases.jsonl" ]] +[[ $(find "${validation_root}/input_images" -type f -name '*.jpg' | wc -l) -eq 25 ]] +[[ $(find "${validation_root}/full" -type f -name '*.mp4' | wc -l) -eq 100 ]] +cp "${repo_root}/configs/validation5_seed20260722.json" "${summary_root}/selection.json" + +declare -a variants=( + "online-same-state|${repo_root}/checkpoints/predictor_v4_timestep_rollout_detached_timestepfp32_blocks1-52_overfit8_seed20260721_fp32master_8gpu_${config_tag}/pure/predictor_step_01000.safetensors" + "grouped-offline-teacher-forcing|${repo_root}/checkpoints/predictor_v4_grouped-teacher-forcing-offline-prefeature_blocks1-52_overfit8_seed20260721_fp32master_8gpu_${config_tag}/predictor_step_01000.safetensors" + "offline-ffff-target-rollout|${repo_root}/checkpoints/predictor_v4_detached-timestep-rollout-offline-ffff-targets_blocks1-52_overfit8_seed20260721_fp32master_8gpu_${config_tag}/predictor_step_01000.safetensors" +) + +run_workers() { + local log_dir="$1" + local phase="$2" + shift 2 + local pids=() + local worker_id worker_pid failed=0 + mkdir -p "${log_dir}" + for worker_id in $(seq 0 7); do + CUDA_VISIBLE_DEVICES="${worker_id}" \ + "${python_bin}" -u "$@" \ + --worker_id "${worker_id}" --num_workers 8 \ + >"${log_dir}/${phase}_worker_${worker_id}.log" 2>&1 & + worker_pid=$! + pids+=("${worker_pid}") + echo "${worker_pid}" >"${log_dir}/${phase}_worker_${worker_id}.pid" + done + for worker_pid in "${pids[@]}"; do + if ! wait "${worker_pid}"; then + failed=1 + fi + done + if [[ "${failed}" -ne 0 ]]; then + echo "${phase} worker failure; inspect ${log_dir}/${phase}_worker_*.log" >&2 + return 1 + fi +} + +run_with_retries() { + local stage="$1" + shift + local attempt=1 + while true; do + echo "[pipeline] stage=${stage} attempt=${attempt} start=$(date -u +%FT%TZ)" + if "$@"; then + echo "[pipeline] stage=${stage} complete=$(date -u +%FT%TZ)" + return 0 + fi + if [[ "${attempt}" -ge 5 ]]; then + echo "[pipeline] stage=${stage} failed after ${attempt} attempts" >&2 + return 1 + fi + attempt=$((attempt + 1)) + sleep 30 + done +} + +evaluate_one() { + local variant="$1" + local weights="$2" + local schedule="$3" + local replace_steps="$4" + local eval_name="eval_${subset_tag}_${variant}_schedule-${schedule}_blocks1-52_ckptstep1000_${config_tag}" + local eval_root="${dataset_root}/evaluations/${eval_name}" + local log_dir="${eval_root}/logs" + mkdir -p "${log_dir}" + [[ -f "${weights}" ]] + + local video_count=0 + if [[ -d "${eval_root}/predictor" ]]; then + video_count=$(find "${eval_root}/predictor" -type f -name '*.mp4' | wc -l) + fi + if [[ ! -f "${eval_root}/generation.completed" || "${video_count}" -ne 20 ]]; then + date -u +%FT%TZ >"${eval_root}/generation.started" + run_workers "${log_dir}" generation \ + "${repo_root}/tools/generate_predictor_v4_train_case_eval.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --weights "${weights}" \ + --source_blocks 1,52 \ + --predictor_variant v4 \ + --case_ids "${case_ids}" \ + --modes predictor \ + --replace_steps "${replace_steps}" + [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' | wc -l) -eq 20 ]] + date -u +%FT%TZ >"${eval_root}/generation.completed" + else + echo "[skip] generation complete: ${eval_name}" + fi + + local metric_records=0 + if [[ -f "${eval_root}/metrics.json" ]]; then + metric_records=$("${python_bin}" -c \ + 'import json,sys; print(len(json.load(open(sys.argv[1]))["records"]))' \ + "${eval_root}/metrics.json") + fi + if [[ ! -f "${eval_root}/metrics.completed" || "${metric_records}" -ne 20 ]]; then + date -u +%FT%TZ >"${eval_root}/metrics.started" + run_workers "${log_dir}" metrics \ + "${repo_root}/tools/compute_predictor_v4_train_case_metrics.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --num_cases 25 \ + --case_ids "${case_ids}" \ + --reference_subdir full \ + --modes predictor \ + --metric_batch_size 4 + "${python_bin}" "${repo_root}/tools/merge_predictor_v4_train_case_metrics.py" \ + --output_root "${eval_root}" \ + --num_workers 8 \ + --num_cases 5 \ + --modes predictor \ + --replace_steps "${replace_steps}" \ + --reference_label "Full-DiT validation videos; ${subset_tag}" \ + >"${log_dir}/metrics_merge.log" 2>&1 + [[ $("${python_bin}" -c \ + 'import json,sys; print(len(json.load(open(sys.argv[1]))["records"]))' \ + "${eval_root}/metrics.json") -eq 20 ]] + date -u +%FT%TZ >"${eval_root}/metrics.completed" + else + echo "[skip] metrics complete: ${eval_name}" + fi + + aws s3 sync "${eval_root}" "${durable_eval_uri}/${eval_name}" --only-show-errors + date -u +%FT%TZ >"${eval_root}/sync.completed" + aws s3 cp "${eval_root}/sync.completed" \ + "${durable_eval_uri}/${eval_name}/sync.completed" --only-show-errors +} + +for spec in "${variants[@]}"; do + IFS='|' read -r variant weights <<<"${spec}" + run_with_retries "${variant}_fppf" evaluate_one "${variant}" "${weights}" fppf "1,2" + run_with_retries "${variant}_fppp" evaluate_one "${variant}" "${weights}" fppp "1,2,3" +done + +date -u +%FT%TZ >"${summary_root}/pipeline.completed" +aws s3 sync "${summary_root}" "${durable_eval_uri}/$(basename "${summary_root}")" --only-show-errors +echo "[pipeline] all validation5 evaluations completed" diff --git a/tools/run_stage1_continued_validation25_evaluation.sh b/tools/run_stage1_continued_validation25_evaluation.sh new file mode 100644 index 0000000000000000000000000000000000000000..d11c1be56855c9356bd1eb78e891112dcc438782 --- /dev/null +++ b/tools/run_stage1_continued_validation25_evaluation.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +vbench_repo="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +weights="${repo_root}/checkpoints/predictor_v4_prefeature-continued-from-stage1step2000_blocks1-52_fulltrain100_fp32master_8gpu_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs16_swanlab/predictor_step_02000.safetensors" +eval_prefix="eval_validation25_stage1-continued-offline-tf" +eval_suffix="blocks1-52_ckptstep2000-initstage1step2000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs16" +fppf_name="${eval_prefix}_schedule-fppf_${eval_suffix}" +fppp_name="${eval_prefix}_schedule-fppp_${eval_suffix}" +run_name="validation25_stage1_continued_offline_tf_step2000" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" + +export PYTHONPATH="${repo_root}" + +bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${fppf_name}" "${weights}" "1,52" v4 "1,2" +bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${fppp_name}" "${weights}" "1,52" v4 "1,2,3" + +export PYTHONPATH="${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/vbench" +export TORCH_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/torch" +export HF_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/huggingface" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +group_fppf="stage1_continued_offline_tf_fppf" +group_fppp="stage1_continued_offline_tf_fppp" +declare -a groups=( + "${group_fppf}|${dataset_root}/evaluations/${fppf_name}/predictor" + "${group_fppp}|${dataset_root}/evaluations/${fppp_name}/predictor" +) + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne 200 ]]; then + echo "Expected 200 staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 200 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json +import sys + +path, dimension = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == 200, (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + --group "${group_fppf}" \ + --group "${group_fppp}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +echo "[validation] all standard and VBench metrics complete: ${run_root}" diff --git a/tools/run_stage1_previous_condition_ablation.sh b/tools/run_stage1_previous_condition_ablation.sh new file mode 100644 index 0000000000000000000000000000000000000000..8789702a899d95d5dcb4a976dcf817b970e94f70 --- /dev/null +++ b/tools/run_stage1_previous_condition_ablation.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root=/mnt/local_nvme/cz/projects/HY-WorldPlay-DEV-Predictor +predictor_env=/mnt/local_nvme/cz/envs/hy_worldplay_july20260720 +python_bin="$predictor_env/bin/python" +manifest="$repo_root/datasets/hyworldplay_predictor_prefeature_v4/train_eval_manifest.jsonl" +teacher=/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub/models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/diffusion_pytorch_model.safetensors +suite_root="$repo_root/checkpoints/predictor_v4_stage1_previous_condition_ablation_seed0" +log_root="$repo_root/logs/predictor_v4_stage1_previous_condition_ablation_seed0" +s3_root=s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor +modes=(none previous_step both) + +export PATH="$predictor_env/bin:$PATH" +export PYTHONPATH="$repo_root" +export OMP_NUM_THREADS=4 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 +export SSL_CERT_FILE="$predictor_env/lib/python3.10/site-packages/certifi/cacert.pem" +export REQUESTS_CA_BUNDLE="$SSL_CERT_FILE" + +mkdir -p "$suite_root" "$log_root" + +"$python_bin" - <<'PY' | tee "$log_root/environment.log" +import importlib.util +import torch + +assert torch.__version__ == "2.6.0+cu124", torch.__version__ +assert torch.version.cuda == "12.4", torch.version.cuda +assert importlib.util.find_spec("flash_attn") is None, "flash_attn must not be installed" +print(f"torch={torch.__version__} cuda={torch.version.cuda} flash_attn=False") +PY + +for mode in "${modes[@]}"; do + output_dir="$suite_root/$mode" + log_file="$log_root/${mode}.log" + complete_marker="$output_dir/training.complete" + if [[ -f "$complete_marker" && -f "$output_dir/predictor_step_02000.safetensors" ]]; then + echo "[skip] $mode already complete" + continue + fi + mkdir -p "$output_dir" + resume_args=() + if [[ -f "$output_dir/training_latest.pt" ]]; then + resume_args+=(--resume "$output_dir/training_latest.pt") + fi + echo "[start] mode=$mode $(date -u +%FT%TZ)" | tee -a "$log_file" + "$python_bin" -m torch.distributed.run \ + --standalone --nproc_per_node=8 \ + "$repo_root/train_predictor_v4.py" \ + --manifest "$manifest" \ + --teacher_checkpoint "$teacher" \ + --source_blocks 1,52 \ + --previous_condition_mode "$mode" \ + --output_dir "$output_dir" \ + --batch_size 16 \ + --gradient_accumulation_steps 1 \ + --max_steps 2000 \ + --seed 0 \ + --fusion_lr 1e-4 \ + --blocks_lr 1e-5 \ + --scheduler_warmup_steps 100 \ + --block_warmup_steps 100 \ + --block_warmup_mode delayed \ + --fp32_trainable_params \ + --gradient_checkpointing \ + --bucket_by_context \ + --save_every 100 \ + --keep_snapshots 20 \ + --log_every 10 \ + --swanlab_mode cloud \ + --swanlab_experiment "predictor-v4-stage1-previous-${mode}-seed0-blocks1-52-fp32master-8gpu-torch26-noflash" \ + "${resume_args[@]}" \ + 2>&1 | tee -a "$log_file" + test -f "$output_dir/predictor_step_02000.safetensors" + date -u +%FT%TZ >"$complete_marker" + aws s3 sync "$output_dir" "$s3_root/checkpoints/predictor_v4_stage1_previous_condition_ablation_seed0/$mode" --only-show-errors + aws s3 sync "$log_root" "$s3_root/logs/predictor_v4_stage1_previous_condition_ablation_seed0" --only-show-errors + echo "[complete] mode=$mode $(date -u +%FT%TZ)" | tee -a "$log_file" +done + +date -u +%FT%TZ >"$suite_root/suite.complete" +aws s3 cp "$suite_root/suite.complete" "$s3_root/checkpoints/predictor_v4_stage1_previous_condition_ablation_seed0/suite.complete" --only-show-errors +echo "[complete] all modes $(date -u +%FT%TZ)" diff --git a/tools/run_stage1_previous_condition_retrain_validation.sh b/tools/run_stage1_previous_condition_retrain_validation.sh new file mode 100644 index 0000000000000000000000000000000000000000..f91170a06769536b8d4184d59262ad0040d09804 --- /dev/null +++ b/tools/run_stage1_previous_condition_retrain_validation.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +repo_root=/mnt/local_nvme/cz/projects/HY-WorldPlay-DEV-Predictor +predictor_env=/mnt/local_nvme/cz/envs/hy_worldplay_july20260720 +python_bin="$predictor_env/bin/python" +base_model="$repo_root/outputs/hf_snapshot_repaired_9b49404" +checkpoint_root="$repo_root/checkpoints/predictor_v4_stage1_previous_condition_ablation_seed0" +run_tag=torch26-retrain_20260824 +run_root="$repo_root/logs/stage1_previous_condition_validation_${run_tag}" +status_file="$run_root/status.tsv" + +export PATH="$predictor_env/bin:$PATH" +export PYTHONPATH="$repo_root" +export PREDICTOR_PYTHON_BIN="$python_bin" +export PREDICTOR_BASE_MODEL="$base_model" +export PREDICTOR_SKIP_SYNC=0 + +mkdir -p "$run_root" +exec > >(tee -a "$run_root/sequential.log") 2>&1 + +record_status() { + printf '%s\t%s\t%s\n' "$1" "$(date -u +%FT%TZ)" "$2" >>"$status_file" +} + +on_error() { + local exit_code="$?" + record_status failed "line=${BASH_LINENO[0]} exit=${exit_code}" + exit "$exit_code" +} +trap on_error ERR + +"$python_bin" - <<'PY' +import importlib.util +import torch + +assert torch.__version__ == "2.6.0+cu124", torch.__version__ +assert torch.version.cuda == "12.4", torch.version.cuda +assert importlib.util.find_spec("flash_attn") is None, "flash_attn must not be installed" +print(f"torch={torch.__version__} cuda={torch.version.cuda} flash_attn=False") +PY +test -d "$base_model" +record_status preflight "torch=2.6.0+cu124 cuda=12.4 flash_attn=False" + +run_eval() { + local label="$1" + local mode="$2" + local checkpoint="$checkpoint_root/$mode/predictor_step_02000.safetensors" + local eval_name="eval_validation25_previous-condition-${label}_${run_tag}" + test -f "$checkpoint" + record_status start "$eval_name" + bash "$repo_root/tools/launch_vbench_validation_predictor.sh" \ + "$eval_name" "$checkpoint" 1,52 v4 1,2 "$mode" concat + record_status complete "$eval_name" +} + +run_eval none none +run_eval previous-step previous_step +run_eval both both + +record_status complete all +date -u +%FT%TZ >"$run_root/all.completed" +aws s3 sync "$run_root" \ + "s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/logs/stage1_previous_condition_validation_${run_tag}" \ + --only-show-errors +echo "[complete] all previous-condition validation evaluations" diff --git a/tools/run_vbench_atc_validation25.sh b/tools/run_vbench_atc_validation25.sh new file mode 100644 index 0000000000000000000000000000000000000000..0ecd8c63a452bce5b08bd6f37f5426ef438f4678 --- /dev/null +++ b/tools/run_vbench_atc_validation25.sh @@ -0,0 +1,136 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/mnt/local_nvme/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_july20260720/bin/python}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/vbench_eval/lib/python3.10/site-packages}" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-${repo_root}/datasets/${dataset_name}}" +run_name="validation25_atc_ablation_oldenv" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" +overlay_root="${run_root}/site_overlay" +self_forcing_site_packages="${SELF_FORCING_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/self_forcing/lib/python3.10/site-packages}" + +mkdir -p "${overlay_root}" +for package in easydict timm open_clip; do + [[ -d "${self_forcing_site_packages}/${package}" ]] + ln -sfn "${self_forcing_site_packages}/${package}" "${overlay_root}/${package}" +done + +export PYTHONPATH="${overlay_root}:${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="${VBENCH_CACHE_DIR:-/mnt/local_nvme/cz/.cache/vbench}" +export TORCH_HOME="${TORCH_HOME:-/mnt/local_nvme/cz/.cache/torch}" +export HF_HOME="${HF_HOME:-/mnt/local_nvme/cz/.cache/huggingface}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" + +[[ -x "${python_bin}" ]] +[[ -f "${vbench_repo}/evaluate.py" ]] +[[ -f "${dataset_root}/validation/cases.jsonl" ]] +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +groups=() +group_names=() +for variant in full anchor_only raw_prev no_gate no_distance; do + for schedule in fppf fppp; do + group="atc_${variant}_${schedule}" + eval_name="eval_validation25_atc-${variant}_schedule-${schedule}_blocks1-52_step2000_oldenv" + groups+=("${group}|${dataset_root}/evaluations/${eval_name}/predictor") + group_names+=("${group}") + done +done + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +expected_records=1000 +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne "${expected_records}" ]]; then + echo "Expected ${expected_records} staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count "${expected_records}" + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" "${expected_records}" <<'PY' +import json +import sys + +path, dimension, expected_records = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == int(expected_records), (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + fi + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +summary_args=() +for group in "${group_names[@]}"; do + summary_args+=(--group "${group}") +done +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + "${summary_args[@]}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +fi +echo "[vbench] all ATC dimensions complete: ${run_root}" diff --git a/tools/run_vbench_c1_threshold_validation25.sh b/tools/run_vbench_c1_threshold_validation25.sh new file mode 100644 index 0000000000000000000000000000000000000000..788889436194cfff65bbd870a4ea92838d89663a --- /dev/null +++ b/tools/run_vbench_c1_threshold_validation25.sh @@ -0,0 +1,92 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/mnt/local_nvme/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/self_forcing/bin/python3.10}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/vbench_eval/lib/python3.10/site-packages}" +dataset_root="${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50" +run_root="${dataset_root}/vbench/validation25_c1_thresholds" +staging_root="${dataset_root}/vbench_staging/validation25_c1_thresholds" + +export PYTHONPATH="${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="${VBENCH_CACHE_DIR:-/mnt/local_nvme/cz/.cache/vbench}" +export TORCH_HOME="${TORCH_HOME:-/mnt/local_nvme/cz/.cache/torch}" +export HF_HOME="${HF_HOME:-/mnt/local_nvme/cz/.cache/huggingface}" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +groups=( + "c1_r10|${dataset_root}/evaluations/eval_validation25_stage1-c1-threshold-r10_blocks1-52_fp32master_step2000/predictor" + "c1_r25|${dataset_root}/evaluations/eval_validation25_stage1-c1-threshold-r25_blocks1-52_fp32master_step2000/predictor" + "c1_r33|${dataset_root}/evaluations/eval_validation25_stage1-c1-threshold-r33_blocks1-52_fp32master_step2000/predictor" + "c1_r50|${dataset_root}/evaluations/eval_validation25_stage1-c1-threshold-r50_blocks1-52_fp32master_step2000/predictor" +) + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + [[ $(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) -eq 100 ]] + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done +[[ $(find "${staging_root}" -maxdepth 1 -type l | wc -l) -eq 400 ]] + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 400 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + if [[ -f "${marker}" ]] \ + && [[ $(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) -eq 1 ]]; then + echo "[skip] ${dimension}" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json, sys +path, dimension = sys.argv[1:] +records = json.load(open(path))[dimension][1] +assert len(records) == 400, (dimension, len(records)) +print(dimension, len(records)) +PY + date -u +%FT%TZ >"${marker}" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + --group c1_r10 --group c1_r25 --group c1_r33 --group c1_r50 \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +echo "[vbench] all C1 threshold dimensions complete: ${run_root}" diff --git a/tools/run_vbench_previous_condition_none_step0300_validation25.sh b/tools/run_vbench_previous_condition_none_step0300_validation25.sh new file mode 100644 index 0000000000000000000000000000000000000000..098c4f13efc97366084896758386031f4ecdbb4e --- /dev/null +++ b/tools/run_vbench_previous_condition_none_step0300_validation25.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/mnt/local_nvme/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_july20260720/bin/python}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/vbench_eval/lib/python3.10/site-packages}" +self_forcing_site_packages="${SELF_FORCING_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/self_forcing/lib/python3.10/site-packages}" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-${repo_root}/datasets/${dataset_name}}" +run_name="validation25_previous_condition_none_step0300_torch26_20260825" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +source_dir="${dataset_root}/evaluations/eval_validation25_previous-condition-none-step0300_force-first-chunk-ffff_torch26_20260824/predictor" +group="none_step0300_first_chunk_ffff" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" +overlay_root="${run_root}/site_overlay" + +mkdir -p "${overlay_root}" +for package in easydict timm open_clip; do + [[ -d "${self_forcing_site_packages}/${package}" ]] + ln -sfn "${self_forcing_site_packages}/${package}" "${overlay_root}/${package}" +done + +export PYTHONPATH="${overlay_root}:${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="${VBENCH_CACHE_DIR:-/mnt/local_nvme/cz/.cache/vbench}" +export TORCH_HOME="${TORCH_HOME:-/mnt/local_nvme/cz/.cache/torch}" +export HF_HOME="${HF_HOME:-/mnt/local_nvme/cz/.cache/huggingface}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" + +[[ -x "${python_bin}" ]] +[[ -f "${vbench_repo}/evaluate.py" ]] +[[ -f "${dataset_root}/validation/cases.jsonl" ]] +[[ $(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) -eq 100 ]] +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +find "${staging_root}" -maxdepth 1 -type l -delete +while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${group}__$(basename "${video_path}")" +done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +[[ $(find "${staging_root}" -maxdepth 1 -type l | wc -l) -eq 100 ]] + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 100 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json +import sys + +path, dimension = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == 100, (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + fi + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + --group "${group}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +fi +echo "[vbench] none step0300 complete: ${run_root}" diff --git a/tools/run_vbench_previous_condition_validation25.sh b/tools/run_vbench_previous_condition_validation25.sh new file mode 100644 index 0000000000000000000000000000000000000000..80746f06c3402358787646c952cbd1a505dc3725 --- /dev/null +++ b/tools/run_vbench_previous_condition_validation25.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/mnt/local_nvme/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay_july20260720/bin/python}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/vbench_eval/lib/python3.10/site-packages}" +self_forcing_site_packages="${SELF_FORCING_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/self_forcing/lib/python3.10/site-packages}" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-${repo_root}/datasets/${dataset_name}}" +run_name="validation25_previous_condition_torch26_20260824" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" +overlay_root="${run_root}/site_overlay" + +mkdir -p "${overlay_root}" +for package in easydict timm open_clip; do + [[ -d "${self_forcing_site_packages}/${package}" ]] + ln -sfn "${self_forcing_site_packages}/${package}" "${overlay_root}/${package}" +done + +export PYTHONPATH="${overlay_root}:${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="${VBENCH_CACHE_DIR:-/mnt/local_nvme/cz/.cache/vbench}" +export TORCH_HOME="${TORCH_HOME:-/mnt/local_nvme/cz/.cache/torch}" +export HF_HOME="${HF_HOME:-/mnt/local_nvme/cz/.cache/huggingface}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" + +[[ -x "${python_bin}" ]] +[[ -f "${vbench_repo}/evaluate.py" ]] +[[ -f "${dataset_root}/validation/cases.jsonl" ]] +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +groups=( + "same_step_step2000|${dataset_root}/evaluations/eval_validation25_v4_blocks1-52_fp32master_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/predictor" + "previous_step_step2000|${dataset_root}/evaluations/eval_validation25_previous-condition-previous-step_torch26-retrain_20260824/predictor" + "both_step2000|${dataset_root}/evaluations/eval_validation25_previous-condition-both_torch26-retrain_20260824/predictor" + "none_step2000_first_chunk_fppf|${dataset_root}/evaluations/eval_validation25_previous-condition-none_torch26-retrain_20260824/predictor" + "none_step0500_first_chunk_ffff|${dataset_root}/evaluations/eval_validation25_previous-condition-none-step0500_force-first-chunk-ffff_torch26_20260824/predictor" + "none_step1000_first_chunk_ffff|${dataset_root}/evaluations/eval_validation25_previous-condition-none-step1000_force-first-chunk-ffff_torch26_20260824/predictor" + "none_step2000_first_chunk_ffff|${dataset_root}/evaluations/eval_validation25_previous-condition-none_force-first-chunk-ffff_torch26_20260824/predictor" +) +group_names=() + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + group_names+=("${label}") + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +expected_records=700 +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne "${expected_records}" ]]; then + echo "Expected ${expected_records} staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count "${expected_records}" + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" "${expected_records}" <<'PY' +import json +import sys + +path, dimension, expected_records = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == int(expected_records), (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + fi + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +summary_args=() +for group in "${group_names[@]}"; do + summary_args+=(--group "${group}") +done +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + "${summary_args[@]}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +if [[ "${PREDICTOR_SKIP_SYNC:-0}" != "1" ]]; then + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +fi +echo "[vbench] all previous-condition dimensions complete: ${run_root}" diff --git a/tools/run_vbench_token_gate_validation25.sh b/tools/run_vbench_token_gate_validation25.sh new file mode 100644 index 0000000000000000000000000000000000000000..f075801c6287c8871d110a17758647fc724d2c2e --- /dev/null +++ b/tools/run_vbench_token_gate_validation25.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/mnt/local_nvme/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/self_forcing/bin/python3.10}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/mnt/local_nvme/cz/envs/vbench_eval/lib/python3.10/site-packages}" +dataset_root="${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50" +source_dir="${dataset_root}/evaluations/eval_validation25_stage1-token-gate-prev_schedule-fppf_blocks1-52_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/predictor" +run_root="${dataset_root}/vbench/validation25_token_gate_prev_fppf" +staging_root="${dataset_root}/vbench_staging/validation25_token_gate_prev_fppf" +group=token_gate_prev_fppf + +export PYTHONPATH="${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="${VBENCH_CACHE_DIR:-/mnt/local_nvme/cz/.cache/vbench}" +export TORCH_HOME="${TORCH_HOME:-/mnt/local_nvme/cz/.cache/torch}" +export HF_HOME="${HF_HOME:-/mnt/local_nvme/cz/.cache/huggingface}" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" +find "${staging_root}" -maxdepth 1 -type l -delete +[[ $(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) -eq 100 ]] +while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${group}__$(basename "${video_path}")" +done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +[[ $(find "${staging_root}" -maxdepth 1 -type l | wc -l) -eq 100 ]] + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 100 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + if [[ -f "${marker}" ]] \ + && [[ $(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) -eq 1 ]]; then + echo "[skip] ${dimension}" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json, sys +path, dimension = sys.argv[1:] +records = json.load(open(path))[dimension][1] +assert len(records) == 100, (dimension, len(records)) +print(dimension, len(records)) +PY + date -u +%FT%TZ >"${marker}" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" \ + --run_root "${run_root}" \ + --group "${group}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +echo "[vbench] token-gate validation complete: ${run_root}" diff --git a/tools/run_vbench_validation25_chunk_timestep_step2000.sh b/tools/run_vbench_validation25_chunk_timestep_step2000.sh new file mode 100644 index 0000000000000000000000000000000000000000..4110875503da2b69046cfdbe151dcdc3ad68b465 --- /dev/null +++ b/tools/run_vbench_validation25_chunk_timestep_step2000.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +vbench_repo="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +run_name="validation25_chunk_timestep_step2000" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/${run_name}" + +export PYTHONPATH="${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/vbench" +export TORCH_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/torch" +export HF_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/huggingface" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +declare -a groups=( + "chunk_timestep_step2000_fppf|${dataset_root}/evaluations/eval_validation25_detached-chunk-timestep-rollout-onpolicy-fppp-fulltrain100_schedule-fppf_blocks1-52_ckptstep2000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" + "chunk_timestep_step2000_fppp|${dataset_root}/evaluations/eval_validation25_detached-chunk-timestep-rollout-onpolicy-fppp-fulltrain100_schedule-fppp_blocks1-52_ckptstep2000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" +) + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne 200 ]]; then + echo "Expected 200 staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" \ + --expected_count 200 + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json +import sys + +path, dimension = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == 200, (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +"${python_bin}" \ + "${repo_root}/tools/summarize_vbench_validation25_chunk_timestep_step2000.py" \ + --run_root "${run_root}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +echo "[vbench] all dimensions complete: ${run_root}" diff --git a/tools/run_vbench_validation25_core.sh b/tools/run_vbench_validation25_core.sh new file mode 100644 index 0000000000000000000000000000000000000000..6221d51be82a9e734102776a27168a2bf3d4fe04 --- /dev/null +++ b/tools/run_vbench_validation25_core.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +vbench_repo="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/VBench" +dataset_name="hyworldplay_predictor_vbench_val25_test50" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +run_root="${dataset_root}/vbench/validation25_core8" +staging_root="${dataset_root}/vbench_staging/validation25_core8" +python_bin="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/hy_worldplay/bin/python" +durable_uri="s3://s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}/vbench/validation25_core8" + +export PYTHONPATH="${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/vbench" +export TORCH_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/torch" +export HF_HOME="/mnt/s3files/s3-us-west2-default/zoubin/cz/.cache/huggingface" + +mkdir -p "${run_root}/raw" "${run_root}/logs" "${staging_root}" + +declare -a groups=( + "full_ffff|${dataset_root}/validation/full" + "reuse_frrf|${dataset_root}/validation/reuse" + "stage1_fppf|${dataset_root}/evaluations/eval_validation25_v4_blocks1-52_fp32master_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/predictor" + "stage1_fppp|${dataset_root}/evaluations/eval_validation25_stage1-v4_schedule-fppp_blocks1-52_fp32master_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100/predictor" + "timestep_onpolicy_fppf|${dataset_root}/evaluations/eval_validation25_detached-timestep-rollout-onpolicy-fulltrain100_schedule-fppf_blocks1-52_ckptstep1000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" + "timestep_onpolicy_fppp|${dataset_root}/evaluations/eval_validation25_detached-timestep-rollout-onpolicy-fulltrain100_schedule-fppp_blocks1-52_ckptstep1000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" + "chunk_timestep_fppf|${dataset_root}/evaluations/eval_validation25_detached-chunk-timestep-rollout-onpolicy-fppp-fulltrain100_schedule-fppf_blocks1-52_ckptstep1000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" + "chunk_timestep_fppp|${dataset_root}/evaluations/eval_validation25_detached-chunk-timestep-rollout-onpolicy-fppp-fulltrain100_schedule-fppp_blocks1-52_ckptstep1000_flr5e-5_blr5e-6_schedsteps2000_swarmup100_bwarmup100linear_bs1/predictor" +) + +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${groups[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + video_count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + if [[ "${video_count}" -ne 100 ]]; then + echo "${label}: expected 100 videos, found ${video_count}: ${source_dir}" >&2 + exit 1 + fi + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) +done + +staged_count=$(find "${staging_root}" -maxdepth 1 -type l | wc -l) +if [[ "${staged_count}" -ne 800 ]]; then + echo "Expected 800 staged videos, found ${staged_count}" >&2 + exit 1 +fi + +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" \ + --staging_root "${staging_root}" \ + --output "${prompt_file}" + +dimensions=( + subject_consistency + background_consistency + temporal_flickering + motion_smoothness + dynamic_degree + aesthetic_quality + imaging_quality + overall_consistency +) + +for dimension in "${dimensions[@]}"; do + output_dir="${run_root}/raw/${dimension}" + marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + result_count=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) + if [[ -f "${marker}" && "${result_count}" -eq 1 ]]; then + echo "[skip] ${dimension} already complete" + continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run \ + --standalone \ + --nproc_per_node=8 \ + "${vbench_repo}/evaluate.py" \ + --videos_path "${staging_root}" \ + --dimension "${dimension}" \ + --mode custom_input \ + --prompt_file "${prompt_file}" \ + --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer \ + --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" <<'PY' +import json +import sys + +path, dimension = sys.argv[1:] +payload = json.load(open(path)) +records = payload[dimension][1] +assert len(records) == 800, (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done + +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_core.py" \ + --run_root "${run_root}" \ + >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +aws s3 sync "${run_root}" "${durable_uri}" --only-show-errors +echo "[vbench] all dimensions complete: ${run_root}" diff --git a/tools/run_vbench_validation25_groups.sh b/tools/run_vbench_validation25_groups.sh new file mode 100644 index 0000000000000000000000000000000000000000..d22782160328fb1e09763dac6f57d52a847792a8 --- /dev/null +++ b/tools/run_vbench_validation25_groups.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# Generic VBench (8 dimensions) over staged validation25 groups. +# RUN_NAME= NPROC=4 VBENCH_GROUPS="label|dir ..." bash tools/run_vbench_validation25_groups.sh +# VBENCH_GROUPS entries are "label|/predictor" separated by spaces; each dir +# must hold exactly 100 mp4s. Local paths; no S3 sync. +set -euo pipefail +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +vbench_repo="${VBENCH_REPO:-/local/zoubin/cz/projects/VBench}" +python_bin="${PREDICTOR_PYTHON_BIN:-/local/zoubin/cz/envs/hy_worldplay_july20260720/bin/python}" +vbench_site_packages="${VBENCH_SITE_PACKAGES:-/local/zoubin/cz/envs/vbench_eval/lib/python3.10/site-packages}" +self_forcing_site_packages="${SELF_FORCING_SITE_PACKAGES:-/local/zoubin/cz/envs/self_forcing/lib/python3.10/site-packages}" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50}" +run_name="${RUN_NAME:?set RUN_NAME}" +nproc="${NPROC:-8}" +run_root="${dataset_root}/vbench/${run_name}" +staging_root="${dataset_root}/vbench_staging/${run_name}" +overlay_root="${run_root}/site_overlay" +mkdir -p "${overlay_root}" "${run_root}/raw" "${run_root}/logs" "${staging_root}" +for package in easydict timm open_clip; do + [[ -d "${self_forcing_site_packages}/${package}" ]] + ln -sfn "${self_forcing_site_packages}/${package}" "${overlay_root}/${package}" +done +export PYTHONPATH="${overlay_root}:${vbench_site_packages}:${vbench_repo}:${repo_root}" +export VBENCH_CACHE_DIR="/local/zoubin/cz/.cache/vbench" +export TORCH_HOME="/local/zoubin/cz/.cache/torch" +export TORCH_HUB="${TORCH_HOME}/hub" # torch.hub reads TORCH_HUB first; the login shell sets it to /mnt/local_nvme +export HF_HOME="/local/zoubin/cz/.cache/huggingface" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" +[[ -x "${python_bin}" && -f "${vbench_repo}/evaluate.py" && -f "${dataset_root}/validation/cases.jsonl" ]] +read -r -a group_specs <<< "${VBENCH_GROUPS:?set VBENCH_GROUPS}" +group_names=() +find "${staging_root}" -maxdepth 1 -type l -delete +for spec in "${group_specs[@]}"; do + IFS='|' read -r label source_dir <<<"${spec}" + count=$(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | wc -l) + [[ "${count}" -eq 100 ]] || { echo "${label}: expected 100 videos, found ${count}: ${source_dir}" >&2; exit 1; } + while IFS= read -r video_path; do + ln -s "${video_path}" "${staging_root}/${label}__$(basename "${video_path}")" + done < <(find "${source_dir}" -maxdepth 1 -type f -name '*.mp4' | sort) + group_names+=("${label}") +done +expected_records=$(( ${#group_names[@]} * 100 )) +[[ $(find "${staging_root}" -maxdepth 1 -type l | wc -l) -eq "${expected_records}" ]] +prompt_file="${run_root}/validation25_prompt_map.json" +"${python_bin}" "${repo_root}/tools/build_vbench_validation25_prompt_map.py" \ + --cases_jsonl "${dataset_root}/validation/cases.jsonl" --staging_root "${staging_root}" \ + --output "${prompt_file}" --expected_count "${expected_records}" +for dimension in subject_consistency background_consistency temporal_flickering motion_smoothness dynamic_degree aesthetic_quality imaging_quality overall_consistency; do + output_dir="${run_root}/raw/${dimension}"; marker="${run_root}/${dimension}.completed" + mkdir -p "${output_dir}" + if [[ -f "${marker}" && $(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json' | wc -l) -eq 1 ]]; then + echo "[skip] ${dimension} already complete"; continue + fi + find "${output_dir}" -maxdepth 1 -type f -delete + echo "[vbench] ${dimension} start $(date -u +%FT%TZ)" + "${python_bin}" -m torch.distributed.run --standalone --nproc_per_node="${nproc}" \ + "${vbench_repo}/evaluate.py" --videos_path "${staging_root}" --dimension "${dimension}" \ + --mode custom_input --prompt_file "${prompt_file}" --load_ckpt_from_local True \ + --imaging_quality_preprocessing_mode longer --output_path "${output_dir}" \ + >"${run_root}/logs/${dimension}.log" 2>&1 + result_json=$(find "${output_dir}" -maxdepth 1 -type f -name '*_eval_results.json') + "${python_bin}" - "${result_json}" "${dimension}" "${expected_records}" <<'PY' +import json, sys +path, dimension, expected = sys.argv[1:] +payload = json.load(open(path)); records = payload[dimension][1] +assert len(records) == int(expected), (dimension, len(records)) +print(dimension, len(records), payload[dimension][0]) +PY + date -u +%FT%TZ >"${marker}" + echo "[vbench] ${dimension} complete $(date -u +%FT%TZ)" +done +summary_args=(); for group in "${group_names[@]}"; do summary_args+=(--group "${group}"); done +"${python_bin}" "${repo_root}/tools/summarize_vbench_validation25_pair.py" --run_root "${run_root}" "${summary_args[@]}" >"${run_root}/logs/summary.log" 2>&1 +date -u +%FT%TZ >"${run_root}/pipeline.completed" +echo "[vbench] complete: ${run_root}" diff --git a/tools/run_wide_c1_r33_validation25_reference.sh b/tools/run_wide_c1_r33_validation25_reference.sh new file mode 100644 index 0000000000000000000000000000000000000000..e9db7ba858f7f956e9c7ca21ce2dbfda04134b56 --- /dev/null +++ b/tools/run_wide_c1_r33_validation25_reference.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay/bin/python3.10}" +dataset_root="${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50" +validation_root="${dataset_root}/validation" +base_model="${repo_root}/outputs/hf_snapshot_repaired_9b49404" +action_ckpt="/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub/models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/diffusion_pytorch_model.safetensors" +weights="${repo_root}/checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" +confidence_checkpoint="${repo_root}/confidence_validation_wide4096x2_c1/checkpoints/confidence_c1.pt" +confidence_threshold="-4.308430990" +eval_name="eval_validation25_stage1-wide4096x2-c1-threshold-r33_blocks1-52_fp32master_step2000" +eval_root="${dataset_root}/evaluations/${eval_name}" +log_dir="${eval_root}/logs" + +export PYTHONPATH="${repo_root}" +export HY_PROFILE_TIMING=1 +mkdir -p "${log_dir}" + +run_workers() { + local phase="$1" + shift + local -a pids=() + local worker_id + for worker_id in $(seq 0 7); do + CUDA_VISIBLE_DEVICES="${worker_id}" "${python_bin}" -u "$@" \ + --worker_id "${worker_id}" --num_workers 8 \ + >"${log_dir}/${phase}_worker_${worker_id}.log" 2>&1 & + pids+=("$!") + done + local failed=0 pid + for pid in "${pids[@]}"; do + if ! wait "${pid}"; then failed=1; fi + done + if [[ "${failed}" -ne 0 ]]; then + echo "${phase} worker failure; inspect ${log_dir}/${phase}_worker_*.log" >&2 + return 1 + fi +} + +if [[ ! -f "${eval_root}/generation.completed" ]] \ + || [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' 2>/dev/null | wc -l) -ne 100 ]]; then + date -u +%FT%TZ >"${eval_root}/generation.started" + run_workers generate \ + "${repo_root}/tools/generate_predictor_v4_train_case_eval.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --base_model "${base_model}" \ + --action_ckpt "${action_ckpt}" \ + --weights "${weights}" \ + --source_blocks 1,52 \ + --predictor_variant v4 \ + --previous_condition_mode same_step \ + --replace_steps 1,2,3 \ + --modes predictor \ + --confidence_checkpoint "${confidence_checkpoint}" \ + --confidence_threshold "${confidence_threshold}" + [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' | wc -l) -eq 100 ]] + date -u +%FT%TZ >"${eval_root}/generation.completed" +fi + +if [[ ! -f "${eval_root}/metrics.completed" ]]; then + date -u +%FT%TZ >"${eval_root}/metrics.started" + run_workers metrics \ + "${repo_root}/tools/compute_predictor_v4_train_case_metrics.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --num_cases 25 \ + --reference_subdir full \ + --modes predictor \ + --metric_batch_size 4 + "${python_bin}" "${repo_root}/tools/merge_predictor_v4_train_case_metrics.py" \ + --output_root "${eval_root}" \ + --num_workers 8 \ + --num_cases 25 \ + --modes predictor \ + --replace_steps 1,2,3 \ + --reference_label "Full-DiT validation videos" \ + >"${log_dir}/metrics_merge.log" 2>&1 + date -u +%FT%TZ >"${eval_root}/metrics.completed" +fi + +echo "Validation evaluation complete: ${eval_root}" diff --git a/tools/run_wide_c2_hybrid_validation25_reference.sh b/tools/run_wide_c2_hybrid_validation25_reference.sh new file mode 100644 index 0000000000000000000000000000000000000000..0410d03b9c7d1fe9271703ce24f4c0f765183555 --- /dev/null +++ b/tools/run_wide_c2_hybrid_validation25_reference.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +python_bin="${PREDICTOR_PYTHON_BIN:-/mnt/local_nvme/cz/envs/hy_worldplay/bin/python3.10}" +dataset_root="${repo_root}/datasets/hyworldplay_predictor_vbench_val25_test50" +validation_root="${dataset_root}/validation" +base_model="${repo_root}/outputs/hf_snapshot_repaired_9b49404" +action_ckpt="/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub/models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/diffusion_pytorch_model.safetensors" +weights="${repo_root}/checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" +confidence_checkpoint="${repo_root}/confidence_validation_wide4096x2_c2/checkpoints/confidence_c2.pt" + +tags=(p12r10 p12r25) +thresholds=( + "1:-4.136674309,2:-3.798280764" + "1:-4.400576830,2:-4.049407601" +) +gpu_groups=("0 1 2 3" "4 5 6 7") + +export PYTHONPATH="${repo_root}" +export HY_PROFILE_TIMING=1 + +run_group() { + local phase="$1" group_index="$2" + local tag="${tags[$group_index]}" + local threshold_spec="${thresholds[$group_index]}" + local eval_name="eval_validation25_stage1-wide4096x2-c2-hybrid-p3full-${tag}_blocks1-52_fp32master_step2000" + local eval_root="${dataset_root}/evaluations/${eval_name}" + local log_dir="${eval_root}/logs" + local -a gpus=(${gpu_groups[$group_index]}) + mkdir -p "${log_dir}" + + if [[ "${phase}" == "generate" ]] \ + && [[ -f "${eval_root}/generation.completed" ]] \ + && [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' 2>/dev/null | wc -l) -eq 100 ]]; then + return + fi + if [[ "${phase}" == "metrics" && -f "${eval_root}/metrics.completed" ]]; then + return + fi + + local worker_id gpu failed=0 pid + local -a pids=() + for worker_id in "${!gpus[@]}"; do + gpu="${gpus[$worker_id]}" + if [[ "${phase}" == "generate" ]]; then + CUDA_VISIBLE_DEVICES="${gpu}" "${python_bin}" -u \ + "${repo_root}/tools/generate_predictor_v4_train_case_eval.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --base_model "${base_model}" \ + --action_ckpt "${action_ckpt}" \ + --weights "${weights}" \ + --source_blocks 1,52 \ + --predictor_variant v4 \ + --previous_condition_mode same_step \ + --replace_steps 1,2 \ + --modes predictor \ + --confidence_checkpoint "${confidence_checkpoint}" \ + --confidence_thresholds "${threshold_spec}" \ + --worker_id "${worker_id}" --num_workers "${#gpus[@]}" \ + >"${log_dir}/generate_worker_${worker_id}.log" 2>&1 & + else + CUDA_VISIBLE_DEVICES="${gpu}" "${python_bin}" -u \ + "${repo_root}/tools/compute_predictor_v4_train_case_metrics.py" \ + --dataset_root "${validation_root}" \ + --output_root "${eval_root}" \ + --num_cases 25 \ + --reference_subdir full \ + --modes predictor \ + --metric_batch_size 4 \ + --worker_id "${worker_id}" --num_workers "${#gpus[@]}" \ + >"${log_dir}/metrics_worker_${worker_id}.log" 2>&1 & + fi + pids+=("$!") + done + for pid in "${pids[@]}"; do + if ! wait "${pid}"; then failed=1; fi + done + if [[ "${failed}" -ne 0 ]]; then + echo "${phase} failed for ${tag}; inspect ${log_dir}" >&2 + return 1 + fi + + if [[ "${phase}" == "generate" ]]; then + [[ $(find "${eval_root}/predictor" -type f -name '*.mp4' | wc -l) -eq 100 ]] + date -u +%FT%TZ >"${eval_root}/generation.completed" + else + "${python_bin}" "${repo_root}/tools/merge_predictor_v4_train_case_metrics.py" \ + --output_root "${eval_root}" \ + --num_workers "${#gpus[@]}" \ + --num_cases 25 \ + --modes predictor \ + --replace_steps 1,2 \ + --reference_label "Full-DiT validation videos" \ + >"${log_dir}/metrics_merge.log" 2>&1 + date -u +%FT%TZ >"${eval_root}/metrics.completed" + fi +} + +run_phase() { + local phase="$1" failed=0 pid group_index + local -a pids=() + for group_index in "${!tags[@]}"; do + run_group "${phase}" "${group_index}" & + pids+=("$!") + done + for pid in "${pids[@]}"; do + if ! wait "${pid}"; then failed=1; fi + done + [[ "${failed}" -eq 0 ]] +} + +run_phase generate +run_phase metrics +echo "Wide C2 hybrid validation25 evaluation complete" diff --git a/tools/summarize_atc_rollout_ckpt_sweep.py b/tools/summarize_atc_rollout_ckpt_sweep.py new file mode 100644 index 0000000000000000000000000000000000000000..5dead111633dafa75a7494773c75a1d831d6315e --- /dev/null +++ b/tools/summarize_atc_rollout_ckpt_sweep.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python +"""Full-reference metrics + VBench Core5 table for the ATC rollout checkpoint sweep.""" +from __future__ import annotations +import argparse, json +from pathlib import Path + +CORE5 = ("subject_consistency", "background_consistency", "motion_smoothness", "aesthetic_quality", "imaging_quality") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--dataset_root", type=Path, required=True) + p.add_argument("--vbench_run", required=True) + p.add_argument("--group", action="append", required=True, help="label|eval_dir_name") + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + vb = json.load(open(args.dataset_root / "vbench" / args.vbench_run / "summary.json")) + groups = vb.get("groups", vb) + rows = [] + for spec in args.group: + label, eval_name = spec.split("|", 1) + m = json.load(open(args.dataset_root / "evaluations" / eval_name / "metrics.json"))["summary_mean"]["predictor"] + g = groups.get(label, {}).get("mean", {}) + core5 = sum(g[d] / (100.0 if d == "imaging_quality" and g[d] > 1.5 else 1.0) for d in CORE5) / 5 if all(d in g for d in CORE5) else float("nan") + rows.append({"label": label, "psnr": m["psnr_db"], "frame_psnr": m["frame_psnr_mean_db"], "ssim": m["ssim"], + "lpips": m["lpips_alex"], "mse": m["pixel_mse"], "temporal_mse": m["temporal_delta_mse"], + "vbench_core5": core5, **{f"vbench_{d}": g.get(d, float("nan")) for d in g}}) + lines = ["| 组 | PSNR ↑ | Frame PSNR ↑ | SSIM ↑ | LPIPS ↓ | MSE ↓ | Temporal MSE ↓ | VBench Core5 ↑ |", "|---|---:|---:|---:|---:|---:|---:|---:|"] + for r in rows: + lines.append(f"| {r['label']} | {r['psnr']:.4f} | {r['frame_psnr']:.4f} | {r['ssim']:.5f} | {r['lpips']:.5f} | {r['mse']:.6f} | {r['temporal_mse']:.6f} | {r['vbench_core5']:.6f} |") + args.output.with_suffix(".json").write_text(json.dumps(rows, indent=2) + "\n") + args.output.write_text("\n".join(lines) + "\n") + print("\n".join(lines)) + + +if __name__ == "__main__": + main() diff --git a/tools/summarize_c1_threshold_validation25.py b/tools/summarize_c1_threshold_validation25.py new file mode 100644 index 0000000000000000000000000000000000000000..41b8807541f1d359c1bf8db6b9e139c8d94228f4 --- /dev/null +++ b/tools/summarize_c1_threshold_validation25.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Summarize Stage-1 C1-threshold validation25 reference metrics and routing.""" + +from __future__ import annotations + +import json +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +DATASET = ROOT / "datasets/hyworldplay_predictor_vbench_val25_test50" +EVALUATIONS = DATASET / "evaluations" +OUTPUT = ROOT / "confidence_validation/c1_threshold_validation25" +BASELINES = { + "FFFF": None, + "FPPF": "eval_validation25_v4_blocks1-52_fp32master_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100", + "FPPP": "eval_validation25_stage1-v4_schedule-fppp_blocks1-52_fp32master_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100", +} +THRESHOLDS = { + "C1-r10": ("eval_validation25_stage1-c1-threshold-r10_blocks1-52_fp32master_step2000", -3.523588657), + "C1-r25": ("eval_validation25_stage1-c1-threshold-r25_blocks1-52_fp32master_step2000", -3.915276289), + "C1-r33": ("eval_validation25_stage1-c1-threshold-r33_blocks1-52_fp32master_step2000", -4.329442501), + "C1-r50": ("eval_validation25_stage1-c1-threshold-r50_blocks1-52_fp32master_step2000", -4.798926830), +} +METRICS = ( + "psnr_db", "frame_psnr_mean_db", "ssim", "lpips_alex", "pixel_mae", + "pixel_mse", "pixel_rmse", "temporal_delta_mse", +) +VBENCH_DIMENSIONS = ( + "subject_consistency", "background_consistency", "temporal_flickering", + "motion_smoothness", "dynamic_degree", "aesthetic_quality", + "imaging_quality", "overall_consistency", +) +VBENCH_CORE5 = ( + "subject_consistency", "background_consistency", "motion_smoothness", + "aesthetic_quality", "imaging_quality", +) + + +def read_metrics(eval_name: str) -> dict: + payload = json.loads((EVALUATIONS / eval_name / "metrics.json").read_text()) + return payload["summary_mean"]["predictor"] + + +def read_decisions(eval_name: str) -> tuple[list[dict], list[float]]: + root = EVALUATIONS / eval_name + decisions = [] + for path in sorted((root / "confidence_decisions").glob("worker_*.jsonl")): + decisions.extend(json.loads(line) for line in path.read_text().splitlines()) + elapsed = [] + for path in sorted((root / "logs").glob("worker_*.jsonl")): + elapsed.extend(json.loads(line)["elapsed_s"] for line in path.read_text().splitlines()) + return decisions, elapsed + + +def main() -> None: + rows = [] + rows.append({ + "method": "FFFF", "threshold_log_c1": None, "fallback_rate": 1.0, + **{name: (float("inf") if "psnr" in name else 1.0 if name == "ssim" else 0.0) for name in METRICS}, + }) + for method in ("FPPF", "FPPP"): + fallback = 1 / 3 if method == "FPPF" else 0.0 + rows.append({ + "method": method, "threshold_log_c1": None, + "fallback_rate": fallback, **read_metrics(BASELINES[method]), + }) + for method, (eval_name, threshold) in THRESHOLDS.items(): + decisions, elapsed = read_decisions(eval_name) + fallback = [item for item in decisions if item["fallback_to_full"]] + by_step = { + f"p{step}_fallback_rate": sum( + item["fallback_to_full"] for item in decisions + if item["denoise_step"] == step + ) / sum(item["denoise_step"] == step for item in decisions) + for step in (1, 2, 3) + } + rows.append({ + "method": method, + "threshold_log_c1": threshold, + "fallback_rate": len(fallback) / len(decisions), + **by_step, + "mean_end_to_end_s": sum(elapsed) / len(elapsed), + **read_metrics(eval_name), + }) + core_vbench = json.loads( + (DATASET / "vbench/validation25_core8/summary.json").read_text() + )["groups"] + threshold_vbench = json.loads( + (DATASET / "vbench/validation25_c1_thresholds/summary.json").read_text() + )["groups"] + vbench_groups = { + "FFFF": core_vbench["full_ffff"]["mean"], + "FPPF": core_vbench["stage1_fppf"]["mean"], + "FPPP": core_vbench["stage1_fppp"]["mean"], + "C1-r10": threshold_vbench["c1_r10"]["mean"], + "C1-r25": threshold_vbench["c1_r25"]["mean"], + "C1-r33": threshold_vbench["c1_r33"]["mean"], + "C1-r50": threshold_vbench["c1_r50"]["mean"], + } + for row in rows: + scores = vbench_groups[row["method"]] + row.update({f"vbench_{name}": scores[name] for name in VBENCH_DIMENSIONS}) + row["vbench_core5_mean"] = sum(scores[name] for name in VBENCH_CORE5) / len(VBENCH_CORE5) + by_method = {row["method"]: row for row in rows} + comparison_metrics = ( + "psnr_db", "ssim", "lpips_alex", "pixel_mse", "temporal_delta_mse", + *[f"vbench_{name}" for name in VBENCH_DIMENSIONS], "vbench_core5_mean", + ) + comparisons = {} + for candidate in ("C1-r33", "C1-r50"): + comparisons[f"{candidate}_vs_FPPF"] = { + name: by_method[candidate][name] - by_method["FPPF"][name] + for name in comparison_metrics + } + OUTPUT.mkdir(parents=True, exist_ok=True) + payload = { + "protocol": "validation25, seed 0, Stage-1 v4 FP32-master step 2000", + "rows": rows, + "comparisons": comparisons, + } + (OUTPUT / "summary.json").write_text(json.dumps(payload, indent=2, allow_nan=True) + "\n") + columns = [ + "method", "threshold_log_c1", "fallback_rate", "p1_fallback_rate", + "p2_fallback_rate", "p3_fallback_rate", "mean_end_to_end_s", *METRICS, + *[f"vbench_{name}" for name in VBENCH_DIMENSIONS], "vbench_core5_mean", + ] + lines = [",".join(columns)] + for row in rows: + lines.append(",".join(str(row.get(column, "")) for column in columns)) + (OUTPUT / "summary.csv").write_text("\n".join(lines) + "\n") + reference_lines = [ + "# C1 threshold validation25 comparison", + "", + "## Reference metrics against FFFF", + "", + "| Method | Fallback | P1/P2/P3 fallback | PSNR | SSIM | LPIPS | Pixel MSE | Temporal MSE |", + "| --- | ---: | --- | ---: | ---: | ---: | ---: | ---: |", + ] + for row in rows: + routing = "-" if "p1_fallback_rate" not in row else "/".join( + f"{100 * row[f'p{step}_fallback_rate']:.1f}%" for step in (1, 2, 3) + ) + psnr = "∞" if row["psnr_db"] == float("inf") else f"{row['psnr_db']:.4f}" + reference_lines.append( + f"| {row['method']} | {100 * row['fallback_rate']:.1f}% | {routing} | " + f"{psnr} | {row['ssim']:.5f} | {row['lpips_alex']:.5f} | " + f"{row['pixel_mse']:.6f} | {row['temporal_delta_mse']:.6f} |" + ) + reference_lines.extend([ + "", "## VBench", "", + "| Method | Subject | Background | Flicker | Motion | Dynamic | Aesthetic | Imaging | Overall | Core5 mean |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", + ]) + for row in rows: + values = " | ".join( + f"{row[f'vbench_{name}']:.6f}" for name in VBENCH_DIMENSIONS + ) + reference_lines.append( + f"| {row['method']} | {values} | {row['vbench_core5_mean']:.6f} |" + ) + (OUTPUT / "report.md").write_text("\n".join(reference_lines) + "\n") + print(json.dumps(payload, indent=2, allow_nan=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/summarize_predictor_timestep_rollout_overfit8.py b/tools/summarize_predictor_timestep_rollout_overfit8.py new file mode 100644 index 0000000000000000000000000000000000000000..e6289b982b41b38f26008830512aca0fa23da35d --- /dev/null +++ b/tools/summarize_predictor_timestep_rollout_overfit8.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Summarize the six FPPP/FPPF evaluations from the replay sweep.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +SCHEDULES = ("fppp", "fppf") +METRICS = ( + "psnr_db", "frame_psnr_mean_db", "ssim", "lpips_alex", + "pixel_mse", "pixel_rmse", "pixel_mae", "temporal_delta_mse", +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--eval_parent", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--experiment", + default="v4 timestep-only FPPP rollout overfit8, 200 optimizer steps", + ) + parser.add_argument( + "--variants", + default="replay0:0.0 replay0p2:0.2 replay0p4:0.4", + ) + args = parser.parse_args() + variants = tuple( + (spec.split(":", 1)[0], float(spec.split(":", 1)[1])) + for spec in args.variants.split() + ) + rows = [] + for label, replay_weight in variants: + for schedule in SCHEDULES: + path = args.eval_parent / f"{label}_{schedule}" / "metrics.json" + with path.open("r", encoding="utf-8") as handle: + payload = json.load(handle) + summary = payload["summary_mean"]["predictor"] + rows.append( + { + "variant": label, + "replay_weight": replay_weight, + "evaluation_schedule": schedule.upper(), + "replace_steps": payload["config"]["replace_steps"], + "videos": payload["config"]["videos_per_mode"], + **{name: summary[name] for name in METRICS}, + "metrics_path": str(path), + } + ) + result = { + "experiment": args.experiment, + "training_variants": [ + {"variant": label, "replay_weight": weight} for label, weight in variants + ], + "evaluations_per_variant": ["FPPP", "FPPF"], + "videos_per_evaluation": 32, + "rows": rows, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(result, indent=2, ensure_ascii=False, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(result, indent=2, ensure_ascii=False, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/summarize_vbench_validation25_chunk_timestep_step2000.py b/tools/summarize_vbench_validation25_chunk_timestep_step2000.py new file mode 100644 index 0000000000000000000000000000000000000000..47024345911eca12ce95856fd5de143d4690aafa --- /dev/null +++ b/tools/summarize_vbench_validation25_chunk_timestep_step2000.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +"""Summarize VBench results for chunk+timestep step-2000 FPPF/FPPP videos.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +from collections import defaultdict +from pathlib import Path +from statistics import fmean + + +DIMENSIONS = [ + "subject_consistency", + "background_consistency", + "temporal_flickering", + "motion_smoothness", + "dynamic_degree", + "aesthetic_quality", + "imaging_quality", + "overall_consistency", +] + +GROUP_ORDER = [ + "chunk_timestep_step2000_fppf", + "chunk_timestep_step2000_fppp", +] + +ACTION_RE = re.compile(r"_action_(\d{2})$") + + +def load_dimension_result(run_root: Path, dimension: str) -> dict: + candidates = sorted((run_root / "raw" / dimension).glob("*_eval_results.json")) + if len(candidates) != 1: + raise RuntimeError( + f"Expected exactly one result JSON for {dimension}, got {len(candidates)}" + ) + payload = json.loads(candidates[0].read_text()) + if dimension not in payload: + raise KeyError(f"{dimension} missing from {candidates[0]}") + overall, records = payload[dimension] + if len(records) != 200: + raise RuntimeError(f"{dimension}: expected 200 records, got {len(records)}") + return { + "source": str(candidates[0]), + "official_overall": float(overall), + "records": records, + } + + +def normalize_score(dimension: str, value: float) -> float: + if dimension == "imaging_quality": + return float(value) / 100.0 + return float(value) + + +def parse_record(record: dict, dimension: str) -> dict: + video_path = Path(record["video_path"]) + stem = video_path.stem + if "__" not in stem: + raise ValueError(f"Missing group prefix in staged video name: {video_path}") + group, original_stem = stem.split("__", 1) + if group not in GROUP_ORDER: + raise ValueError(f"Unknown VBench group {group}: {video_path}") + action_match = ACTION_RE.search(original_stem) + if action_match is None: + raise ValueError(f"Cannot parse action id from {video_path}") + return { + "group": group, + "video_id": original_stem, + "action_id": int(action_match.group(1)), + "video_path": str(video_path), + "score": normalize_score(dimension, record["video_results"]), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run_root", type=Path, required=True) + args = parser.parse_args() + run_root = args.run_root.resolve() + + parsed_by_dimension: dict[str, list[dict]] = {} + official_overall: dict[str, float] = {} + sources: dict[str, str] = {} + for dimension in DIMENSIONS: + result = load_dimension_result(run_root, dimension) + parsed_by_dimension[dimension] = [ + parse_record(record, dimension) for record in result["records"] + ] + official_overall[dimension] = result["official_overall"] + sources[dimension] = result["source"] + + per_video: dict[str, dict[str, dict]] = defaultdict(dict) + for dimension, records in parsed_by_dimension.items(): + counts = defaultdict(int) + for record in records: + counts[record["group"]] += 1 + item = per_video[record["group"]].setdefault( + record["video_id"], + { + "action_id": record["action_id"], + "video_id": record["video_id"], + "scores": {}, + }, + ) + item["scores"][dimension] = record["score"] + for group in GROUP_ORDER: + if counts[group] != 100: + raise RuntimeError( + f"{dimension}/{group}: expected 100 records, got {counts[group]}" + ) + + group_summary = {} + for group in GROUP_ORDER: + items = list(per_video[group].values()) + if len(items) != 100: + raise RuntimeError(f"{group}: expected 100 paired videos, got {len(items)}") + mean = { + dimension: fmean(item["scores"][dimension] for item in items) + for dimension in DIMENSIONS + } + mean_by_action_id = {} + for action_id in range(4): + action_items = [item for item in items if item["action_id"] == action_id] + if len(action_items) != 25: + raise RuntimeError( + f"{group}/action{action_id}: expected 25 videos, " + f"got {len(action_items)}" + ) + mean_by_action_id[str(action_id)] = { + dimension: fmean( + item["scores"][dimension] for item in action_items + ) + for dimension in DIMENSIONS + } + group_summary[group] = { + "record_count": len(items), + "mean": mean, + "mean_by_action_id": mean_by_action_id, + } + + payload = { + "dimensions": DIMENSIONS, + "group_order": GROUP_ORDER, + "official_overall_all_200_videos": official_overall, + "raw_sources": sources, + "groups": group_summary, + } + (run_root / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + with (run_root / "summary.csv").open("w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["group", "video_count", *DIMENSIONS]) + for group in GROUP_ORDER: + writer.writerow( + [ + group, + group_summary[group]["record_count"], + *[ + f"{group_summary[group]['mean'][dimension]:.8f}" + for dimension in DIMENSIONS + ], + ] + ) + + header = ( + "| 版本 | 视频数 | Subject ↑ | Background ↑ | Flickering ↑ | " + "Motion smoothness ↑ | Dynamic degree ↑ | Aesthetic ↑ | Imaging ↑ | " + "Overall consistency ↑ |" + ) + divider = ( + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + rows = [header, divider] + for group in GROUP_ORDER: + mean = group_summary[group]["mean"] + rows.append( + "| " + + " | ".join( + [ + group, + str(group_summary[group]["record_count"]), + f"{mean['subject_consistency']:.6f}", + f"{mean['background_consistency']:.6f}", + f"{mean['temporal_flickering']:.6f}", + f"{mean['motion_smoothness']:.6f}", + f"{mean['dynamic_degree']:.6f}", + f"{mean['aesthetic_quality']:.6f}", + f"{mean['imaging_quality']:.6f}", + f"{mean['overall_consistency']:.6f}", + ] + ) + + " |" + ) + (run_root / "summary.md").write_text("\n".join(rows) + "\n") + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/summarize_vbench_validation25_core.py b/tools/summarize_vbench_validation25_core.py new file mode 100644 index 0000000000000000000000000000000000000000..982f9f5aa7413288ce1814a85b030d3151186f32 --- /dev/null +++ b/tools/summarize_vbench_validation25_core.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""Summarize the five VBench metrics for the validation25 core comparison.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +from collections import defaultdict +from pathlib import Path +from statistics import fmean + + +DIMENSIONS = [ + "subject_consistency", + "background_consistency", + "temporal_flickering", + "motion_smoothness", + "dynamic_degree", + "aesthetic_quality", + "imaging_quality", + "overall_consistency", +] + +GROUP_ORDER = [ + "full_ffff", + "reuse_frrf", + "stage1_fppf", + "stage1_fppp", + "timestep_onpolicy_fppf", + "timestep_onpolicy_fppp", + "chunk_timestep_fppf", + "chunk_timestep_fppp", +] + +COMPARISONS = [ + ("timestep_onpolicy_fppf", "stage1_fppf"), + ("timestep_onpolicy_fppp", "stage1_fppp"), + ("chunk_timestep_fppf", "stage1_fppf"), + ("chunk_timestep_fppp", "stage1_fppp"), +] + +ACTION_RE = re.compile(r"_action_(\d{2})$") + + +def load_dimension_result(run_root: Path, dimension: str) -> dict: + candidates = sorted((run_root / "raw" / dimension).glob("*_eval_results.json")) + if len(candidates) != 1: + raise RuntimeError( + f"Expected exactly one result JSON for {dimension}, got {len(candidates)}" + ) + payload = json.loads(candidates[0].read_text()) + if dimension not in payload: + raise KeyError(f"{dimension} missing from {candidates[0]}") + overall, records = payload[dimension] + if len(records) != 800: + raise RuntimeError(f"{dimension}: expected 800 records, got {len(records)}") + return { + "source": str(candidates[0]), + "official_overall": float(overall), + "records": records, + } + + +def normalize_score(dimension: str, value: float) -> float: + # VBench stores per-video MUSIQ scores on the 0..100 scale but divides its + # aggregate imaging_quality result by 100. + if dimension == "imaging_quality": + return float(value) / 100.0 + return float(value) + + +def parse_record(record: dict, dimension: str) -> dict: + video_path = Path(record["video_path"]) + stem = video_path.stem + if "__" not in stem: + raise ValueError(f"Missing group prefix in staged video name: {video_path}") + group, original_stem = stem.split("__", 1) + if group not in GROUP_ORDER: + raise ValueError(f"Unknown VBench group {group}: {video_path}") + action_match = ACTION_RE.search(original_stem) + if action_match is None: + raise ValueError(f"Cannot parse action id from {video_path}") + return { + "group": group, + "video_id": original_stem, + "action_id": int(action_match.group(1)), + "video_path": str(video_path), + "score": normalize_score(dimension, record["video_results"]), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run_root", type=Path, required=True) + args = parser.parse_args() + run_root = args.run_root.resolve() + + parsed_by_dimension: dict[str, list[dict]] = {} + official_overall: dict[str, float] = {} + sources: dict[str, str] = {} + for dimension in DIMENSIONS: + result = load_dimension_result(run_root, dimension) + parsed = [parse_record(record, dimension) for record in result["records"]] + parsed_by_dimension[dimension] = parsed + official_overall[dimension] = result["official_overall"] + sources[dimension] = result["source"] + + per_video: dict[str, dict[str, dict]] = defaultdict(dict) + for dimension, records in parsed_by_dimension.items(): + counts = defaultdict(int) + for record in records: + counts[record["group"]] += 1 + item = per_video[record["group"]].setdefault( + record["video_id"], + { + "action_id": record["action_id"], + "video_id": record["video_id"], + "scores": {}, + }, + ) + item["scores"][dimension] = record["score"] + for group in GROUP_ORDER: + if counts[group] != 100: + raise RuntimeError( + f"{dimension}/{group}: expected 100 records, got {counts[group]}" + ) + + group_summary = {} + for group in GROUP_ORDER: + items = list(per_video[group].values()) + if len(items) != 100: + raise RuntimeError(f"{group}: expected 100 paired videos, got {len(items)}") + mean = { + dimension: fmean(item["scores"][dimension] for item in items) + for dimension in DIMENSIONS + } + mean_by_action_id = {} + for action_id in range(4): + action_items = [item for item in items if item["action_id"] == action_id] + if len(action_items) != 25: + raise RuntimeError( + f"{group}/action{action_id}: expected 25 videos, got {len(action_items)}" + ) + mean_by_action_id[str(action_id)] = { + dimension: fmean( + item["scores"][dimension] for item in action_items + ) + for dimension in DIMENSIONS + } + group_summary[group] = { + "record_count": len(items), + "mean": mean, + "mean_by_action_id": mean_by_action_id, + } + + comparisons = {} + for candidate, baseline in COMPARISONS: + candidate_items = per_video[candidate] + baseline_items = per_video[baseline] + if set(candidate_items) != set(baseline_items): + raise RuntimeError(f"Pairing mismatch: {candidate} vs {baseline}") + key = f"{candidate}_vs_{baseline}" + comparisons[key] = { + "candidate": candidate, + "baseline": baseline, + "delta": {}, + "wins_out_of_100": {}, + } + for dimension in DIMENSIONS: + deltas = [ + candidate_items[video_id]["scores"][dimension] + - baseline_items[video_id]["scores"][dimension] + for video_id in sorted(candidate_items) + ] + comparisons[key]["delta"][dimension] = fmean(deltas) + comparisons[key]["wins_out_of_100"][dimension] = sum( + delta > 0 for delta in deltas + ) + + payload = { + "dimensions": DIMENSIONS, + "group_order": GROUP_ORDER, + "official_overall_all_800_videos": official_overall, + "raw_sources": sources, + "groups": group_summary, + "comparisons": comparisons, + } + (run_root / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + with (run_root / "summary.csv").open("w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["group", "video_count", *DIMENSIONS]) + for group in GROUP_ORDER: + writer.writerow( + [ + group, + group_summary[group]["record_count"], + *[ + f"{group_summary[group]['mean'][dimension]:.8f}" + for dimension in DIMENSIONS + ], + ] + ) + + header = ( + "| 版本 | 视频数 | Subject ↑ | Background ↑ | Flickering ↑ | " + "Motion smoothness ↑ | Dynamic degree ↑ | Aesthetic ↑ | Imaging ↑ | " + "Overall consistency ↑ |" + ) + divider = ( + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |" + ) + rows = [header, divider] + for group in GROUP_ORDER: + mean = group_summary[group]["mean"] + rows.append( + "| " + + " | ".join( + [ + group, + str(group_summary[group]["record_count"]), + f"{mean['subject_consistency']:.6f}", + f"{mean['background_consistency']:.6f}", + f"{mean['temporal_flickering']:.6f}", + f"{mean['motion_smoothness']:.6f}", + f"{mean['dynamic_degree']:.6f}", + f"{mean['aesthetic_quality']:.6f}", + f"{mean['imaging_quality']:.6f}", + f"{mean['overall_consistency']:.6f}", + ] + ) + + " |" + ) + (run_root / "summary.md").write_text("\n".join(rows) + "\n") + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/summarize_vbench_validation25_pair.py b/tools/summarize_vbench_validation25_pair.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f91c20aceb9a902645de6fc9974993318540e5 --- /dev/null +++ b/tools/summarize_vbench_validation25_pair.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +"""Summarize eight VBench dimensions for validation25 video groups.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +from collections import defaultdict +from pathlib import Path +from statistics import fmean + + +DIMENSIONS = [ + "subject_consistency", + "background_consistency", + "temporal_flickering", + "motion_smoothness", + "dynamic_degree", + "aesthetic_quality", + "imaging_quality", + "overall_consistency", +] + +ACTION_RE = re.compile(r"_action_(\d{2})$") + + +def load_dimension_result( + run_root: Path, dimension: str, expected_records: int +) -> dict: + candidates = sorted((run_root / "raw" / dimension).glob("*_eval_results.json")) + if len(candidates) != 1: + raise RuntimeError( + f"Expected exactly one result JSON for {dimension}, got {len(candidates)}" + ) + payload = json.loads(candidates[0].read_text()) + if dimension not in payload: + raise KeyError(f"{dimension} missing from {candidates[0]}") + overall, records = payload[dimension] + if len(records) != expected_records: + raise RuntimeError( + f"{dimension}: expected {expected_records} records, got {len(records)}" + ) + return { + "source": str(candidates[0]), + "official_overall": float(overall), + "records": records, + } + + +def normalize_score(dimension: str, value: float) -> float: + if dimension == "imaging_quality": + return float(value) / 100.0 + return float(value) + + +def parse_record(record: dict, dimension: str, group_order: list[str]) -> dict: + video_path = Path(record["video_path"]) + stem = video_path.stem + if "__" not in stem: + raise ValueError(f"Missing group prefix in staged video name: {video_path}") + group, original_stem = stem.split("__", 1) + if group not in group_order: + raise ValueError(f"Unknown VBench group {group}: {video_path}") + action_match = ACTION_RE.search(original_stem) + if action_match is None: + raise ValueError(f"Cannot parse action id from {video_path}") + return { + "group": group, + "video_id": original_stem, + "action_id": int(action_match.group(1)), + "video_path": str(video_path), + "score": normalize_score(dimension, record["video_results"]), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run_root", type=Path, required=True) + parser.add_argument("--group", action="append", required=True) + parser.add_argument("--videos_per_group", type=int, default=100) + args = parser.parse_args() + run_root = args.run_root.resolve() + group_order = list(args.group) + if not group_order or len(set(group_order)) != len(group_order): + raise ValueError("Pass one or more distinct --group values") + expected_records = len(group_order) * args.videos_per_group + + parsed_by_dimension: dict[str, list[dict]] = {} + official_overall: dict[str, float] = {} + sources: dict[str, str] = {} + for dimension in DIMENSIONS: + result = load_dimension_result(run_root, dimension, expected_records) + parsed_by_dimension[dimension] = [ + parse_record(record, dimension, group_order) + for record in result["records"] + ] + official_overall[dimension] = result["official_overall"] + sources[dimension] = result["source"] + + per_video: dict[str, dict[str, dict]] = defaultdict(dict) + for dimension, records in parsed_by_dimension.items(): + counts = defaultdict(int) + for record in records: + counts[record["group"]] += 1 + item = per_video[record["group"]].setdefault( + record["video_id"], + { + "action_id": record["action_id"], + "video_id": record["video_id"], + "scores": {}, + }, + ) + item["scores"][dimension] = record["score"] + for group in group_order: + if counts[group] != args.videos_per_group: + raise RuntimeError( + f"{dimension}/{group}: expected {args.videos_per_group} " + f"records, got {counts[group]}" + ) + + group_summary = {} + for group in group_order: + items = list(per_video[group].values()) + if len(items) != args.videos_per_group: + raise RuntimeError( + f"{group}: expected {args.videos_per_group} videos, got {len(items)}" + ) + mean = { + dimension: fmean(item["scores"][dimension] for item in items) + for dimension in DIMENSIONS + } + mean_by_action_id = {} + for action_id in range(4): + action_items = [item for item in items if item["action_id"] == action_id] + if len(action_items) * 4 != args.videos_per_group: + raise RuntimeError( + f"{group}/action{action_id}: uneven validation action count " + f"{len(action_items)}" + ) + mean_by_action_id[str(action_id)] = { + dimension: fmean(item["scores"][dimension] for item in action_items) + for dimension in DIMENSIONS + } + group_summary[group] = { + "record_count": len(items), + "mean": mean, + "mean_by_action_id": mean_by_action_id, + } + + summary = { + "dimensions": DIMENSIONS, + "group_order": group_order, + "official_overall_all_videos": official_overall, + "raw_sources": sources, + "groups": group_summary, + } + (run_root / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n" + ) + + with (run_root / "summary.csv").open("w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["group", "video_count", *DIMENSIONS]) + for group in group_order: + writer.writerow( + [ + group, + group_summary[group]["record_count"], + *[ + f"{group_summary[group]['mean'][dimension]:.8f}" + for dimension in DIMENSIONS + ], + ] + ) + + rows = [ + ( + "| 版本 | 视频数 | Subject ↑ | Background ↑ | Flickering ↑ | " + "Motion smoothness ↑ | Dynamic degree ↑ | Aesthetic ↑ | Imaging ↑ | " + "Overall consistency ↑ |" + ), + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for group in group_order: + mean = group_summary[group]["mean"] + rows.append( + "| " + + " | ".join( + [ + group, + str(group_summary[group]["record_count"]), + f"{mean['subject_consistency']:.6f}", + f"{mean['background_consistency']:.6f}", + f"{mean['temporal_flickering']:.6f}", + f"{mean['motion_smoothness']:.6f}", + f"{mean['dynamic_degree']:.6f}", + f"{mean['aesthetic_quality']:.6f}", + f"{mean['imaging_quality']:.6f}", + f"{mean['overall_consistency']:.6f}", + ] + ) + + " |" + ) + (run_root / "summary.md").write_text("\n".join(rows) + "\n") + print(json.dumps(summary, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/supervise_predictor_v4_fp32master_training.sh b/tools/supervise_predictor_v4_fp32master_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..f67be4ab157f64a95213bfc54a42de84441ec93e --- /dev/null +++ b/tools/supervise_predictor_v4_fp32master_training.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +output_dir="${repo_root}/checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab" +final_weights="${output_dir}/predictor_step_02000.safetensors" +resume_path="${output_dir}/training_latest.pt" +lock_path="${output_dir}/supervisor.lock" +mkdir -p "${output_dir}" + +exec 9>"${lock_path}" +if ! flock -n 9; then + echo "[supervisor] another FP32-master v4 supervisor is already active" + exit 0 +fi + +while [[ ! -f "${final_weights}" ]]; do + if pgrep -f "python -u ${repo_root}/train_predictor_v4.py .*--output_dir ${output_dir}" >/dev/null; then + sleep 60 + continue + fi + + echo "[supervisor] FP32-master v4 worker is absent; restarting at $(date -u +%FT%TZ)" + if [[ -f "${resume_path}" ]]; then + bash "${repo_root}/tools/launch_predictor_v4_fp32master_training.sh" \ + "${output_dir}" "${resume_path}" || true + else + bash "${repo_root}/tools/launch_predictor_v4_fp32master_training.sh" \ + "${output_dir}" || true + fi + sleep 30 +done + +echo "[supervisor] step 2000 checkpoint is complete: ${final_weights}" diff --git a/tools/supervise_predictor_v4_rollout_fppf_overfit8_training.sh b/tools/supervise_predictor_v4_rollout_fppf_overfit8_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..3de4d92d05db67d8758754c999a20788e104d503 --- /dev/null +++ b/tools/supervise_predictor_v4_rollout_fppf_overfit8_training.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +output_dir="${repo_root}/checkpoints/predictor_v4_rollout_fppf_blocks1-52_overfit8_seed20260721_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps300_swarmup8_bwarmup8_swanlab" +final_weights="${output_dir}/predictor_step_00300.safetensors" +resume_path="${output_dir}/training_latest.pt" +lock_path="${output_dir}/supervisor.lock" +mkdir -p "${output_dir}" + +exec 9>"${lock_path}" +if ! flock -n 9; then + echo "[supervisor] another FPPF rollout overfit supervisor is already active" + exit 0 +fi + +while [[ ! -f "${final_weights}" ]]; do + if pgrep -f "train_predictor_v4_rollout.py .*--output_dir ${output_dir} .*--rollout_schedule fppf" >/dev/null; then + sleep 60 + continue + fi + + echo "[supervisor] FPPF rollout worker is absent; restarting at $(date -u +%FT%TZ)" + if [[ -f "${resume_path}" ]]; then + bash "${repo_root}/tools/launch_predictor_v4_rollout_fppf_overfit8_training.sh" \ + "${output_dir}" "${resume_path}" 300 || true + else + bash "${repo_root}/tools/launch_predictor_v4_rollout_fppf_overfit8_training.sh" \ + "${output_dir}" "" 300 || true + fi + sleep 30 +done + +echo "[supervisor] FPPF step 300 checkpoint is complete: ${final_weights}" diff --git a/tools/supervise_predictor_v4_rollout_overfit8_training.sh b/tools/supervise_predictor_v4_rollout_overfit8_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..0639ae8f23f84b7e43fe3adb96afb3599a8e0096 --- /dev/null +++ b/tools/supervise_predictor_v4_rollout_overfit8_training.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +output_dir="${repo_root}/checkpoints/predictor_v4_rollout_fppp_blocks1-52_overfit8_seed20260721_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup8_bwarmup8_swanlab" +final_weights="${output_dir}/predictor_step_02000.safetensors" +resume_path="${output_dir}/training_latest.pt" +lock_path="${output_dir}/supervisor.lock" +mkdir -p "${output_dir}" + +exec 9>"${lock_path}" +if ! flock -n 9; then + echo "[supervisor] another rollout overfit supervisor is already active" + exit 0 +fi + +while [[ ! -f "${final_weights}" ]]; do + if pgrep -f "python -u ${repo_root}/train_predictor_v4_rollout.py .*--output_dir ${output_dir}" >/dev/null; then + sleep 60 + continue + fi + + echo "[supervisor] rollout overfit worker is absent; restarting at $(date -u +%FT%TZ)" + if [[ -f "${resume_path}" ]]; then + bash "${repo_root}/tools/launch_predictor_v4_rollout_overfit8_training.sh" \ + "${output_dir}" "${resume_path}" || true + else + bash "${repo_root}/tools/launch_predictor_v4_rollout_overfit8_training.sh" \ + "${output_dir}" || true + fi + sleep 30 +done + +echo "[supervisor] step 2000 checkpoint is complete: ${final_weights}" diff --git a/tools/supervise_predictor_v4_rollout_training.sh b/tools/supervise_predictor_v4_rollout_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..24e5d8c94d8d92fe030a17fe079e244395f846c1 --- /dev/null +++ b/tools/supervise_predictor_v4_rollout_training.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +output_dir="${repo_root}/checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab" +final_weights="${output_dir}/predictor_step_02000.safetensors" +resume_path="${output_dir}/training_latest.pt" +lock_path="${output_dir}/supervisor.lock" +mkdir -p "${output_dir}" + +exec 9>"${lock_path}" +if ! flock -n 9; then + echo "[supervisor] another rollout supervisor is already active" + exit 0 +fi + +while [[ ! -f "${final_weights}" ]]; do + if pgrep -f "python -u ${repo_root}/train_predictor_v4_rollout.py .*--output_dir ${output_dir}" >/dev/null; then + sleep 60 + continue + fi + + echo "[supervisor] rollout worker is absent; restarting at $(date -u +%FT%TZ)" + if [[ -f "${resume_path}" ]]; then + bash "${repo_root}/tools/launch_predictor_v4_rollout_training.sh" \ + "${output_dir}" "${resume_path}" || true + else + bash "${repo_root}/tools/launch_predictor_v4_rollout_training.sh" \ + "${output_dir}" || true + fi + sleep 30 +done + +echo "[supervisor] step 2000 checkpoint is complete: ${final_weights}" diff --git a/tools/supervise_previous_vbench_validation_predictors.sh b/tools/supervise_previous_vbench_validation_predictors.sh new file mode 100644 index 0000000000000000000000000000000000000000..d10b2fbc6b010e432e274d0cb0e6149fd777d7a0 --- /dev/null +++ b/tools/supervise_previous_vbench_validation_predictors.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor" +dataset_root="${PREDICTOR_HOLDOUT_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_vbench_val25_test50}" +eval_base="${dataset_root}/evaluations" +lock_path="${eval_base}/previous_versions_validation.lock" +complete_path="${eval_base}/previous_versions_validation.completed" +mkdir -p "${eval_base}" + +exec 9>"${lock_path}" +if ! flock -n 9; then + echo "[supervisor] another previous-version validation supervisor is active" + exit 0 +fi + +run_until_complete() { + local eval_name="$1" + local weights="$2" + local source_blocks="$3" + local predictor_variant="$4" + local marker="${eval_base}/${eval_name}/sync.completed" + while [[ ! -f "${marker}" ]]; do + echo "[supervisor] starting ${eval_name} at $(date -u +%FT%TZ)" + if ! bash "${repo_root}/tools/launch_vbench_validation_predictor.sh" \ + "${eval_name}" "${weights}" "${source_blocks}" "${predictor_variant}"; then + echo "[supervisor] ${eval_name} failed; retrying in 30 seconds" + sleep 30 + fi + done + echo "[supervisor] completed ${eval_name} at $(date -u +%FT%TZ)" +} + +run_until_complete \ + eval_validation25_v4_blocks0-53_bf16_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100 \ + "${repo_root}/checkpoints/predictor_v4_prefeature_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" \ + 0,53 v4 + +run_until_complete \ + eval_validation25_v4_blocks1-52_bf16_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100 \ + "${repo_root}/checkpoints/predictor_v4_prefeature_blocks1-52_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" \ + 1,52 v4 + +run_until_complete \ + eval_validation25_disca_blocks0-53_bf16_ckptstep2000_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100 \ + "${repo_root}/checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/predictor_step_02000.safetensors" \ + 0,53 disca + +date -u +%FT%TZ >"${complete_path}" +echo "[supervisor] all previous-version validation evaluations are complete" diff --git a/tools/supervise_training_tests_then_helios.sh b/tools/supervise_training_tests_then_helios.sh new file mode 100644 index 0000000000000000000000000000000000000000..a52d8b6b173d7f2fb7c83c6b34aa34719b42b66f --- /dev/null +++ b/tools/supervise_training_tests_then_helios.sh @@ -0,0 +1,244 @@ +#!/usr/bin/env bash +set -uo pipefail + +ROOT=/data3/chenzhuo/workspace/HY-WorldPlay-DEV-Predictor +HELIOS=/data3/chenzhuo/workspace/Helios +PYTHON=/data3/chenzhuo/anaconda3/envs/worldplay/bin/python +SUPERVISOR_LOG="$ROOT/outputs/supervise_training_tests_then_helios.log" + +V2_TRAIN_PID=4073736 +V3_TRAIN_PID=3092818 +V2_WEIGHTS="$ROOT/checkpoints/predictor_v2_1_52/predictor_step_02000.safetensors" +V3_WEIGHTS="$ROOT/checkpoints/predictor_v3_0_53/predictor_step_02000.safetensors" + +log() { + printf '[%s] %s\n' "$(date '+%F %T')" "$*" +} + +latest_step() { + local path=$1 + if [[ -s "$path" ]]; then + tail -n 1 "$path" | sed -n 's/.*"global_step"[ ]*:[ ]*\([0-9]*\).*/\1/p' + else + printf '?' + fi +} + +wait_for_training() { + local label=$1 + local pid=$2 + local train_log=$3 + local weights=$4 + while kill -0 "$pid" 2>/dev/null; do + log "$label is active: step=$(latest_step "$train_log")/2000 pid=$pid" + sleep 60 + done + if [[ ! -s "$weights" ]]; then + log "ERROR: $label exited without final weights: $weights" + return 1 + fi + log "$label completed and final weights exist: $weights" +} + +validate_eval() { + local output_dir=$1 + local expected_runs=$2 + "$PYTHON" - "$output_dir" "$expected_runs" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +expected = int(sys.argv[2]) +payload = json.loads((root / "metrics.json").read_text()) +runs = payload.get("runs", []) +if len(runs) != expected: + raise SystemExit(f"expected {expected} runs, found {len(runs)}") +missing = [run.get("video_path") for run in runs if not Path(run["video_path"]).is_file()] +if missing: + raise SystemExit(f"missing videos: {missing[:3]}") +print(f"validated {len(runs)} runs and videos in {root}") +PY +} + +adopt_or_validate_eval() { + local label=$1 + local output_dir=$2 + local expected_runs=$3 + if validate_eval "$output_dir" "$expected_runs" >/dev/null 2>&1; then + log "$label is already complete: $output_dir" + return 0 + fi + + local pid + pid=$(pgrep -f "eval_predictor_v2_rollout.py.*--output_dir $output_dir" | head -n 1 || true) + if [[ -z "$pid" ]]; then + return 1 + fi + log "adopting active $label process: pid=$pid output=$output_dir" + while kill -0 "$pid" 2>/dev/null; do + log "$label is active: pid=$pid" + sleep 60 + done + validate_eval "$output_dir" "$expected_runs" || return 1 + log "$label completed: $output_dir" +} + +run_v2_test_after_training() { + wait_for_training \ + 'v2 [1,52] GPU4 training' \ + "$V2_TRAIN_PID" \ + "$ROOT/checkpoints/predictor_v2_1_52/train_log.jsonl" \ + "$V2_WEIGHTS" || return 1 + + local output_dir="$ROOT/outputs/predictor_v2/online_1_52_step2000" + local eval_log="$ROOT/outputs/predictor_v2_online_1_52_step2000.log" + if adopt_or_validate_eval 'v2 [1,52] test' "$output_dir" 30; then + return 0 + fi + log "starting v2 [1,52] 10-case Full/Reuse/Predictor test on GPU4" + ( + cd "$ROOT" || exit 1 + unset HTTP_PROXY HTTPS_PROXY ALL_PROXY http_proxy https_proxy all_proxy + export CUDA_VISIBLE_DEVICES=4 PYTHONPATH=. + "$PYTHON" -u eval_predictor_v2_rollout.py \ + --predictor_version v2 \ + --weights_1_52 "$V2_WEIGHTS" \ + --modes full,reuse,predictor_1_52 \ + --case_ids 1-10 \ + --video_frames 253 \ + --output_dir "$output_dir" + ) >"$eval_log" 2>&1 + local status=$? + if (( status != 0 )); then + log "ERROR: v2 [1,52] test failed with status $status; log=$eval_log" + return "$status" + fi + validate_eval "$output_dir" 30 || return 1 + log "v2 [1,52] test completed: $output_dir" +} + +run_v3_test_after_training() { + wait_for_training \ + 'v3 [0,53] GPU5 training' \ + "$V3_TRAIN_PID" \ + "$ROOT/checkpoints/predictor_v3_0_53/train_log.jsonl" \ + "$V3_WEIGHTS" || return 1 + + local smoke_dir="$ROOT/outputs/predictor_v3/online_0_53_step2000_smoke" + local smoke_log="$ROOT/outputs/predictor_v3_online_0_53_step2000_smoke.log" + local status=0 + if ! adopt_or_validate_eval 'v3 [0,53] smoke' "$smoke_dir" 2; then + log "starting v3 [0,53] two-chunk online smoke on GPU5" + ( + cd "$ROOT" || exit 1 + unset HTTP_PROXY HTTPS_PROXY ALL_PROXY http_proxy https_proxy all_proxy + export CUDA_VISIBLE_DEVICES=5 PYTHONPATH=. + "$PYTHON" -u eval_predictor_v2_rollout.py \ + --predictor_version v3 \ + --weights_0_53 "$V3_WEIGHTS" \ + --modes full,predictor_0_53 \ + --case_ids 1 \ + --video_frames 29 \ + --output_dir "$smoke_dir" + ) >"$smoke_log" 2>&1 + status=$? + if (( status != 0 )); then + log "ERROR: v3 [0,53] smoke failed with status $status; log=$smoke_log" + return "$status" + fi + validate_eval "$smoke_dir" 2 || return 1 + fi + log "v3 smoke passed; starting 10-case Full/Reuse/Predictor test on GPU5" + + local output_dir="$ROOT/outputs/predictor_v3/online_0_53_step2000" + local eval_log="$ROOT/outputs/predictor_v3_online_0_53_step2000.log" + if adopt_or_validate_eval 'v3 [0,53] full test' "$output_dir" 30; then + return 0 + fi + ( + cd "$ROOT" || exit 1 + unset HTTP_PROXY HTTPS_PROXY ALL_PROXY http_proxy https_proxy all_proxy + export CUDA_VISIBLE_DEVICES=5 PYTHONPATH=. + "$PYTHON" -u eval_predictor_v2_rollout.py \ + --predictor_version v3 \ + --weights_0_53 "$V3_WEIGHTS" \ + --modes full,reuse,predictor_0_53 \ + --case_ids 1-10 \ + --video_frames 253 \ + --output_dir "$output_dir" + ) >"$eval_log" 2>&1 + status=$? + if (( status != 0 )); then + log "ERROR: v3 [0,53] test failed with status $status; log=$eval_log" + return "$status" + fi + validate_eval "$output_dir" 30 || return 1 + log "v3 [0,53] test completed: $output_dir" +} + +wait_for_existing_v2_0_53_test() { + local pid=1215665 + local output_dir="$ROOT/outputs/predictor_v2/online_0_53_step2000" + while kill -0 "$pid" 2>/dev/null; do + log "v2 [0,53] GPU3 test is active: pid=$pid" + sleep 60 + done + validate_eval "$output_dir" 30 || { + log "ERROR: v2 [0,53] GPU3 test did not complete successfully" + return 1 + } + log "v2 [0,53] GPU3 test completed" +} + +launch_helios() { + local gpu=$1 + local stamp=$2 + local output_root="$HELIOS/output_helios/moviebench_990_gpu${gpu}_post_predictor_${stamp}" + local helios_log="$HELIOS/logs/run_gpu${gpu}_post_predictor_${stamp}.log" + ( + cd "$HELIOS" || exit 1 + unset HTTP_PROXY HTTPS_PROXY ALL_PROXY http_proxy https_proxy all_proxy + export GPU="$gpu" OUTPUT_ROOT="$output_root" + exec bash "$HELIOS/run.sh" + ) >"$helios_log" 2>&1 < /dev/null & + local pid=$! + log "launched Helios/run.sh on GPU$gpu: pid=$pid log=$helios_log" >&2 + printf '%s' "$pid" +} + +main() { + mkdir -p "$ROOT/outputs" "$HELIOS/logs" + log "supervisor started" + + run_v2_test_after_training & + local worker4=$! + run_v3_test_after_training & + local worker5=$! + + wait "$worker4" + local status4=$? + wait "$worker5" + local status5=$? + wait_for_existing_v2_0_53_test + local status3=$? + + if (( status3 != 0 || status4 != 0 || status5 != 0 )); then + log "ERROR: at least one required test failed; Helios will not be launched (gpu3=$status3 gpu4=$status4 gpu5=$status5)" + return 1 + fi + + local stamp + stamp=$(date '+%Y%m%d_%H%M%S') + local helios4 helios5 + helios4=$(launch_helios 4 "$stamp") + helios5=$(launch_helios 5 "$stamp") + sleep 60 + if ! kill -0 "$helios4" 2>/dev/null || ! kill -0 "$helios5" 2>/dev/null; then + log "ERROR: a Helios process exited during startup (gpu4=$helios4 gpu5=$helios5)" + return 1 + fi + log "Helios processes are active on GPU4/GPU5; supervisor finished successfully" +} + +main >>"$SUPERVISOR_LOG" 2>&1 diff --git a/tools/sync_predictor_prefeature_to_s3.sh b/tools/sync_predictor_prefeature_to_s3.sh new file mode 100644 index 0000000000000000000000000000000000000000..d9cbf947acbb3419b6ea24fba861e46d89ec1cb0 --- /dev/null +++ b/tools/sync_predictor_prefeature_to_s3.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +export BUCKET="${BUCKET:-s3-us-west2-default}" +export FS_ID="${FS_ID:-fs-05c60be7352499179}" +export MOUNT_POINT="${MOUNT_POINT:-/mnt/s3files/s3-us-west2-default}" + +dataset_name="hyworldplay_predictor_prefeature_v4" +nvme_root="${PREDICTOR_DATA_NVME_ROOT:-/mnt/local_nvme/zoubin/cz/${dataset_name}}" +s3_uri="${PREDICTOR_DATA_S3_URI:-s3://${BUCKET}/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor/datasets/${dataset_name}}" + +[[ -f "${nvme_root}/manifest.jsonl" ]] || { + echo "ERROR: missing merged manifest in ${nvme_root}" >&2 + exit 1 +} +[[ $(wc -l < "${nvme_root}/manifest.jsonl") -eq 3200 ]] || { + echo "ERROR: manifest.jsonl must contain 3200 records" >&2 + exit 1 +} +[[ $(wc -l < "${nvme_root}/train_eval_manifest.jsonl") -eq 2800 ]] || { + echo "ERROR: train_eval_manifest.jsonl must contain 2800 records" >&2 + exit 1 +} +[[ $(find "${nvme_root}/cases" -type f -name '*.safetensors' | wc -l) -eq 100 ]] || { + echo "ERROR: expected 100 case tensor files" >&2 + exit 1 +} +[[ $(find "${nvme_root}/steps" -type f -name '*.safetensors' | wc -l) -eq 3200 ]] || { + echo "ERROR: expected 3200 step tensor files" >&2 + exit 1 +} +[[ $(find "${nvme_root}/context_prefeature" -type f -name '*.safetensors' | wc -l) -eq 12800 ]] || { + echo "ERROR: expected 12800 Context prefeature files" >&2 + exit 1 +} +[[ $(find "${nvme_root}/videos" -type f -name '*.mp4' | wc -l) -eq 400 ]] || { + echo "ERROR: expected 400 evaluation videos" >&2 + exit 1 +} + +echo "[dataset] validated complete NVMe build: ${nvme_root}" +echo "[dataset] syncing ${nvme_root} -> ${s3_uri}" +aws s3 sync "${nvme_root}" "${s3_uri}" --only-show-errors +echo "[dataset] durable S3 synchronization completed" diff --git a/tools/train_evaluate_predictor_confidence.py b/tools/train_evaluate_predictor_confidence.py new file mode 100644 index 0000000000000000000000000000000000000000..264ca802d39fd8f7227392be0f5e57a4c2d5da12 --- /dev/null +++ b/tools/train_evaluate_predictor_confidence.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""Train C0/C1/C2 confidence heads and report held-out risk coverage.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +import random +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from scipy.stats import pearsonr, spearmanr +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader, TensorDataset +from torch.utils.data.distributed import DistributedSampler + +from models.confidence_head import PredictorConfidenceHead + + +TARGETS = ("c0", "c1", "c2") +RATIOS = (0.0, 0.05, 0.1, 0.2, 0.25, 1 / 3, 0.4, 0.5, 0.6, 0.75, 1.0) + + +def load_features(directory): + shards = [torch.load(path, map_location="cpu", weights_only=False) for path in sorted(directory.glob("features_rank_*.pt"))] + if not shards: + raise FileNotFoundError(f"No feature shards in {directory}") + names = ("pooled_hidden", "log_residual", "true_c0", "true_c1", "true_c2", "scheduler_coeff", "step_id") + result = {name: torch.cat([shard[name] for shard in shards]) for name in names} + result["metadata"] = sum((shard["metadata"] for shard in shards), []) + order = sorted(range(len(result["metadata"])), key=lambda i: tuple(result["metadata"][i][k] for k in ("case_id", "action_id", "chunk_id", "step_id"))) + index = torch.tensor(order) + for name in names: + result[name] = result[name][index] + result["metadata"] = [result["metadata"][i] for i in order] + return result + + +def train_head( + features, target_name, device, epochs, batch_size, seed, rank, world_size, + mlp_dims, +): + torch.manual_seed(seed) + model = PredictorConfidenceHead( + features["pooled_hidden"].shape[1], mlp_dims=mlp_dims + ).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) + dataset = TensorDataset( + features["pooled_hidden"].float(), features["step_id"].long(), + features["log_residual"].float(), + torch.log(features[f"true_{target_name}"].float() + 1e-8), + ) + sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=seed) if world_size > 1 else None + loader = DataLoader(dataset, batch_size=batch_size, shuffle=sampler is None, sampler=sampler) + wrapped = DDP(model, device_ids=[device.index], broadcast_buffers=False) if world_size > 1 else model + wrapped.train() + for epoch in range(epochs): + if sampler is not None: + sampler.set_epoch(epoch) + total = count = 0 + for pooled, step, residual, target in loader: + pooled, step, residual, target = (x.to(device) for x in (pooled, step, residual, target)) + optimizer.zero_grad(set_to_none=True) + loss = F.smooth_l1_loss(wrapped(pooled, step, residual), target) + loss.backward() + optimizer.step() + total += float(loss.detach()) * len(pooled) + count += len(pooled) + if rank == 0 and (epoch == 0 or (epoch + 1) % 10 == 0): + print(f"{target_name} epoch {epoch + 1}/{epochs} loss={total / count:.6f}", flush=True) + if world_size > 1: + dist.barrier() + return model.eval() + + +@torch.inference_mode() +def predict(model, features, device): + result = [] + for start in range(0, len(features["step_id"]), 512): + stop = start + 512 + result.append(model( + features["pooled_hidden"][start:stop].float().to(device), + features["step_id"][start:stop].long().to(device), + features["log_residual"][start:stop].float().to(device), + ).cpu()) + return torch.cat(result).numpy() + + +def correlation(x, y, function): + if len(x) < 2 or np.all(x == x[0]) or np.all(y == y[0]): + return float("nan") + return float(function(x, y).statistic) + + +def metric_rows(split, predictions, features): + rows = [] + steps = features["step_id"].numpy() + for target, score in predictions.items(): + truth = np.log(features[f"true_{target}"].numpy() + 1e-8) + groups = [("overall", np.ones(len(steps), dtype=bool))] + [(f"P{s}", steps == s) for s in (1, 2, 3)] + for label, mask in groups: + rows.append({ + "model": f"Confidence {target.upper()}", "split": split, "step": label, + "spearman": correlation(score[mask], truth[mask], spearmanr), + "pearson": correlation(score[mask], truth[mask], pearsonr), + "log_mae": float(np.abs(score[mask] - truth[mask]).mean()), + }) + return rows + + +def remaining_risk(risk, score, ratio): + count = min(len(risk), int(round(ratio * len(risk)))) + if count == 0: + return float(risk.mean()) + selected = np.argpartition(score, len(score) - count)[-count:] + keep = np.ones(len(risk), dtype=bool) + keep[selected] = False + return float(np.where(keep, risk, 0.0).mean()) + + +def risk_rows(train, validation, predictions, seed, risk_target): + risk = validation[f"true_{risk_target}"].numpy() + train_steps = train["step_id"].numpy() + train_risk = train[f"true_{risk_target}"].numpy() + means = {step: train_risk[train_steps == step].mean() for step in (1, 2, 3)} + methods = { + "Timestep-only": np.array([means[int(s)] for s in validation["step_id"]]), + **{ + f"Confidence {target.upper()}": score + for target, score in predictions.items() + }, + "Oracle": risk, + } + rows = [] + rng = np.random.default_rng(seed) + for ratio in RATIOS: + count = min(len(risk), int(round(ratio * len(risk)))) + random_risks = [] + for _ in range(100): + chosen = rng.choice(len(risk), count, replace=False) if count else [] + keep = np.ones(len(risk), dtype=bool) + keep[chosen] = False + random_risks.append(float(np.where(keep, risk, 0.0).mean())) + rows.append({"method": "Random", "fallback_ratio": ratio, "remaining_risk": np.mean(random_risks), "std_if_random": np.std(random_risks)}) + for method, score in methods.items(): + rows.append({"method": method, "fallback_ratio": ratio, "remaining_risk": remaining_risk(risk, score, ratio), "std_if_random": ""}) + return rows + + +def write_csv(path, rows): + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def make_plots(validation, risks, predictions, output, risk_target): + output.mkdir(parents=True, exist_ok=True) + plt.figure(figsize=(7.2, 5.2)) + confidence_method = f"Confidence {risk_target.upper()}" + for method in ("Random", "Timestep-only", confidence_method, "Oracle"): + rows = [row for row in risks if row["method"] == method] + plt.plot([100 * row["fallback_ratio"] for row in rows], [row["remaining_risk"] for row in rows], marker="o", label=method) + plt.xlabel("Fallback ratio (%)") + plt.ylabel(f"Remaining {risk_target.upper()} risk") + plt.legend() + plt.tight_layout() + plt.savefig(output / "risk_coverage.png", dpi=180) + plt.close() + scores = predictions[risk_target] + count = int(round(len(scores) / 3)) + chosen = np.argpartition(scores, len(scores) - count)[-count:] + steps = validation["step_id"].numpy()[chosen] + plt.figure(figsize=(5.5, 4.4)) + plt.bar(("P1", "P2", "P3"), [float((steps == s).mean()) for s in (1, 2, 3)]) + plt.ylabel("Fraction of selected fallbacks") + plt.tight_layout() + plt.savefig(output / "fallback_distribution_33pct.png", dpi=180) + plt.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--train_features", type=Path, required=True) + parser.add_argument("--validation_features", type=Path, required=True) + parser.add_argument("--output_dir", type=Path, required=True) + parser.add_argument("--epochs", type=int, default=40) + parser.add_argument("--batch_size", type=int, default=256) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--targets", nargs="+", choices=TARGETS, default=list(TARGETS)) + parser.add_argument("--risk_target", choices=TARGETS, default="c1") + parser.add_argument("--mlp_dims", nargs="+", type=int, default=[256]) + args = parser.parse_args() + if args.risk_target not in args.targets: + parser.error("--risk_target must be included in --targets") + if any(dim <= 0 for dim in args.mlp_dims): + parser.error("--mlp_dims values must be positive") + targets = tuple(dict.fromkeys(args.targets)) + mlp_dims = tuple(args.mlp_dims) + random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed) + world_size = int(os.environ.get("WORLD_SIZE", 1)); rank = int(os.environ.get("RANK", 0)); local_rank = int(os.environ.get("LOCAL_RANK", 0)) + if world_size > 1: + torch.cuda.set_device(local_rank); dist.init_process_group("nccl") + device = torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu") + train = load_features(args.train_features); validation = load_features(args.validation_features) + if rank == 0: + for child in ("checkpoints", "plots", "predictions", "tables"): + (args.output_dir / child).mkdir(parents=True, exist_ok=True) + if world_size > 1: dist.barrier() + train_predictions = {}; validation_predictions = {} + for target in targets: + model = train_head( + train, target, device, args.epochs, args.batch_size, args.seed, + rank, world_size, mlp_dims, + ) + if rank == 0: + torch.save( + { + "model": model.state_dict(), + "target": target, + "hidden_dim": train["pooled_hidden"].shape[1], + "timestep_dim": 32, + "mlp_dims": list(mlp_dims), + }, + args.output_dir / "checkpoints" / f"confidence_{target}.pt", + ) + train_predictions[target] = predict(model, train, device) + validation_predictions[target] = predict(model, validation, device) + if world_size > 1: dist.barrier() + if rank != 0: + dist.destroy_process_group(); return + metrics = metric_rows("train", train_predictions, train) + metric_rows("validation", validation_predictions, validation) + risks = risk_rows( + train, validation, validation_predictions, args.seed, args.risk_target + ) + write_csv(args.output_dir / "tables" / "confidence_metrics.csv", metrics) + write_csv(args.output_dir / "tables" / "risk_coverage.csv", risks) + prediction_rows = [{**item, **{f"pred_{target}": float(validation_predictions[target][i]) for target in targets}, **{f"true_{target}": float(validation[f"true_{target}"][i]) for target in targets}} for i, item in enumerate(validation["metadata"])] + write_csv(args.output_dir / "predictions" / "validation_predictions.csv", prediction_rows) + make_plots( + validation, risks, validation_predictions, args.output_dir / "plots", + args.risk_target, + ) + at_33 = {row["method"]: float(row["remaining_risk"]) for row in risks if math.isclose(row["fallback_ratio"], 1 / 3)} + confidence_method = f"Confidence {args.risk_target.upper()}" + target_metrics = {row["step"]: row for row in metrics if row["split"] == "validation" and row["model"] == confidence_method} + scores = validation_predictions[args.risk_target]; count = int(round(len(scores) / 3)); chosen = np.argpartition(scores, len(scores) - count)[-count:] + chosen_mask = np.zeros(len(scores), dtype=bool); chosen_mask[chosen] = True; steps = validation["step_id"].numpy() + summary = { + "protocol": "WorldPlay fulltrain100 (8400 transitions); independent validation25 (2100 transitions); no validation tuning", + "targets": list(targets), "risk_target": args.risk_target, + "mlp_dims": list(mlp_dims), + "train_samples": len(train["step_id"]), "validation_samples": len(validation["step_id"]), + "risk_target_validation_metrics": target_metrics, + "remaining_risk_at_33pct": at_33, + "gain_vs_random": (at_33["Random"] - at_33[confidence_method]) / at_33["Random"], + "oracle_gap": (at_33[confidence_method] - at_33["Oracle"]) / max(at_33["Oracle"], 1e-12), + "routing_at_33pct": {f"P{s}": {"selection_fraction": float((steps[chosen_mask] == s).mean()), "within_step_fallback_rate": float(chosen_mask[steps == s].mean())} for s in (1, 2, 3)}, + } + if args.risk_target == "c1": + summary["c1_validation_metrics"] = target_metrics + (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") + (args.output_dir / "summary.md").write_text( + "# HY-WorldPlay offline confidence validation\n\n" + summary["protocol"] + "\n\n" + + f"## {args.risk_target.upper()} validation Spearman\n\n" + "\n".join(f"- {key}: {value['spearman']:.4f}" for key, value in target_metrics.items()) + + f"\n\n## 33.3% fallback remaining {args.risk_target.upper()} risk\n\n" + "\n".join(f"- {key}: {value:.8g}" for key, value in at_33.items()) + "\n" + ) + print(json.dumps(summary, indent=2), flush=True) + if world_size > 1: dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/validate_predictor_dataset.py b/tools/validate_predictor_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..6d4fb2a485b75d04b934d8d3e942df807bc6813a --- /dev/null +++ b/tools/validate_predictor_dataset.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""Validate Predictor safetensors files without loading the Teacher.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from safetensors.torch import load_file + +from predictor_data.schema import SCHEMA_VERSION, validate_case_tensors, validate_chunk_tensors + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", default="datasets/predictor_v1") + parser.add_argument("--expected_cases", type=int, default=None) + parser.add_argument("--expected_chunks_per_case", type=int, default=None) + args = parser.parse_args() + + root = Path(args.root).resolve() + manifest = root / "manifest.jsonl" + if not manifest.is_file(): + raise FileNotFoundError(manifest) + records = [json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines() if line] + if len({(r["case_id"], r["seed"], r["chunk_id"]) for r in records}) != len(records): + raise ValueError("Duplicate manifest keys") + + seen_cases = set() + per_case: dict[int, int] = {} + for record in records: + if record["schema_version"] != SCHEMA_VERSION: + raise ValueError(f"Unexpected schema: {record['schema_version']}") + case_id = int(record["case_id"]) + if case_id not in seen_cases: + validate_case_tensors(load_file(str(root / record["case_tensor_file"]), device="cpu")) + seen_cases.add(case_id) + validate_chunk_tensors(load_file(str(root / record["tensor_file"]), device="cpu")) + per_case[case_id] = per_case.get(case_id, 0) + 1 + + if args.expected_cases is not None and len(seen_cases) != args.expected_cases: + raise ValueError(f"Expected {args.expected_cases} cases, found {len(seen_cases)}") + if args.expected_chunks_per_case is not None: + bad = {case: count for case, count in per_case.items() if count != args.expected_chunks_per_case} + if bad: + raise ValueError(f"Unexpected chunks per case: {bad}") + print(json.dumps({"records": len(records), "cases": len(seen_cases), "chunks_per_case": per_case}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validate_predictor_prefeature_dataset.py b/tools/validate_predictor_prefeature_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..10f41e5e1b22909dfdd076fc5bc68ef612400612 --- /dev/null +++ b/tools/validate_predictor_prefeature_dataset.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""Validate BF16 prefeature files, targets, task cardinality, and videos.""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter +from pathlib import Path + +from safetensors.torch import load_file + +from predictor_data.prefeature_schema import ( + CONTEXT_BLOCK_IDS, + SCHEMA_VERSION_PREFEATURE, + validate_case_tensors, + validate_context_prefeature, + validate_step_tensors, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--manifest", + type=Path, + default=Path("/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4/manifest.jsonl"), + ) + parser.add_argument("--expected_records", type=int, default=3200) + parser.add_argument("--check_videos", action=argparse.BooleanOptionalAction, default=True) + args = parser.parse_args() + manifest = args.manifest.resolve() + root = manifest.parent + with manifest.open("r", encoding="utf-8") as handle: + records = [json.loads(line) for line in handle if line.strip()] + if len(records) != args.expected_records: + raise ValueError(f"Found {len(records)} records, expected {args.expected_records}") + keys = Counter() + cases_checked = set() + frame_histogram = Counter() + for record in records: + if record["schema_version"] != SCHEMA_VERSION_PREFEATURE: + raise ValueError("Schema mismatch") + key = (int(record["case_id"]), int(record["action_id"]), int(record["chunk_id"])) + keys[key] += 1 + if key[0] not in cases_checked: + validate_case_tensors(load_file(str(root / record["case_tensor_file"]), device="cpu")) + cases_checked.add(key[0]) + validate_step_tensors(load_file(str(root / record["step_tensor_file"]), device="cpu")) + block_frames = [] + selected = None + for block_id in CONTEXT_BLOCK_IDS: + value = load_file(str(root / record["context_tensor_files"][str(block_id)]), device="cpu") + block_frames.append(validate_context_prefeature(block_id, value)) + current = value["selected_frame_indices"] + if selected is not None and not current.equal(selected): + raise ValueError(f"Selected frame mismatch across blocks for {key}") + selected = current + if len(set(block_frames)) != 1 or block_frames[0] != int(record["context_frames"]): + raise ValueError(f"Context frame mismatch for {key}") + frame_histogram[block_frames[0]] += 1 + duplicates = [key for key, count in keys.items() if count != 1] + if duplicates: + raise ValueError(f"Duplicate keys: {duplicates[:10]}") + if args.check_videos: + missing_videos = [ + (case_id, action_id) + for case_id in range(100) for action_id in range(4) + if not (root / "videos" / f"case_{case_id:04d}_action_{action_id:02d}.mp4").is_file() + ] + if missing_videos: + raise ValueError(f"Missing videos: {missing_videos[:10]}") + print(json.dumps({ + "status": "ok", "records": len(records), "cases": len(cases_checked), + "frame_histogram": dict(sorted(frame_histogram.items())), + }, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validate_predictor_text_kv_all54.py b/tools/validate_predictor_text_kv_all54.py new file mode 100644 index 0000000000000000000000000000000000000000..de815e4d441cb7068c750dead40fbe6fc21a05f7 --- /dev/null +++ b/tools/validate_predictor_text_kv_all54.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""Validate cardinality, schema, shape, dtype, and provenance of Text K/V cache files.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from safetensors import safe_open +from safetensors.torch import load_file + +from predictor_data.text_kv_schema import ( + SCHEMA_VERSION_TEXT_KV_ALL54, + validate_text_kv_all54, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, required=True) + parser.add_argument("--expected_cases", type=int, default=100) + args = parser.parse_args() + root = args.root.resolve() + files = sorted(root.glob("case_*_text_kv_exact_all54_bf16.safetensors")) + if len(files) != args.expected_cases: + raise ValueError(f"Expected {args.expected_cases} files, found {len(files)}") + total_bytes = 0 + for case_id, path in enumerate(files): + expected_name = f"case_{case_id:04d}_text_kv_exact_all54_bf16.safetensors" + if path.name != expected_name: + raise ValueError(f"Expected {expected_name}, found {path.name}") + with safe_open(str(path), framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + if metadata.get("schema_version") != SCHEMA_VERSION_TEXT_KV_ALL54: + raise ValueError(f"Schema mismatch: {path}") + if metadata.get("case_id") != str(case_id): + raise ValueError(f"Case metadata mismatch: {path}") + validate_text_kv_all54(load_file(str(path), device="cpu")) + total_bytes += path.stat().st_size + print(f"[ok] {case_id + 1}/{len(files)} {path.name}", flush=True) + print( + json.dumps( + { + "status": "ok", + "schema_version": SCHEMA_VERSION_TEXT_KV_ALL54, + "cases": len(files), + "total_bytes": total_bytes, + "total_gib": total_bytes / 2**30, + }, + indent=2, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validate_predictor_v2_dataset.py b/tools/validate_predictor_v2_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..025c77fec77378c42a061bad2f1bfa20d7a0e749 --- /dev/null +++ b/tools/validate_predictor_v2_dataset.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""Validate Predictor-v2 manifests, BF16 targets, and four-layer context KV.""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter, defaultdict +from pathlib import Path + +from safetensors.torch import load_file + +from predictor_data.v2_schema import ( + CONTEXT_BLOCK_IDS, + SCHEMA_VERSION_V2, + validate_case_tensors_v2, + validate_step_tensors_v2, + validate_vision_context_v2, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", default="datasets/predictor_v2/manifest.jsonl") + parser.add_argument("--expected_cases", type=int, default=10) + parser.add_argument("--expected_chunks_per_case", type=int, default=16) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + manifest = Path(args.manifest).resolve() + root = manifest.parent + with manifest.open("r", encoding="utf-8") as handle: + records = [json.loads(line) for line in handle if line.strip()] + if len(records) != args.expected_cases * args.expected_chunks_per_case: + raise ValueError(f"Unexpected record count: {len(records)}") + + counts = Counter() + seen_cases = set() + frame_histogram: dict[int, Counter] = defaultdict(Counter) + text_padded_tokens = set() + text_valid_tokens: dict[int, int] = {} + context_padded_tokens = set() + for record in records: + if record["schema_version"] != SCHEMA_VERSION_V2: + raise ValueError(f"Unexpected schema: {record['schema_version']}") + key = (int(record["case_id"]), int(record["seed"]), int(record["chunk_id"])) + counts[key] += 1 + case_id = key[0] + if case_id not in seen_cases: + case = load_file(str(root / record["case_tensor_file"]), device="cpu") + text_valid_tokens[case_id] = validate_case_tensors_v2(case) + text_padded_tokens.add(int(case["text_valid_mask"].shape[1])) + seen_cases.add(case_id) + steps = load_file(str(root / record["step_tensor_file"]), device="cpu") + validate_step_tensors_v2(steps) + frame_counts = [] + for block_id in CONTEXT_BLOCK_IDS: + context = load_file( + str(root / record["context_tensor_files"][str(block_id)]), + device="cpu", + ) + frames = validate_vision_context_v2(block_id, context) + context_padded_tokens.add(int(context["context_valid_mask"].shape[1])) + frame_counts.append(frames) + frame_histogram[block_id][frames] += 1 + if len(set(frame_counts)) != 1 or frame_counts[0] != int(record["context_frames"]): + raise ValueError(f"Context frame mismatch for {key}: {frame_counts}") + duplicates = [key for key, count in counts.items() if count != 1] + if duplicates: + raise ValueError(f"Duplicate manifest keys: {duplicates[:5]}") + if len(text_padded_tokens) != 1: + raise ValueError(f"Text KV is not uniformly padded: {sorted(text_padded_tokens)}") + if len(context_padded_tokens) != 1: + raise ValueError( + f"Context Vision KV is not uniformly padded: {sorted(context_padded_tokens)}" + ) + payload = { + "root": str(root), + "records": len(records), + "cases": len(seen_cases), + "context_blocks": list(CONTEXT_BLOCK_IDS), + "text_padded_tokens": next(iter(text_padded_tokens)), + "text_valid_tokens": dict(sorted(text_valid_tokens.items())), + "context_padded_tokens": next(iter(context_padded_tokens)), + "context_frame_histogram": { + str(block_id): dict(sorted(counter.items())) + for block_id, counter in frame_histogram.items() + }, + "status": "ok", + } + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/video_review_server.py b/tools/video_review_server.py new file mode 100644 index 0000000000000000000000000000000000000000..b56ca1369c2f49d7fcccbdb10faa52c28915ee1a --- /dev/null +++ b/tools/video_review_server.py @@ -0,0 +1,1031 @@ +#!/usr/bin/env python3 +"""A zero-dependency video review and download server.""" + +from __future__ import annotations + +import argparse +import json +import mimetypes +import os +import re +import secrets +import threading +import time +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path, PurePosixPath +from urllib.parse import parse_qs, quote, urlparse + + +VIDEO_EXTENSIONS = {".mp4", ".webm", ".mov", ".mkv", ".avi", ".m4v"} +CHUNK_SIZE = 256 * 1024 +DEFAULT_REFRESH_INTERVAL = 30.0 +DEFAULT_SETTLE_SECONDS = 3.0 + +INDEX_HTML = r""" + + + + + Projects 视频审阅 + + + +
+ + +
+
+
+
尚未选择视频
+
点击文件在当前窗口打开，或拖入工作区新增窗口
+
+
+ + + + + + +
0 个窗口
+
+
+
+
+ 拖入视频开始多窗口比较 + 点击文件使用单窗口；拖动多个文件到这里可并排查看。 +
+
+
+ 就绪 + 拖入新增窗口 · 点击替换活动窗口 · J/K/L 控制活动窗口 +
+
+
+ + + +""" + + +class ReviewState: + def __init__( + self, + root: Path, + refresh_interval: float = DEFAULT_REFRESH_INTERVAL, + settle_seconds: float = DEFAULT_SETTLE_SECONDS, + ): + self.root = root.resolve() + self.refresh_interval = max(0.0, refresh_interval) + self.settle_seconds = max(0.0, settle_seconds) + self.lease_lock = threading.Lock() + self.catalog_lock = threading.Lock() + self.refresh_guard = threading.Lock() + self.leases: dict[str, tuple[str, str]] = {} + self.refreshing = False + self.last_scan_started = time.time() + initial = self._scan_raw() + self.videos = [self._public_item(item) for item in initial] + self.paths = {str(item["path"]) for item in self.videos} + self.signatures = { + str(item["path"]): (int(item["size"]), int(item["mtime_ns"])) + for item in initial + } + self.total_size = sum(int(item["size"]) for item in self.videos) + self.revision = 1 + self.last_scan_finished = time.time() + if self.refresh_interval > 0: + threading.Thread( + target=self._auto_refresh_loop, + name="video-catalog-auto-refresh", + daemon=True, + ).start() + + @staticmethod + def _public_item(item: dict[str, object]) -> dict[str, object]: + return {"path": str(item["path"]), "size": int(item["size"])} + + def _scan_raw(self) -> list[dict[str, object]]: + videos = [] + for directory, dirnames, filenames in os.walk(self.root): + dirnames.sort() + for filename in sorted(filenames): + path = Path(directory, filename) + if path.suffix.lower() not in VIDEO_EXTENSIONS: + continue + try: + stat = path.stat() + relative = path.relative_to(self.root).as_posix() + except (OSError, ValueError): + continue + videos.append( + {"path": relative, "size": stat.st_size, "mtime_ns": stat.st_mtime_ns} + ) + videos.sort(key=lambda item: str(item["path"])) + return videos + + def _stable_after_delay( + self, items: list[dict[str, object]] + ) -> set[str]: + if not items: + return set() + if self.settle_seconds: + time.sleep(self.settle_seconds) + stable = set() + for item in items: + relative = str(item["path"]) + path = self.root / PurePosixPath(relative) + try: + stat = path.stat() + except OSError: + continue + if ( + stat.st_size == int(item["size"]) + and stat.st_mtime_ns == int(item["mtime_ns"]) + ): + stable.add(relative) + return stable + + def _refresh_worker(self) -> None: + try: + with self.catalog_lock: + self.refreshing = True + self.last_scan_started = time.time() + old_items = {str(item["path"]): item for item in self.videos} + old_signatures = dict(self.signatures) + + scanned = self._scan_raw() + changed = [ + item + for item in scanned + if old_signatures.get(str(item["path"])) + != (int(item["size"]), int(item["mtime_ns"])) + ] + stable_changed = self._stable_after_delay(changed) + + accepted = [] + accepted_signatures: dict[str, tuple[int, int]] = {} + for item in scanned: + relative = str(item["path"]) + signature = (int(item["size"]), int(item["mtime_ns"])) + if old_signatures.get(relative) == signature or relative in stable_changed: + accepted.append(self._public_item(item)) + accepted_signatures[relative] = signature + elif relative in old_items: + accepted.append(old_items[relative]) + accepted_signatures[relative] = old_signatures[relative] + + accepted.sort(key=lambda item: str(item["path"])) + with self.catalog_lock: + catalog_changed = accepted != self.videos + self.videos = accepted + self.paths = {str(item["path"]) for item in accepted} + self.signatures = accepted_signatures + self.total_size = sum(int(item["size"]) for item in accepted) + if catalog_changed: + self.revision += 1 + self.last_scan_finished = time.time() + print( + f"Catalog refresh complete: {len(self.videos)} videos " + f"(revision {self.revision})", + flush=True, + ) + except Exception as error: + print(f"Catalog refresh failed: {error}", flush=True) + finally: + with self.catalog_lock: + self.refreshing = False + self.last_scan_finished = time.time() + self.refresh_guard.release() + + def request_refresh(self) -> bool: + if not self.refresh_guard.acquire(blocking=False): + return False + with self.catalog_lock: + self.refreshing = True + self.last_scan_started = time.time() + threading.Thread( + target=self._refresh_worker, + name="video-catalog-refresh", + daemon=True, + ).start() + return True + + def _auto_refresh_loop(self) -> None: + while True: + time.sleep(self.refresh_interval) + self.request_refresh() + + def status(self) -> dict[str, object]: + with self.catalog_lock: + return { + "revision": self.revision, + "video_count": len(self.videos), + "total_size": self.total_size, + "refreshing": self.refreshing, + "last_scan_started": self.last_scan_started, + "last_scan_finished": self.last_scan_finished, + "refresh_interval": self.refresh_interval, + } + + def catalog(self) -> dict[str, object]: + with self.catalog_lock: + result = { + "videos": list(self.videos), + "revision": self.revision, + "video_count": len(self.videos), + "total_size": self.total_size, + "refreshing": self.refreshing, + "last_scan_started": self.last_scan_started, + "last_scan_finished": self.last_scan_finished, + "refresh_interval": self.refresh_interval, + } + return result + + def resolve_video(self, relative: str) -> Path | None: + normalized = PurePosixPath(relative).as_posix() + with self.catalog_lock: + known = normalized in self.paths + if not known: + return None + candidate = (self.root / normalized).resolve() + try: + candidate.relative_to(self.root) + except ValueError: + return None + return candidate if candidate.is_file() else None + + def select(self, relative: str, slot: str) -> str: + token = secrets.token_urlsafe(24) + with self.lease_lock: + if slot not in self.leases and len(self.leases) >= 256: + self.leases.pop(next(iter(self.leases))) + self.leases[slot] = (relative, token) + return token + + def release(self, slot: str, token: str) -> bool: + with self.lease_lock: + lease = self.leases.get(slot) + if lease is None or not secrets.compare_digest(token, lease[1]): + return False + del self.leases[slot] + return True + + def is_active(self, relative: str, slot: str, token: str) -> bool: + with self.lease_lock: + lease = self.leases.get(slot) + return ( + lease is not None + and relative == lease[0] + and secrets.compare_digest(token, lease[1]) + ) + + +class VideoReviewHandler(BaseHTTPRequestHandler): + server_version = "HYVideoReview/1.0" + + @property + def state(self) -> ReviewState: + return self.server.review_state # type: ignore[attr-defined] + + def log_message(self, fmt: str, *args: object) -> None: + print(f"{self.address_string()} - {fmt % args}", flush=True) + + def send_bytes(self, data: bytes, content_type: str, status: int = 200) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(data))) + self.send_header("Cache-Control", "no-store") + self.send_header("X-Content-Type-Options", "nosniff") + self.end_headers() + if self.command != "HEAD": + self.wfile.write(data) + + def send_json(self, value: object, status: int = 200) -> None: + self.send_bytes( + json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode(), + "application/json; charset=utf-8", + status, + ) + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path == "/": + self.send_bytes(INDEX_HTML.encode(), "text/html; charset=utf-8") + elif parsed.path == "/api/videos": + self.send_json(self.state.catalog()) + elif parsed.path == "/api/status": + self.send_json(self.state.status()) + elif parsed.path == "/media": + self.serve_media(parsed.query) + elif parsed.path == "/download": + self.serve_download(parsed.query) + elif parsed.path == "/healthz": + self.send_json({"ok": True, **self.state.status()}) + else: + self.send_error(HTTPStatus.NOT_FOUND) + + def do_HEAD(self) -> None: + parsed = urlparse(self.path) + if parsed.path == "/media": + self.serve_media(parsed.query) + elif parsed.path == "/download": + self.serve_download(parsed.query) + elif parsed.path == "/healthz": + self.send_json({"ok": True, "videos": len(self.state.videos)}) + else: + self.send_error(HTTPStatus.NOT_FOUND) + + def do_POST(self) -> None: + endpoint = urlparse(self.path).path + if endpoint == "/api/refresh": + started = self.state.request_refresh() + self.send_json({"started": started, **self.state.status()}, HTTPStatus.ACCEPTED) + return + if endpoint not in {"/api/select", "/api/release"}: + self.send_error(HTTPStatus.NOT_FOUND) + return + try: + length = int(self.headers.get("Content-Length", "0")) + if length > 64 * 1024: + raise ValueError("request too large") + payload = json.loads(self.rfile.read(length)) + slot = str(payload["slot"]) + if not slot or len(slot) > 128: + raise ValueError("invalid slot") + except (ValueError, KeyError, TypeError, json.JSONDecodeError): + self.send_json({"error": "invalid request"}, HTTPStatus.BAD_REQUEST) + return + + if endpoint == "/api/release": + try: + token = str(payload["token"]) + except (KeyError, TypeError): + self.send_json({"error": "invalid request"}, HTTPStatus.BAD_REQUEST) + return + self.send_json({"released": self.state.release(slot, token)}) + return + + try: + relative = str(payload["path"]) + except (KeyError, TypeError): + self.send_json({"error": "invalid request"}, HTTPStatus.BAD_REQUEST) + return + if self.state.resolve_video(relative) is None: + self.send_json({"error": "video not found"}, HTTPStatus.NOT_FOUND) + return + self.send_json({"token": self.state.select(relative, slot)}) + + def serve_media(self, query: str) -> None: + params = parse_qs(query) + relative = params.get("path", [""])[0] + slot = params.get("slot", [""])[0] + token = params.get("token", [""])[0] + path = self.state.resolve_video(relative) + if path is None: + self.send_error(HTTPStatus.NOT_FOUND) + return + if not self.state.is_active(relative, slot, token): + self.send_error(HTTPStatus.GONE, "video lease expired") + return + + size = path.stat().st_size + start, end = 0, size - 1 + range_header = self.headers.get("Range") + if range_header: + match = re.fullmatch(r"bytes=(\d*)-(\d*)", range_header.strip()) + if not match: + self.send_error(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE) + return + first, last = match.groups() + if first: + start = int(first) + end = min(int(last), size - 1) if last else size - 1 + elif last: + suffix = int(last) + start = max(0, size - suffix) + if start >= size or end < start: + self.send_response(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE) + self.send_header("Content-Range", f"bytes */{size}") + self.end_headers() + return + + length = end - start + 1 + content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream" + self.send_response(HTTPStatus.PARTIAL_CONTENT if range_header else HTTPStatus.OK) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(length)) + self.send_header("Accept-Ranges", "bytes") + self.send_header("Cache-Control", "private, no-store") + if range_header: + self.send_header("Content-Range", f"bytes {start}-{end}/{size}") + self.end_headers() + if self.command == "HEAD": + return + + remaining = length + try: + with path.open("rb") as source: + source.seek(start) + while remaining > 0 and self.state.is_active(relative, slot, token): + chunk = source.read(min(CHUNK_SIZE, remaining)) + if not chunk: + break + self.wfile.write(chunk) + remaining -= len(chunk) + except (BrokenPipeError, ConnectionResetError): + pass + + def serve_download(self, query: str) -> None: + relative = parse_qs(query).get("path", [""])[0] + path = self.state.resolve_video(relative) + if path is None: + self.send_error(HTTPStatus.NOT_FOUND) + return + + size = path.stat().st_size + content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream" + fallback = re.sub(r"[^A-Za-z0-9._-]", "_", path.name) or "video" + encoded = quote(path.name, safe="") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(size)) + self.send_header( + "Content-Disposition", + f'attachment; filename="{fallback}"; filename*=UTF-8\'\'{encoded}', + ) + self.send_header("Cache-Control", "private, no-store") + self.send_header("X-Content-Type-Options", "nosniff") + self.end_headers() + if self.command == "HEAD": + return + + try: + with path.open("rb") as source: + while chunk := source.read(CHUNK_SIZE): + self.wfile.write(chunk) + except (BrokenPipeError, ConnectionResetError): + pass + + +class ReviewServer(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + def __init__(self, address: tuple[str, int], state: ReviewState): + super().__init__(address, VideoReviewHandler) + self.review_state = state + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8765) + parser.add_argument("--refresh-interval", type=float, default=DEFAULT_REFRESH_INTERVAL) + parser.add_argument("--settle-seconds", type=float, default=DEFAULT_SETTLE_SECONDS) + args = parser.parse_args() + + state = ReviewState(args.root, args.refresh_interval, args.settle_seconds) + server = ReviewServer((args.host, args.port), state) + print( + f"Reviewing {len(state.videos)} videos under {state.root}\n" + f"Listening on http://{args.host}:{args.port}\n" + f"Auto-refresh every {state.refresh_interval:g}s; " + f"new files settle for {state.settle_seconds:g}s", + flush=True, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + pass + finally: + server.server_close() + + +if __name__ == "__main__": + main() diff --git a/train.md b/train.md new file mode 100644 index 0000000000000000000000000000000000000000..10de3191f86fa47cf02ca04e7b5ed3db42c715c6 --- /dev/null +++ b/train.md @@ -0,0 +1,592 @@ +# Train + +## 实验命名规范 + +从 2026-07-22 起，checkpoint 目录、独立日志、SwanLab experiment 和评测输出统一显式记录两组 +学习率、训练总步数与两种 warmup： + +```text +_flr_blr_schedsteps_swarmup_bwarmup +``` + +- `flr`：image fusion/residual 参数组的峰值学习率； +- `blr`：两个 Predictor blocks 参数组的峰值学习率； +- `schedsteps`：创建 LR scheduler 时使用的训练总步数，不代表某次运行实际走到的步数； +- `swarmup`：全局 cosine scheduler 的 warmup steps； +- `bwarmup`：blocks 参数组额外冻结/延迟启用的 warmup steps；旧 rollout 没有该独立参数时写 + `bwarmup-none`； +- `ckptstep`：仅用于评测输出，表示本次评测实际加载的 checkpoint step。 + +例如，计划训练 2000 steps、实际使用 step 300 权重进行评测时，训练目录保留 +`schedsteps2000`，评测目录同时写 `ckptstep300` 与 `schedsteps2000`。目录内的 `train_log.jsonl`、 +`torchrun.log`、`train_config.json` 和 `swanlab/` 不重复拼接参数，它们继承父目录的完整命名。 + +当前主要配置如下： + +| 实验组 | `flr` | `blr` | `schedsteps` | `swarmup` | `bwarmup` | +| --- | ---: | ---: | ---: | ---: | ---: | +| v4 / predictor-disca 离线 Stage 1 | 1e-4 | 1e-5 | 2000 | 100 | 100 | +| 早期整轨迹 FPPP rollout | 1e-5 | 1e-6 | 2000 | 50 | none | +| overfit8 FPPP rollout | 1e-4 | 1e-5 | 2000 | 8 | 8 | +| overfit8 FPPF rollout | 1e-4 | 1e-5 | 300 | 8 | 8 | +| 旧 timestep-only sweep | 1e-4 | 1e-5 | 200 | 8 | 8 | +| detached timestep-only | 1e-4 | 1e-5 | 1000 | 8 | 8 | + +## 1. 监督目标 + +当前 v4 数据为每个 chunk 保存四步 Full-DiT Teacher hidden 和 velocity。正式 v4 与 +predictor-disca 使用全部相邻监督 pair： + +```text +step 0 final hidden + step 1 input -> step 1 hidden / velocity +step 1 final hidden + step 2 input -> step 2 hidden / velocity +step 2 final hidden + step 3 input -> step 3 hidden / velocity +``` + +v1/v2/v3 的历史训练只使用 `0→1` 和 `2→3`，对应 F-P-F-P；从 v4 开始加入 `1→2`，使 Predictor +具备连续替代和任意选择后三步的能力。 + +## 2. Loss + +所有版本使用相同的监督损失： + +```python +hidden_loss = mse(pred_hidden.float(), target_hidden.float()) +velocity_loss = mse(pred_velocity.float(), target_velocity.float()) +loss = 0.1 * hidden_loss + 1.0 * velocity_loss +``` + +- hidden 和 velocity 的 forward 使用 BF16； +- MSE 转为 FP32 计算； +- velocity MSE 是主目标； +- hidden MSE 是 final-layer 前 feature 的辅助约束； +- frozen final layer 只冻结参数，forward 不能放进 `no_grad()`，否则 velocity loss 无法回传到 + Predictor hidden。 + +## 3. 优化方法 + +当前正式训练使用 AdamW： + +| 参数组 | 学习率 | +| --- | ---: | +| image fusion + residual head | 1e-4 | +| 两个 Predictor blocks | 1e-5 | + +共同配置： + +```yaml +betas: [0.9, 0.95] +weight_decay: 0.01 +grad_clip: 1.0 +precision: bfloat16 +loss_precision: float32 +max_steps: 2000 +scheduler: cosine +scheduler_warmup_steps: 100 +block_warmup_steps: 100 +gradient_checkpointing: true +``` + +前 100 optimizer steps 中 Predictor block 参数组的学习率为 0，只训练 fusion/residual；之后 block +组按较低学习率参与训练。两个参数组都经过 100-step warmup 和 cosine decay，step 2000 衰减到 0。 + +## 4. 各版本训练方式 + +### v1 + +- 入口：`train_predictor.py`； +- 单 GPU、micro-batch 1； +- gradient accumulation 4，有效 batch 4； +- 2000 steps； +- pairs：`0→1`、`2→3`； +- source blocks：`[1,52]`； +- 每 500 steps 保存，默认保留 2 个快照。 + +旧机器上的 v1 实验结果不再纳入当前文档，这里只保留训练方法。 + +### v2/v3 + +- 入口：`train_predictor_v2.py --predictor_version v2|v3`； +- 单 GPU，默认 batch 1、gradient accumulation 4； +- pairs：`0→1`、`2→3`； +- source blocks：`[0,53]` 或 `[1,52]`； +- 2000 steps、100-step block warmup； +- SwanLab 可选。 + +v3 Dataset 额外加载前一 chunk 的 target-step final hidden，并排除 chunk 0。迁移后的当前目录未保留 +v2/v3 正式 checkpoint，因此这里只记录方法，不将缺失产物写成可复现实验结果。 + +### v4 + +- 入口：`train_predictor_v4.py`； +- 启动：`tools/launch_predictor_v4_training.sh`； +- 8 卡 DDP； +- 每卡 batch size 16，全局 batch size 128； +- gradient accumulation 1； +- 2800 records × 3 transitions = 8400 pairs； +- 按 Context frames 分桶，减少动态 padding； +- 2000 steps； +- 每 100 steps 保存，保留 20 个 safetensors； +- SwanLab 仅由 rank 0 上报。 + +已训练 blocks `[0,53]` 和 `[1,52]` 两版。 + +### predictor-disca + +- 入口：`train_predictor_disca.py`； +- 启动：`tools/launch_predictor_disca_training.sh`； +- 8 卡 DDP、每卡 batch size 16、全局 batch size 128； +- 8400 pairs、2000 steps； +- 不加载 Context pre-feature 或 previous-chunk hidden； +- 不按 Context frames 分桶，使用普通 DistributedSampler； +- 每 100 steps 保存，保留 20 个 safetensors； +- 当前正式版使用 blocks `[0,53]`。 + +## 5. 当前正式训练结果 + +下表是训练日志在 step 2000 记录窗口内的均值，不是完整数据集 validation loss： + +| 版本 | blocks | total loss | hidden MSE | velocity MSE | peak memory/GPU | avg micro time | +| --- | --- | ---: | ---: | ---: | ---: | ---: | +| v4 | `[0,53]` | 0.011895 | 0.011744 | 0.010721 | 62.82 GiB | 4.47 s | +| v4 | `[1,52]` | 0.011537 | 0.011732 | 0.010364 | 62.82 GiB | 4.49 s | +| predictor-disca | `[0,53]` | 0.011053 | 0.011121 | 0.009941 | 31.18 GiB | 1.66 s | + +训练 loss 不能直接替代 rollout 视频质量。predictor-disca 的训练 loss 更低，但其训练集视频指标明显 +低于 v4，说明缺失 AR Context 后更容易拟合单步监督，却在连续在线 rollout 中累积了更大误差。 + +## 6. Checkpoint 与 resume + +正式训练同时保存： + +```text +predictor_step_XXXXX.safetensors # 仅模型推理权重 +training_latest.pt # model、optimizer、scheduler、step、epoch、配置 +train_config.json +train_log.jsonl +torchrun.log +swanlab/ +``` + +resume 会恢复 model、optimizer、scheduler、global/micro step 和 epoch，并检查 source blocks 与 +Predictor variant 是否匹配。例如： + +```bash +python train_predictor_disca.py \ + --resume checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt +``` + +正式云端监控使用： + +```bash +export SWANLAB_API_KEY='' +``` + +token 只放在进程环境或用户登录缓存中，不写入仓库、配置或日志。 + +## 7. 在线训练集评测 + +统一评测 100 个训练 case × 4 actions，共 400 个 125-frame 视频： + +```text +reference: Full-DiT F-F-F-F +reuse: Full steps 0、3；steps 1、2 reuse +predictor: Full steps 0、3；steps 1、2 Predictor +``` + +视频指标包括 full-resolution PSNR、frame-mean PSNR、SSIM、AlexNet LPIPS、pixel +MSE/RMSE/MAE 和 temporal-delta MSE。 + +| 方法 | blocks | PSNR | SSIM | LPIPS | temporal delta MSE | +| --- | --- | ---: | ---: | ---: | ---: | +| Direct Reuse | — | 18.4639 | 0.6429 | 0.2083 | 0.01212 | +| v4 | `[0,53]` | 21.9337 | 0.7606 | 0.1204 | 0.00865 | +| v4 | `[1,52]` | 22.1090 | 0.7670 | 0.1168 | 0.00832 | +| predictor-disca | `[0,53]` | 19.4084 | 0.6748 | 0.1816 | 0.01061 | + +这些是 train-set/in-sample 指标，只能比较架构在已训练图文 pair 上的拟合和 rollout 行为，不能证明 +对新图、新 prompt、新 seed 或新动作的泛化。validation/test 应使用 `data.md` 中互斥的 25/50 pair +holdout。 + +## 8. v4 递归 F-P-P-P Self-Rollout 微调方案 + +### 8.1 定位与实现状态 + +这是一阶段离线 Teacher-forcing 训练之后的二阶段 rollout 微调。初始化应加载已有 v4 Predictor +权重，而不是重新从两个 Teacher blocks 开始训练。当前训练集视频结果中 `[1,52]` 最好，因此正式 +实验优先从以下离线 v4 FP32-master 训练状态中的 `model` 初始化；optimizer/scheduler 不继承， +rollout 使用全新的优化状态： + +```text +checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt +``` + +`[0,53]` 可以用同一套流程作为独立对照，但 source blocks 必须与初始 checkpoint 一致。rollout +训练已作为独立入口实现，不改变现有离线 `train_predictor_v4.py`： + +```text +train_predictor_v4_rollout.py +predictor_training/trajectory_dataset.py +predictor_training/rollout_utils.py +tools/launch_predictor_v4_rollout_training.sh +``` + +低学习率 `[1,52]` 基线输出到以下目录，并已按实验要求在 step 600 停止： + +```text +checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-5_blr1e-6_schedsteps2000_swarmup50_bwarmup-none_swanlab +``` + +使用第一阶段学习率的正式入口改用明确 LR 命名，避免用含义模糊的 `stage1lr`： + +```text +checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab +``` + +### 8.2 目标轨迹 + +每条 item 是一个完整的 `(case_id, action_id)` trajectory，而不是一个独立 transition： + +```text +Chunk 0：F-F-F-F，只建立 self-generated history 和 previous-chunk hidden，不训练 +Chunk 1～7：F-P-P-P，三个 Predictor step 都计算 loss +``` + +在训练 chunk 内： + +```text +initial x0 + -> Full step 0: h0, v0 + -> Euler update: x1 + -> Predictor step 1: pred_h1, pred_v1 + -> differentiable Euler update: x2 + -> Predictor step 2: pred_h2, pred_v2 + -> differentiable Euler update: x3 + -> Predictor step 3: pred_h3, pred_v3 + -> differentiable Euler update: final chunk latent +``` + +同一 chunk 内 `pred_h1 → pred_h2 → pred_h3` 是递归 same-chunk anchor，三个 Predictor 调用保留 +一张完整计算图。chunk 之间则同时 detach： + +- final chunk latent detach 后加入 history； +- `pred_h1/pred_h2/pred_h3` detach 后作为下一 chunk 的同 step previous-chunk hidden。 + +因此，v4 的 previous-chunk 输入不能只用最终 latent 代替。Chunk 0 虽不训练，也必须保存 Full-DiT +step 1、2、3 的 final hidden，分别作为 Chunk 1 Predictor step 1、2、3 的 previous-chunk hidden。 + +### 8.3 轻量样本表与全在线构建 + +Self-rollout 不依赖之前构建的 F-F-F-F tensor 数据集。Dataset 只保存一张轻量样本表，每个 item +包含： + +```text +case_id, source_image, caption, action_id, pose, seed +``` + +100 个图文 pair × 4 actions 组成 400 条 trajectory。`__getitem__` 不读取 step tensor、Teacher +hidden/velocity、Context pre-feature 或 K/V；每个 rank 拿到样本描述后，使用当前正式 WorldPlay +pipeline 在线完成： + +1. 按与数据构建/评测相同的规则将输入图像 resize/crop 到 480×832； +2. 使用 `pose_to_input(pose, 32)` 生成完整 8-chunk viewmats、Ks 和 action； +3. 用指定 seed 调用 `prepare_latents`，一次生成 `[1,32,32,30,52]` 初始噪声，再按 4 latent + frames 切成 8 个 chunk； +4. 从现有 case tensor 读取静态 `image_condition_latent`，在线构造 condition/mask，避免重复 VAE + encode； +5. 从离线数据集的 `text_kv_exact_all54_bf16/` 读取全部 54 层精确 BF16 Text K/V，并按 + `text_valid_mask` 去掉磁盘 padding； +6. 从 Chunk 0 开始在线执行整条 F-F-F-F/F-P-P-P trajectory，并在线生成所有监督 target。 + +每一步仍按当前 pipeline 语义构造 65 通道输入：32 通道 current noisy latent 加 32 通道 image +condition 和 1 通道 mask。只有 Chunk 0 的第一个 latent frame 含 image condition；Chunk 1～7 的 +33 个 condition 通道为零。history prefill 同样拼接对应历史帧的固定 condition，不能只把 32 通道 +self latent 送入 Teacher。 + +因此 rollout 训练不读取或生成持久化的 `step_1/2/3_noisy_sample`、Teacher target、Context +pre-feature、previous-step hidden 或全层 Text K/V sidecar。Text K/V 与 Predictor policy 和 trajectory +状态无关，按 case 存在主离线数据集内；所有 policy-dependent 状态仍在当前参数、当前 self history 上 +产生，语义是严格 on-policy。图像/caption/action 样本表不是 Teacher trajectory 数据，可以直接从 +当前 100 个训练 pair 的来源清单生成；留出集不参与训练。 + +静态 Text K/V 构建时来自 prompt、vision 和 ByT5 编码，训练主循环不再加载这三个编码器，也不再 +重复 text prefill。 + +### 8.4 v4 Context 的准确语义 + +每个 chunk 使用当前 detached、self-generated history 运行与正式推理相同的 +`select_aligned_memory_frames`。选择结果只由相机轨迹决定；被选中的 latent 内容必须来自当前 +self history。 + +随后冻结 Full Teacher 对这些 history latent 执行一次 54-layer AR vision prefill，得到精确 Context +Vision K/V。该 cache： + +- 与当前 trajectory 的全层 Text K/V 合并； +- 在当前 chunk 的四个 denoise steps 间只读共享； +- Teacher target 使用全部 54 层 cache； +- v4 Predictor 直接读取其 source blocks 对应的同一份 exact K/V； +- 全部 detach，且不对 Context 建图。 + +这里不走 `context_prefeature -> frozen projector -> K/V` 路径。pre-feature 是离线 pair 训练的磁盘 +压缩方案；在线 rollout 已经由 Teacher prefill 持有 exact K/V，行为应与当前推理 pipeline 一致。 + +### 8.5 单个 chunk 的 Teacher 与 Predictor 顺序 + +Chunk 0 在 `torch.no_grad()` 下执行四次 Full-DiT，并保存 final latent 和 Full hidden-by-step。 +Chunk 1～7 的顺序固定为： + +1. 从 detached history 选择 Context，并由 Teacher 生成 exact AR K/V； +2. Full Teacher 在 step 0 输入上产生 `h0/v0`，用正式 Euler scheduler 更新到 `x1`； +3. 对 step 1～3，先在 `current_student_latent.detach()` 上运行 Full Teacher，获得 + `target_hidden/target_velocity`； +4. Predictor 使用相同 noisy latent、condition、相机参数和 Context K/V； +5. step 1 的 anchor 是 `h0`，step 2 是 `pred_h1`，step 3 是 `pred_h2`； +6. 每个 Predictor velocity 经可微 scheduler 更新到下一 latent； +7. 三步 loss 求均值，当前 chunk 立即 backward； +8. final latent 与三步 pred hidden detach 后加入下一 chunk 状态。 + +Teacher target 必须来自 Predictor 真正访问的状态。对应当前 Transformer API 的核心调用是： + +```python +with torch.no_grad(): + teacher_model_input = torch.cat( + [current_student_latent.detach(), current_condition_33ch], dim=1 + ) + teacher_output = full_teacher( + bi_inference=False, + ar_txt_inference=False, + ar_vision_inference=True, + hidden_states=teacher_model_input, + timestep=current_timestep, + timestep_r=None, + mask_type=task_type, + action=current_action, + viewmats=current_viewmats, + Ks=current_Ks, + kv_cache=exact_context_kv, + cache_vision=False, + rope_temporal_size=rope_temporal_size, + start_rope_start_idx=start_rope_start_idx, + return_final_hidden=True, + ) + target_velocity = teacher_output[0] + target_hidden = teacher_output[2] +``` + +不能读取原 F-F-F-F trajectory 中相同步的 target。Teacher 只用于构造 stop-gradient label，其输出 +立即释放。Predictor 使用的 `target_frame_condition` 继续由冻结 Teacher 的 +`time_in(timestep) + action_in(action)` 计算，不从旧 target 文件推断。 + +### 8.6 完整 timestep BPTT 的代码前置条件 + +当前 `HYWorldPlayPredictorV2.forward`（v4 继承该实现）使用 `torch.no_grad()` 包裹冻结 `img_in`。 +这会阻断 `x2/x3 -> img_in -> Predictor` 的输入梯度，不满足完整 timestep BPTT。rollout 实现前必须 +改为： + +```python +# img_in 参数保持 requires_grad=False，但 forward 不进入 no_grad。 +current_img = self.img_in(target_model_input) +``` + +冻结参数不会产生 parameter gradient；当离线训练输入不需要梯度时也不会额外建图，但 rollout 的 +step 2、3 可以对前一步 scheduler 输出求导。冻结 final layer 同理：参数不训练，forward 仍必须保留 +对 `pred_hidden` 的梯度。Context projector 保持 `no_grad()`，因为 Context 本来就应 detach。 + +### 8.7 可微 FlowMatch Euler update + +当前项目的 `FlowMatchDiscreteScheduler.step` 在 `solver=euler` 时核心公式已经是 out-of-place 且可微： + +```python +x_next = x.float() + pred_velocity.float() * ( + sigmas[step_index + 1] - sigmas[step_index] +) +``` + +但正式 pipeline 会把结果原位写回共享 `latents`，不应直接复制该容器逻辑。训练应实现无状态的 +`differentiable_euler_step`：调用相同 `set_timesteps(4)` 后读取同一组 sigmas，每次只接收和返回当前 +chunk local tensor，无 `detach`、`no_grad` 或 inplace slice assignment。其 FP32 update 和输出 dtype +必须与正式 scheduler 数值一致。 + +step 2、3 的 loss 应同时通过两类路径影响前面的 Predictor 调用： + +```text +hidden anchor path: pred_h1 -> Predictor step 2 -> pred_h2 -> Predictor step 3 +latent path: pred_v1 -> x2 -> Predictor step 2 -> pred_v2 -> x3 -> Predictor step 3 +``` + +### 8.8 Loss 与归一化 + +每个 Predictor step 继续使用 v4 的 FP32 MSE： + +```python +step_loss = ( + mse(pred_velocity.float(), target_velocity.float()) + + 0.1 * mse(pred_hidden.float(), target_hidden.float()) +) +chunk_loss = (step1_loss + step2_loss + step3_loss) / 3 +``` + +一条 trajectory 有七个训练 chunk，所以每个 chunk 执行： + +```python +(chunk_loss / 7).backward() +``` + +不要再除以 world size，DDP 会对 8 个 rank 自动求平均。若以后额外累计多条 trajectory 才更新一次， +还需要再除以 `gradient_accumulation_trajectories`。日志应分别记录 step 1/2/3 hidden MSE、velocity +MSE、chunk id 均值和整条 trajectory loss，不能只记录合并 total loss。 + +### 8.9 8 卡 DDP 与同步边界 + +每个 rank 持有一份冻结的 54-layer Teacher、一份 DDP v4 Predictor 和一条独立 trajectory。只包装 +Predictor： + +```python +predictor = DDP( + predictor, + device_ids=[local_rank], + output_device=local_rank, + broadcast_buffers=False, + gradient_as_bucket_view=True, + find_unused_parameters=False, +) +``` + +400 条 trajectory 可被 8 个 rank 整除。`DistributedSampler(..., drop_last=True)` 每个 epoch 给每个 +rank 50 条，调用 `sampler.set_epoch(epoch)` 保证 shuffle 一致且不同 rank 不重叠。micro-batch 固定为 +每卡一条 trajectory，因此每个 optimizer step 的全局 batch 是 8 trajectories。 + +Chunk 1～6 只累计本地梯度，Chunk 7 backward 时执行一次 all-reduce： + +```python +optimizer.zero_grad(set_to_none=True) + +for chunk_id in range(1, 8): + sync_ctx = nullcontext() if chunk_id == 7 else predictor.no_sync() + with sync_ctx: # 必须同时覆盖三个 DDP forward 和当前 backward + chunk_loss, final_latent, hidden_by_step = train_fppp_chunk(...) + (chunk_loss / 7).backward() + + history.append(final_latent.detach()) + previous_hidden_by_step = { + step: hidden.detach() for step, hidden in hidden_by_step.items() + } + del chunk_loss, final_latent, hidden_by_step + +clip_grad_norm_(predictor.parameters(), 1.0) +optimizer.step() +lr_scheduler.step() +``` + +七个 chunk 中间不能 `optimizer.step()` 或 `zero_grad()`。现有 Predictor checkpoint 使用 +`use_reentrant=False`，应继续保留；三个 DDP forward 合并一次 backward 以及七次 chunk backward 的 +组合必须先通过单卡/DDP 小规模梯度一致性测试。 + +### 8.10 显存与训练配置 + +每个 rank 加载 Transformer Teacher 和 Predictor。image-condition latent 直接从 case tensor 读取； +全层 Text K/V 从主离线数据集读取并去掉 padding。训练不加载 prompt encoder、vision encoder、 +ByT5、VAE、SR 或视频编码组件。可微 rollout 阶段只让 Teacher、Predictor、当前 Text/Context cache +和当前 chunk graph 驻留 GPU。 + +Teacher 全程 eval、`requires_grad_(False)` 和 `no_grad()`；Predictor forward activation 使用 BF16 +autocast，所有可训练 Predictor 参数及 AdamW moments 保持 FP32，loss 使用 FP32，并开启现有 +non-reentrant gradient checkpointing。冻结参数和 Teacher 权重仍为 BF16。每个 chunk backward 后 +立即释放三步图，history、previous hidden 和 Context K/V 全部 detach。 + +checkpoint 分成两种精度语义：`training_latest.pt` 保留 FP32 trainable master 参数、FP32 optimizer +状态和 scheduler，用于无损 resume；`predictor_step_XXXXX.safetensors` 显式转换为 BF16，只用于推理 +和评测。resume 时会检查 checkpoint 的 FP32-master 配置，禁止把旧 BF16 rollout optimizer 静默载入 +当前训练。 + +建议从已有 v4 权重开始一个全新的 optimizer/scheduler，不加载 Teacher-forcing +`training_latest.pt` 中的 optimizer 状态。rollout checkpoint 自身仍保存 model、optimizer、scheduler、 +global optimizer step、epoch 和 sampler/trajectory 进度，并每 100 optimizer steps 保存一次。 + +低学习率基线采用 fusion/residual `1e-5`、blocks `1e-6` 和 50-step warmup，运行至 step 600 后， +完整 epoch mean loss 只下降约 3%，因此停止。后续正式配置改为与第一阶段相同的 LR 策略： +fusion/residual 峰值 `1e-4`，blocks 峰值 `1e-5`；fusion 前 100 steps warmup，blocks 前 100 steps +学习率为 0，之后两组共同 cosine decay 到 step 2000。这里的一个 step 指完整 8-rank trajectory +更新，不是一个 chunk。 + +每 100 steps 保存一次 checkpoint。训练入口持续断言每步 8 个 rank 的 trajectory 主键互异，并检查 +step 2/3 中间 latent 的梯度非零。 + +正式启动前的 8 卡完整 8-chunk smoke test 已通过：fresh step 1 的 loss 为 `0.013256`， +`latent_grad_l1_min=1.63e-4`，峰值显存 `49.13 GiB/GPU`；从该 checkpoint resume 到 step 2 后， +`latent_grad_l1_min=2.98e-4`，峰值显存 `50.60 GiB/GPU`。训练 checkpoint 中 48 个可训练参数张量及 +Adam moments 均为 FP32，推理 safetensors 中 80 个张量均为 BF16。 + +为区分学习率问题与梯度/模型容量问题，另固定随机种子 `20260721` 从 100 个训练 case 中选择 +`[4,11,17,28,58,80,94,96]`，每个 case 保留四种 action，共 32 条 trajectory。8 卡每个 epoch +恰好 4 steps；临时将 fusion warmup 和 block freeze 都缩短为 8 steps，LR 仍为 `1e-4/1e-5`。 +截至 epoch 5，完整 overfit mean loss 从 `0.01969` 降至 `0.01836`，下降约 6.8%，证明递归梯度链路 +和优化器能够学习。该实验输出到： + +```text +checkpoints/predictor_v4_rollout_fppp_blocks1-52_overfit8_seed20260721_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup8_bwarmup8_swanlab +``` + +### 8.11 分阶段验收 + +实现必须依次通过： + +1. Dataset 只有 400 条互异的 `(case_id, action_id)` 元数据，8 个 rank 每步读取不同主键； +2. Chunk 0 全部由 Full Teacher 计算、无参数梯度，并产生 step 1/2/3 previous hidden； +3. step 2 anchor 与 `pred_h1`、step 3 anchor 与 `pred_h2` 数值和对象来源正确； +4. 每个 Teacher target 的输入等于当前 student latent 的 detached copy，不读取任何离线 FFFF 状态； +5. self history 生成的 Context latent 正确，selected indices 与正式推理函数一致； +6. 训练 Euler update 与正式 `FlowMatchDiscreteScheduler.step` 在四个 step 上误差满足 FP32 容差； +7. 只保留 step 3 loss 时，`x2/x3.grad` 非零，且前两次 Predictor 调用对参数有非零贡献； +8. Chunk 2 loss 不会回传到 Chunk 1 final latent 或 previous hidden； +9. Chunk 1～6 不触发 all-reduce，Chunk 7 每个 optimizer step 只同步一次； +10. 单卡小样本中，分 chunk backward 与保留七个 chunk 后一次 backward 的梯度在容差内一致； +11. 先完成 1-trajectory forward、1-step backward、8 卡 2-step DDP smoke test，再做 20-step overfit； +12. 无 NaN/Inf，step 1/2/3 loss 均能下降，显存峰值有日志记录。 + +最终评测必须使用 Chunk 0 `F-F-F-F`、Chunk 1～7 `F-P-P-P`，即 +`replace_steps=(1,2,3)`。在同一训练集和 25/50 holdout 上对比：Full-DiT、Direct Reuse、rollout +微调前 v4 checkpoint 和微调后 checkpoint，并同时报告视频质量、端到端耗时和组件耗时。只有留出集 +F-P-P-P 优于微调前权重，才能说明 self-rollout 微调带来泛化收益。 + +## 9. 训练入口速查 + +v4： + +```bash +bash tools/launch_predictor_v4_training.sh +``` + +v4 F-P-P-P on-policy rollout： + +```bash +bash tools/launch_predictor_v4_rollout_training.sh +``` + +resume： + +```bash +bash tools/launch_predictor_v4_rollout_training.sh \ + checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab \ + checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt +``` + +8-case overfit： + +```bash +bash tools/launch_predictor_v4_rollout_overfit8_training.sh +``` + +predictor-disca `[1,52]`： + +```bash +bash tools/launch_predictor_disca_training.sh +``` + +predictor-disca `[0,53]`： + +```bash +bash tools/launch_predictor_disca_training.sh \ + checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab \ + 0,53 +``` diff --git a/train_predictor.py b/train_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..d4948962efbea6f24dffaa0ad998c0dd01cb5bfa --- /dev/null +++ b/train_predictor.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 +"""Single-GPU training for the dense image/text HY-WorldPlay Predictor.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import time +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader + +from hyvideo.commons.infer_state import initialize_infer_state +from models import HYWorldPlayPredictor +from predictor_training.checkpoint import ( + load_predictor_weights, + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.dataset import ( + PredictorPairDataset, + move_pair_to_device, + predictor_collate, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", default="datasets/predictor_v1/train_eval_manifest.jsonl") + parser.add_argument( + "--teacher_checkpoint", + default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors", + ) + parser.add_argument( + "--output_dir", + default=( + "checkpoints/predictor_v1_blocks1-52_bf16_1gpu_" + "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100" + ), + ) + parser.add_argument("--max_steps", type=int, default=2000) + parser.add_argument("--gradient_accumulation_steps", type=int, default=4) + parser.add_argument("--warmup_blocks_steps", type=int, default=100) + parser.add_argument("--scheduler_warmup_steps", type=int, default=100) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_every", type=int, default=10) + parser.add_argument("--save_every", type=int, default=500) + parser.add_argument("--keep_snapshots", type=int, default=2) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num_workers", type=int, default=0) + parser.add_argument("--max_records", type=int, default=None) + parser.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--save_optimizer_state", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--resume", default=None) + return parser.parse_args() + + +def initialize_runtime_state() -> None: + class Args: + sage_blocks_range = "0-53" + use_sageattn = False + enable_torch_compile = False + use_fp8_gemm = False + quant_type = "fp8-per-block" + include_patterns = "double_blocks" + use_vae_parallel = False + + initialize_infer_state(Args()) + + +def make_scheduler(optimizer: torch.optim.Optimizer, warmup: int, total: int) -> LambdaLR: + def factor(step: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = min(1.0, float(step - warmup) / max(1, total - warmup)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + return LambdaLR(optimizer, factor) + + +def set_blocks_trainable(model: HYWorldPlayPredictor, enabled: bool) -> None: + model.double_blocks.requires_grad_(enabled) + + +def prune_snapshots(output_dir: Path, keep: int) -> None: + snapshots = sorted(output_dir.glob("predictor_step_*.safetensors")) + for path in snapshots[:-keep] if keep > 0 else snapshots: + path.unlink() + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + initialize_runtime_state() + device = torch.device("cuda:0") + dtype = torch.bfloat16 + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + log_path = output_dir / "train_log.jsonl" + + dataset = PredictorPairDataset(args.manifest, max_records=args.max_records) + generator = torch.Generator().manual_seed(args.seed) + loader = DataLoader( + dataset, + batch_size=1, + shuffle=True, + num_workers=args.num_workers, + pin_memory=True, + collate_fn=predictor_collate, + generator=generator, + ) + + model = HYWorldPlayPredictor() + model.load_teacher_initialization(args.teacher_checkpoint) + model.to(device=device, dtype=dtype) + model.enable_gradient_checkpointing(args.gradient_checkpointing) + model.train() + + optimizer = AdamW( + [ + { + "params": list(model.img_fusion.parameters()) + + list(model.txt_fusion.parameters()) + + list(model.residual_out.parameters()), + "lr": args.fusion_lr, + }, + {"params": model.double_blocks.parameters(), "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), + weight_decay=args.weight_decay, + ) + scheduler = make_scheduler(optimizer, args.scheduler_warmup_steps, args.max_steps) + global_step = 0 + micro_step = 0 + set_blocks_trainable(model, args.warmup_blocks_steps <= 0) + + if args.resume: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + model.load_state_dict(state["model"], strict=False) + optimizer.load_state_dict(state["optimizer"]) + scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + micro_step = int(state.get("micro_step", global_step * args.gradient_accumulation_steps)) + if "torch_rng_state" in state: + torch.set_rng_state(state["torch_rng_state"]) + if "cuda_rng_state" in state: + torch.cuda.set_rng_state_all(state["cuda_rng_state"]) + set_blocks_trainable(model, global_step >= args.warmup_blocks_steps) + + print( + json.dumps( + { + "dataset_pairs": len(dataset), + "trainable_parameters": model.trainable_parameter_count(), + "breakdown": model.trainable_parameter_breakdown(), + "start_step": global_step, + }, + indent=2, + ), + flush=True, + ) + + optimizer.zero_grad(set_to_none=True) + running: list[dict[str, float | str]] = [] + epoch = 0 + while global_step < args.max_steps: + epoch += 1 + for item in loader: + if global_step >= args.max_steps: + break + if global_step == args.warmup_blocks_steps: + set_blocks_trainable(model, True) + + inputs, target_hidden, target_velocity = move_pair_to_device( + item, device=device, dtype=dtype + ) + started = time.perf_counter() + with torch.autocast(device_type="cuda", dtype=dtype): + output = model(**inputs) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), target_hidden.float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), target_velocity.float() + ) + loss = ( + args.hidden_weight * hidden_loss + + args.velocity_weight * velocity_loss + ) + scaled_loss = loss / args.gradient_accumulation_steps + scaled_loss.backward() + micro_step += 1 + running.append( + { + "pair_id": item["pair_id"], + "loss": float(loss.detach()), + "hidden_mse": float(hidden_loss.detach()), + "velocity_mse": float(velocity_loss.detach()), + "micro_time_s": time.perf_counter() - started, + } + ) + del output, loss, scaled_loss, hidden_loss, velocity_loss + + if micro_step % args.gradient_accumulation_steps: + continue + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + global_step += 1 + + if global_step % args.log_every == 0 or global_step == 1: + torch.cuda.synchronize() + count = len(running) + record = { + "global_step": global_step, + "epoch": epoch, + "loss": sum(float(x["loss"]) for x in running) / count, + "hidden_mse": sum(float(x["hidden_mse"]) for x in running) / count, + "velocity_mse": sum(float(x["velocity_mse"]) for x in running) / count, + "grad_norm": float(grad_norm), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated() / 2**30, + "avg_micro_time_s": sum(float(x["micro_time_s"]) for x in running) / count, + } + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + print(json.dumps(record, sort_keys=True), flush=True) + running.clear() + + should_save = global_step % args.save_every == 0 or global_step == args.max_steps + if should_save: + weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" + save_predictor_weights(model, weights_path) + if args.save_optimizer_state: + latest = { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "global_step": global_step, + "micro_step": micro_step, + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state_all(), + "args": vars(args), + "predictor_config": model.config_dict, + "weights_path": str(weights_path), + } + save_training_state(output_dir / "training_latest.pt", latest) + prune_snapshots(output_dir, args.keep_snapshots) + + +if __name__ == "__main__": + main() diff --git a/train_predictor_disca.py b/train_predictor_disca.py new file mode 100644 index 0000000000000000000000000000000000000000..b4057e2380a3ef0c04d83b4e56b2ceb9c904e8a9 --- /dev/null +++ b/train_predictor_disca.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +"""Eight-GPU training entry point for the project-local Predictor-DisCa.""" + +from __future__ import annotations + +import sys + +from train_predictor_v4 import main + + +def _has_option(name: str) -> bool: + return any(argument == name or argument.startswith(f"{name}=") for argument in sys.argv[1:]) + + +if __name__ == "__main__": + defaults = ( + ("--predictor_variant", "disca"), + ( + "--output_dir", + ( + "checkpoints/predictor_disca_blocks0-53_bf16_8gpu_" + "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab" + ), + ), + ("--source_blocks", "0,53"), + ("--swanlab_project", "HY-WorldPlay-Predictor-DisCa"), + ( + "--swanlab_experiment", + ( + "predictor-disca-blocks0-53-bf16-8gpu-flr1e-4-blr1e-5-" + "schedsteps2000-swarmup100-bwarmup100" + ), + ), + ) + injected: list[str] = [] + for option, value in defaults: + if not _has_option(option): + injected.extend((option, value)) + if not _has_option("--bucket_by_context") and not _has_option("--no-bucket_by_context"): + injected.append("--no-bucket_by_context") + sys.argv[1:1] = injected + main() diff --git a/train_predictor_v2.py b/train_predictor_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..b92f3f9ad49252cf05882dd6938df5e492e26003 --- /dev/null +++ b/train_predictor_v2.py @@ -0,0 +1,374 @@ +#!/usr/bin/env python3 +"""Single-GPU training for AR-context HY-WorldPlay Predictor-v2.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import time +from pathlib import Path + +import swanlab +import torch +import torch.nn.functional as F +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader + +from hyvideo.commons.infer_state import initialize_infer_state +from models import HYWorldPlayPredictorV2, HYWorldPlayPredictorV3, PredictorV2Config +from predictor_training.checkpoint import ( + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.dataset_v2 import ( + PredictorV2PairDataset, + move_v2_pair_to_device, + predictor_v2_collate, +) +from predictor_training.dataset_v3 import ( + PredictorV3PairDataset, + move_v3_pair_to_device, + predictor_v3_collate, +) + + +def parse_source_blocks(value: str) -> tuple[int, int]: + result = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + if result not in {(0, 53), (1, 52)}: + raise argparse.ArgumentTypeError("source_blocks must be 0,53 or 1,52") + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", default="datasets/predictor_v2/train_eval_manifest.jsonl") + parser.add_argument("--predictor_version", choices=("v2", "v3"), default="v2") + parser.add_argument("--source_blocks", type=parse_source_blocks, required=True) + parser.add_argument( + "--teacher_checkpoint", + default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors", + ) + parser.add_argument("--output_dir", required=True) + parser.add_argument("--max_steps", type=int, default=2000) + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--gradient_accumulation_steps", type=int, default=4) + parser.add_argument("--warmup_blocks_steps", type=int, default=100) + parser.add_argument("--scheduler_warmup_steps", type=int, default=100) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_every", type=int, default=10) + parser.add_argument("--save_every", type=int, default=500) + parser.add_argument("--keep_snapshots", type=int, default=2) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num_workers", type=int, default=0) + parser.add_argument("--max_records", type=int, default=None) + parser.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--save_optimizer_state", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--resume", default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v2") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument("--swanlab_experiment", default=None) + parser.add_argument("--swanlab_mode", choices=("cloud", "offline", "local", "disabled"), default="cloud") + return parser.parse_args() + + +def initialize_runtime_state() -> None: + class Args: + sage_blocks_range = "0-53" + use_sageattn = False + enable_torch_compile = False + use_fp8_gemm = False + quant_type = "fp8-per-block" + include_patterns = "double_blocks" + use_vae_parallel = False + + initialize_infer_state(Args()) + + +def make_scheduler(optimizer: torch.optim.Optimizer, warmup: int, total: int) -> LambdaLR: + def factor(step: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = min(1.0, float(step - warmup) / max(1, total - warmup)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + return LambdaLR(optimizer, factor) + + +def set_ar_blocks_trainable(model: HYWorldPlayPredictorV2, enabled: bool) -> None: + for block in model.double_blocks: + for name, parameter in block.named_parameters(): + parameter.requires_grad_(enabled and not name.startswith("txt_")) + + +def prune_snapshots(output_dir: Path, keep: int) -> None: + snapshots = sorted(output_dir.glob("predictor_step_*.safetensors")) + for path in snapshots[:-keep] if keep > 0 else snapshots: + path.unlink() + + +def init_swanlab(args: argparse.Namespace, config: dict): + if args.swanlab_mode == "cloud": + api_key = os.environ.get("SWANLAB_API_KEY") + swanlab.login(api_key=api_key, save=False) if api_key else swanlab.login() + fusion_lr = f"{args.fusion_lr:.0e}".replace("e-0", "e-").replace("e+0", "e+") + blocks_lr = f"{args.blocks_lr:.0e}".replace("e-0", "e-").replace("e+0", "e+") + experiment = args.swanlab_experiment or ( + f"predictor-{args.predictor_version}-blocks" + + "-".join(str(value) for value in args.source_blocks) + + f"-flr{fusion_lr}-blr{blocks_lr}-schedsteps{args.max_steps}" + + f"-swarmup{args.scheduler_warmup_steps}" + + f"-bwarmup{args.warmup_blocks_steps}" + ) + return swanlab.init( + project=args.swanlab_project, + workspace=args.swanlab_workspace, + experiment_name=experiment, + description=( + f"HY-WorldPlay {args.predictor_version} AR-context Predictor, " + "fresh exact 10-case targets" + ), + tags=[ + f"predictor-{args.predictor_version}", + "ar-context-kv", + f"blocks-{args.source_blocks[0]}-{args.source_blocks[1]}", + ], + config=config, + logdir=str(Path(args.output_dir).resolve() / "swanlab"), + mode=args.swanlab_mode, + ) + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + initialize_runtime_state() + device = torch.device("cuda:0") + dtype = torch.bfloat16 + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + log_path = output_dir / "train_log.jsonl" + + dataset_class = ( + PredictorV3PairDataset if args.predictor_version == "v3" else PredictorV2PairDataset + ) + collate_fn = ( + predictor_v3_collate if args.predictor_version == "v3" else predictor_v2_collate + ) + move_pair = ( + move_v3_pair_to_device + if args.predictor_version == "v3" + else move_v2_pair_to_device + ) + dataset = dataset_class( + args.manifest, + source_block_ids=args.source_blocks, + max_records=args.max_records, + ) + generator = torch.Generator().manual_seed(args.seed) + loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=False, + collate_fn=collate_fn, + generator=generator, + ) + + model_class = ( + HYWorldPlayPredictorV3 if args.predictor_version == "v3" else HYWorldPlayPredictorV2 + ) + model = model_class(PredictorV2Config(source_block_ids=args.source_blocks)) + model.load_teacher_initialization(args.teacher_checkpoint) + model.to(device=device, dtype=dtype) + model.enable_gradient_checkpointing(args.gradient_checkpointing) + model.train() + + fusion_parameters = list(model.img_fusion.parameters()) + list(model.residual_out.parameters()) + block_parameters = [ + parameter for parameter in model.double_blocks.parameters() if parameter.requires_grad + ] + intended_trainable_parameters = sum( + parameter.numel() for parameter in fusion_parameters + block_parameters + ) + intended_parameter_breakdown = model.trainable_parameter_breakdown() + optimizer = AdamW( + [ + {"params": fusion_parameters, "lr": args.fusion_lr}, + {"params": block_parameters, "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), + weight_decay=args.weight_decay, + ) + scheduler = make_scheduler(optimizer, args.scheduler_warmup_steps, args.max_steps) + global_step = 0 + micro_step = 0 + set_ar_blocks_trainable(model, args.warmup_blocks_steps <= 0) + + if args.resume: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != args.source_blocks: + raise ValueError("Resume checkpoint source blocks do not match") + model.load_state_dict(state["model"], strict=False) + optimizer.load_state_dict(state["optimizer"]) + scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + micro_step = int(state.get("micro_step", global_step * args.gradient_accumulation_steps)) + if "torch_rng_state" in state: + torch.set_rng_state(state["torch_rng_state"]) + if "cuda_rng_state" in state: + torch.cuda.set_rng_state_all(state["cuda_rng_state"]) + set_ar_blocks_trainable(model, global_step >= args.warmup_blocks_steps) + + run_config = { + **vars(args), + "source_blocks": list(args.source_blocks), + "manifest": str(Path(args.manifest).resolve()), + "teacher_checkpoint": str(Path(args.teacher_checkpoint).resolve()), + "dataset_pairs": len(dataset), + "effective_batch_size": args.batch_size * args.gradient_accumulation_steps, + "trainable_parameters": intended_trainable_parameters, + "parameter_breakdown": intended_parameter_breakdown, + "dtype": "bfloat16", + "predictor_version": args.predictor_version, + "gpu": torch.cuda.get_device_name(device), + } + run = init_swanlab(args, run_config) + print( + json.dumps( + { + "dataset_pairs": len(dataset), + "source_blocks": args.source_blocks, + "trainable_parameters": intended_trainable_parameters, + "breakdown": intended_parameter_breakdown, + "start_step": global_step, + "swanlab_run_id": getattr(run, "id", None), + }, + indent=2, + ), + flush=True, + ) + + optimizer.zero_grad(set_to_none=True) + running: list[dict[str, float | str]] = [] + epoch = 0 + try: + while global_step < args.max_steps: + epoch += 1 + for item in loader: + if global_step >= args.max_steps: + break + if global_step == args.warmup_blocks_steps: + set_ar_blocks_trainable(model, True) + + inputs, target_hidden, target_velocity = move_pair( + item, device=device, dtype=dtype + ) + started = time.perf_counter() + with torch.autocast(device_type="cuda", dtype=dtype): + output = model(**inputs) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), target_hidden.float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), target_velocity.float() + ) + loss = args.hidden_weight * hidden_loss + args.velocity_weight * velocity_loss + scaled_loss = loss / args.gradient_accumulation_steps + scaled_loss.backward() + micro_step += 1 + running.append( + { + "pair_id": ",".join(item["pair_id"]), + "loss": float(loss.detach()), + "hidden_mse": float(hidden_loss.detach()), + "velocity_mse": float(velocity_loss.detach()), + "micro_time_s": time.perf_counter() - started, + "context_frames": sum(item["context_frames"]) + / len(item["context_frames"]), + } + ) + del ( + output, + loss, + scaled_loss, + hidden_loss, + velocity_loss, + inputs, + target_hidden, + target_velocity, + item, + ) + + if micro_step % args.gradient_accumulation_steps: + continue + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + global_step += 1 + + if global_step % args.log_every == 0 or global_step == 1: + torch.cuda.synchronize() + count = len(running) + record = { + "global_step": global_step, + "epoch": epoch, + "loss": sum(float(x["loss"]) for x in running) / count, + "hidden_mse": sum(float(x["hidden_mse"]) for x in running) / count, + "velocity_mse": sum(float(x["velocity_mse"]) for x in running) / count, + "grad_norm": float(grad_norm), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated() / 2**30, + "avg_micro_time_s": sum(float(x["micro_time_s"]) for x in running) / count, + "avg_context_frames": sum(float(x["context_frames"]) for x in running) / count, + } + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + swanlab.log(record, step=global_step) + print(json.dumps(record, sort_keys=True), flush=True) + running.clear() + + should_save = global_step % args.save_every == 0 or global_step == args.max_steps + if should_save: + weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" + save_predictor_weights(model, weights_path) + if args.save_optimizer_state: + latest = { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "global_step": global_step, + "micro_step": micro_step, + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state_all(), + "args": vars(args), + "predictor_config": model.config_dict, + "weights_path": str(weights_path), + } + save_training_state(output_dir / "training_latest.pt", latest) + prune_snapshots(output_dir, args.keep_snapshots) + except BaseException as error: + swanlab.finish(error=error) + raise + else: + swanlab.finish() + + +if __name__ == "__main__": + main() diff --git a/train_predictor_v3.py b/train_predictor_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..c607c5523c0e1106baa795ed861d0e5dc50197a8 --- /dev/null +++ b/train_predictor_v3.py @@ -0,0 +1,14 @@ +#!/usr/bin/env python3 +"""Single-GPU training entry point for previous-chunk Predictor-v3.""" + +from __future__ import annotations + +import sys + +from train_predictor_v2 import main + + +if __name__ == "__main__": + if "--predictor_version" not in sys.argv: + sys.argv[1:1] = ["--predictor_version", "v3"] + main() diff --git a/train_predictor_v4.py b/train_predictor_v4.py new file mode 100644 index 0000000000000000000000000000000000000000..13f761e81d47c9e31cd2e7e9c344050dc5b68a5c --- /dev/null +++ b/train_predictor_v4.py @@ -0,0 +1,1030 @@ +#!/usr/bin/env python3 +"""Eight-GPU DDP training engine for Predictor-v4 and Predictor-DisCa.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +import time +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.nn.functional as F +import swanlab +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from hyvideo.commons.infer_state import initialize_infer_state +from hyvideo.commons.parallel_states import initialize_parallel_state +from models import HYWorldPlayPredictorDisCa, HYWorldPlayPredictorV4, PredictorV2Config +from predictor_training.bucket_sampler import DistributedContextBucketBatchSampler +from predictor_training.checkpoint import ( + load_concat_fusion_initialization, + load_concat_residual_gate_initialization, + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.dataset_prefeature import ( + PredictorPrefeaturePairDataset, + move_prefeature_pair_to_device, + predictor_prefeature_collate, +) +from predictor_training.dataset_disca import ( + PredictorDisCaPairDataset, + move_disca_pair_to_device, + predictor_disca_collate, +) + + +DEFAULT_MANIFEST = ( + "/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4/" + "train_eval_manifest.jsonl" +) +DEFAULT_TEACHER = ( + "/mnt/s3files/s3-us-west2-default/zoubin/cz/checkpoints/hy_worldplay/" + "huggingface/hub/models--tencent--HY-WorldPlay/snapshots/" + "f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/" + "diffusion_pytorch_model.safetensors" +) +DEFAULT_OUTPUT = ( + "/mnt/s3files/s3-us-west2-default/zoubin/cz/projects/" + "HY-WorldPlay-DEV-Predictor/checkpoints/" + "predictor_v4_prefeature_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab" +) + + +def parse_source_blocks(value: str) -> tuple[int, int]: + result = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + if result not in {(0, 53), (1, 52)}: + raise argparse.ArgumentTypeError("source_blocks must be 0,53 or 1,52") + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", default=DEFAULT_MANIFEST) + parser.add_argument("--teacher_checkpoint", default=DEFAULT_TEACHER) + parser.add_argument("--output_dir", default=DEFAULT_OUTPUT) + parser.add_argument("--source_blocks", type=parse_source_blocks, default=(0, 53)) + parser.add_argument( + "--predictor_variant", choices=("v4", "disca"), default="v4", + help="Use v4 Context/previous-chunk inputs or the Predictor-DisCa ablation", + ) + parser.add_argument( + "--previous_condition_mode", + choices=("none", "same_step", "previous_step", "both", "both_delta_gated"), + default="same_step", + help="Explicit previous-chunk hidden condition used by Predictor-v4", + ) + parser.add_argument( + "--predictor_fusion_type", + choices=( + "concat", + "projection_add", + "token_gate_prev", + "concat_residual_gate_prev", + ), + default="concat", + help="Image-token fusion architecture; non-concat modes require same_step", + ) + parser.add_argument("--gate_hidden_dim", type=int, default=256) + parser.add_argument("--gate_initial_probability", type=float, default=0.8) + parser.add_argument("--predictor_arch", choices=("baseline", "atc"), default="baseline") + parser.add_argument( + "--atc_variant", + choices=("full", "anchor_only", "raw_prev", "no_gate", "no_distance"), + default="full", + ) + parser.add_argument("--atc_mlp_hidden_dim", type=int, default=4096) + parser.add_argument("--atc_gate_hidden_dim", type=int, default=512) + parser.add_argument("--atc_gate_initial_probability", type=float, default=0.3) + parser.add_argument( + "--atc_transport_heads", type=int, default=None, + help="Transport attention heads; default = Teacher heads so head_dim matches the Teacher RoPE.", + ) + parser.add_argument( + "--atc_transport_mode", choices=("global", "local"), default="global", + help="global: Causal-Forcing RoPE'd global cross-chunk attention; local: legacy 3x3 window.", + ) + parser.add_argument( + "--atc_previous_scope", choices=("chunk", "last_frame"), default="chunk", + help="Keys of the global transport: whole previous chunk or only its last frame.", + ) + parser.add_argument("--atc_transport_window", type=int, default=3) + parser.add_argument("--atc_transport_query_chunk_size", type=int, default=128) + parser.add_argument("--atc_transport_initial_scale", type=float, default=0.1) + parser.add_argument("--max_steps", type=int, default=2000) + parser.add_argument("--batch_size", type=int, default=16, help="Per-rank batch size") + parser.add_argument("--gradient_accumulation_steps", type=int, default=1) + parser.add_argument("--block_warmup_steps", type=int, default=100) + parser.add_argument( + "--block_warmup_mode", + choices=("delayed", "linear"), + default="delayed", + help=( + "Use the original Stage-1 delayed block start or linearly warm the " + "block group from the beginning" + ), + ) + parser.add_argument("--scheduler_warmup_steps", type=int, default=100) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_every", type=int, default=10) + parser.add_argument("--save_every", type=int, default=100) + parser.add_argument("--keep_snapshots", type=int, default=20) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num_workers", type=int, default=0) + parser.add_argument("--max_records", type=int, default=None) + parser.add_argument( + "--bucket_by_context", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--fp32_trainable_params", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Keep requires_grad=True Predictor parameters and AdamW moments in " + "FP32 while frozen parameters and forward activations remain BF16" + ), + ) + parser.add_argument( + "--train_delta_only", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Freeze the inherited same-step Predictor and optimize only the " + "both_delta_gated residual branch." + ), + ) + parser.add_argument( + "--train_fusion_residual_only", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Freeze the inherited concat fusion and Predictor blocks, and train " + "only the zero-initialized concat_residual_gate_prev branch." + ), + ) + parser.add_argument( + "--save_optimizer_state", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--init_training_state", + default=None, + help=( + "Load only the FP32 Predictor model state from training_latest.pt " + "while starting a fresh optimizer, scheduler, and step counter" + ), + ) + parser.add_argument("--resume", default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v4") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument( + "--swanlab_experiment", + default=( + "predictor-v4-prefeature-blocks0-53-bf16-8gpu-" + "flr1e-4-blr1e-5-schedsteps2000-swarmup100-bwarmup100" + ), + ) + parser.add_argument( + "--swanlab_mode", + choices=("cloud", "offline", "local", "disabled"), + default="cloud", + ) + return parser.parse_args() + + +def initialize_runtime_state() -> None: + class Args: + sage_blocks_range = "0-53" + use_sageattn = False + enable_torch_compile = False + use_fp8_gemm = False + quant_type = "fp8-per-block" + include_patterns = "double_blocks" + use_vae_parallel = False + + initialize_infer_state(Args()) + + +def cosine_factor(step: int, warmup: int, total: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = min(1.0, float(step - warmup) / max(1, total - warmup)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + +def make_scheduler(optimizer: torch.optim.Optimizer, args: argparse.Namespace) -> LambdaLR: + def block_factor(step: int) -> float: + if args.block_warmup_mode == "linear": + return cosine_factor(step, args.block_warmup_steps, args.max_steps) + return ( + 0.0 + if step < args.block_warmup_steps + else cosine_factor(step, args.scheduler_warmup_steps, args.max_steps) + ) + + return LambdaLR( + optimizer, + [ + lambda step: cosine_factor(step, args.scheduler_warmup_steps, args.max_steps), + block_factor, + ], + ) + + +def prune_snapshots(output_dir: Path, keep: int) -> None: + snapshots = sorted(output_dir.glob("predictor_step_*.safetensors")) + for path in snapshots[:-keep] if keep > 0 else snapshots: + path.unlink() + + +def distributed_mean(values: torch.Tensor) -> torch.Tensor: + dist.all_reduce(values, op=dist.ReduceOp.SUM) + values /= dist.get_world_size() + return values + + +def configure_model_precision( + model: torch.nn.Module, + *, + device: torch.device, + activation_dtype: torch.dtype, + fp32_trainable_params: bool, +) -> None: + """Place frozen state in BF16 and optionally retain FP32 master parameters.""" + model.to(device=device, dtype=activation_dtype) + if fp32_trainable_params: + with torch.no_grad(): + for parameter in model.parameters(): + if parameter.requires_grad: + parameter.data = parameter.data.to(dtype=torch.float32) + + +def validate_optimizer_precision( + optimizer: torch.optim.Optimizer, + *, + fp32_trainable_params: bool, + require_state: bool, +) -> None: + expected = torch.float32 if fp32_trainable_params else torch.bfloat16 + for group in optimizer.param_groups: + for parameter in group["params"]: + if parameter.dtype != expected: + raise RuntimeError( + f"Optimizer parameter dtype {parameter.dtype} != expected {expected}" + ) + state = optimizer.state.get(parameter, {}) + for name in ("exp_avg", "exp_avg_sq"): + value = state.get(name) + if value is None: + if require_state: + raise RuntimeError(f"Optimizer state {name} has not been initialized") + continue + if value.dtype != expected: + raise RuntimeError( + f"Optimizer state {name} dtype {value.dtype} != expected {expected}" + ) + + +def configure_delta_only_trainability(model: torch.nn.Module) -> None: + if getattr(model, "previous_condition_mode", None) != "both_delta_gated": + raise ValueError("--train_delta_only requires both_delta_gated mode") + model.requires_grad_(False) + model.img_fusion.previous_step_norm.requires_grad_(True) + model.img_fusion.delta_mlp.requires_grad_(True) + model.img_fusion.delta_gate.requires_grad_(True) + + +def configure_fusion_residual_only_trainability(model: torch.nn.Module) -> None: + fusion_type = getattr( + getattr(model, "predictor_config", None), "predictor_fusion_type", None + ) + if fusion_type != "concat_residual_gate_prev": + raise ValueError( + "--train_fusion_residual_only requires " + "concat_residual_gate_prev fusion" + ) + model.requires_grad_(False) + model.img_fusion.gate_p.requires_grad_(True) + model.img_fusion.residual_mlp.requires_grad_(True) + + +def load_initial_training_state( + model: torch.nn.Module, + path: Path, + *, + source_blocks: tuple[int, int], + predictor_variant: str, + previous_condition_mode: str, + predictor_fusion_type: str, + fp32_trainable_params: bool, +) -> None: + """Load only a prior FP32 Predictor model state for a fresh training branch.""" + state = torch.load(path, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != source_blocks: + raise ValueError("Initialization checkpoint source blocks do not match") + saved_variant = state.get("args", {}).get("predictor_variant", "v4") + if saved_variant != predictor_variant: + raise ValueError("Initialization checkpoint predictor variant does not match") + saved_mode = state.get("args", {}).get( + "previous_condition_mode", + state.get("predictor_config", {}).get("previous_condition_mode", "same_step"), + ) + delta_from_same_step = ( + saved_mode == "same_step" and previous_condition_mode == "both_delta_gated" + ) + if saved_mode != previous_condition_mode and not delta_from_same_step: + raise ValueError( + "Initialization checkpoint previous-condition mode does not match: " + f"saved={saved_mode}, requested={previous_condition_mode}" + ) + saved_fp32_master = bool( + state.get("args", {}).get("fp32_trainable_params", False) + ) + if not saved_fp32_master or not fp32_trainable_params: + raise ValueError( + "Model-only initialization requires both the checkpoint and current " + "run to use FP32 trainable parameters" + ) + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + state.get("args", {}).get("predictor_fusion_type", "concat"), + ) + concat_conversion = ( + saved_fusion_type == "concat" + and predictor_fusion_type in {"projection_add", "token_gate_prev"} + and saved_mode == previous_condition_mode == "same_step" + ) + concat_residual_initialization = ( + saved_fusion_type == "concat" + and predictor_fusion_type == "concat_residual_gate_prev" + and saved_mode == previous_condition_mode == "same_step" + ) + if ( + saved_fusion_type != predictor_fusion_type + and not concat_conversion + and not concat_residual_initialization + ): + raise ValueError( + "Initialization checkpoint fusion type does not match: " + f"saved={saved_fusion_type}, requested={predictor_fusion_type}" + ) + if concat_conversion: + conversion = load_concat_fusion_initialization(model, state["model"]) + print( + json.dumps( + { + "fusion_initialization": "concat_conversion", + "discarded_concat_key_count": len( + conversion["discarded_concat_keys"] + ), + "fresh_fusion_key_count": len(conversion["fresh_fusion_keys"]), + }, + sort_keys=True, + ), + flush=True, + ) + return + if concat_residual_initialization: + conversion = load_concat_residual_gate_initialization( + model, state["model"] + ) + print( + json.dumps( + { + "fusion_initialization": "concat_residual_gate", + "fresh_residual_key_count": len( + conversion["fresh_residual_keys"] + ), + }, + sort_keys=True, + ), + flush=True, + ) + return + + result = model.load_state_dict(state["model"], strict=False) + trainable_names = { + name for name, parameter in model.named_parameters() if parameter.requires_grad + } + allowed_missing = set() + if delta_from_same_step: + allowed_missing = { + name + for name in trainable_names + if name.startswith("img_fusion.previous_step_norm.") + or name.startswith("img_fusion.delta_mlp.") + or name == "img_fusion.delta_gate" + } + with torch.no_grad(): + model.img_fusion.previous_step_norm.load_state_dict( + model.img_fusion.previous_norm.state_dict(), strict=True + ) + missing_trainable = sorted( + trainable_names.intersection(result.missing_keys) - allowed_missing + ) + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Initialization training-state mismatch; " + f"unexpected={result.unexpected_keys}, " + f"missing_trainable={missing_trainable}" + ) + + +def main() -> None: + args = parse_args() + if args.predictor_arch == "atc" and ( + args.predictor_variant != "v4" or args.previous_condition_mode != "same_step" + ): + raise ValueError("ATC requires predictor_variant=v4 and previous_condition_mode=same_step") + if args.predictor_variant != "v4" and args.previous_condition_mode != "same_step": + raise ValueError("previous_condition_mode ablations require predictor_variant=v4") + if args.predictor_variant != "v4" and args.predictor_fusion_type != "concat": + raise ValueError("Non-concat fusion types require predictor_variant=v4") + if ( + args.predictor_fusion_type != "concat" + and args.previous_condition_mode != "same_step" + ): + raise ValueError("Non-concat fusion types require previous_condition_mode=same_step") + if args.resume and args.init_training_state: + raise ValueError("--resume and --init_training_state are mutually exclusive") + if args.train_delta_only and args.train_fusion_residual_only: + raise ValueError( + "--train_delta_only and --train_fusion_residual_only are mutually exclusive" + ) + if args.train_fusion_residual_only and not ( + args.init_training_state or args.resume + ): + raise ValueError( + "--train_fusion_residual_only requires --init_training_state or --resume" + ) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + raise RuntimeError("Launch Predictor-v4 training with torchrun") + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="nccl", device_id=device) + initialize_parallel_state(sp=1) + initialize_runtime_state() + dtype = torch.bfloat16 + + seed = args.seed + rank + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + output_dir = Path(args.output_dir).resolve() + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + dist.barrier() + log_path = output_dir / "train_log.jsonl" + + if args.predictor_variant == "disca": + dataset = PredictorDisCaPairDataset( + args.manifest, + source_block_ids=args.source_blocks, + max_records=args.max_records, + ) + collate_fn = predictor_disca_collate + move_batch = move_disca_pair_to_device + else: + dataset = PredictorPrefeaturePairDataset( + args.manifest, + source_block_ids=args.source_blocks, + require_previous_chunk=True, + previous_condition_mode=args.previous_condition_mode, + max_records=args.max_records, + ) + collate_fn = predictor_prefeature_collate + move_batch = move_prefeature_pair_to_device + bucket_sampler = None + distributed_sampler = None + if args.bucket_by_context: + bucket_sampler = DistributedContextBucketBatchSampler( + dataset, + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + seed=args.seed, + drop_last=True, + ) + loader = DataLoader( + dataset, + batch_sampler=bucket_sampler, + num_workers=args.num_workers, + pin_memory=False, + collate_fn=collate_fn, + ) + else: + distributed_sampler = DistributedSampler( + dataset, + num_replicas=world_size, + rank=rank, + shuffle=True, + seed=args.seed, + drop_last=True, + ) + loader = DataLoader( + dataset, + batch_size=args.batch_size, + sampler=distributed_sampler, + drop_last=True, + num_workers=args.num_workers, + pin_memory=False, + collate_fn=collate_fn, + ) + if len(loader) == 0: + raise ValueError("Training loader has zero batches") + + model_class = ( + HYWorldPlayPredictorDisCa + if args.predictor_variant == "disca" + else HYWorldPlayPredictorV4 + ) + model = model_class( + PredictorV2Config( + source_block_ids=args.source_blocks, + previous_condition_mode=args.previous_condition_mode, + predictor_fusion_type=args.predictor_fusion_type, + gate_hidden_dim=args.gate_hidden_dim, + gate_initial_probability=args.gate_initial_probability, + predictor_arch=args.predictor_arch, + atc_variant=args.atc_variant, + atc_mlp_hidden_dim=args.atc_mlp_hidden_dim, + atc_gate_hidden_dim=args.atc_gate_hidden_dim, + atc_gate_initial_probability=args.atc_gate_initial_probability, + atc_transport_heads=args.atc_transport_heads, + atc_transport_window=args.atc_transport_window, + atc_transport_query_chunk_size=args.atc_transport_query_chunk_size, + atc_transport_initial_scale=args.atc_transport_initial_scale, + atc_transport_mode=args.atc_transport_mode, + atc_previous_scope=args.atc_previous_scope, + ) + ) + model.load_teacher_initialization(args.teacher_checkpoint) + if args.train_delta_only: + configure_delta_only_trainability(model) + if args.train_fusion_residual_only: + configure_fusion_residual_only_trainability(model) + configure_model_precision( + model, + device=device, + activation_dtype=dtype, + fp32_trainable_params=args.fp32_trainable_params, + ) + if args.init_training_state: + initialization_path = Path(args.init_training_state).resolve() + if not initialization_path.is_file(): + raise FileNotFoundError(initialization_path) + load_initial_training_state( + model, + initialization_path, + source_blocks=args.source_blocks, + predictor_variant=args.predictor_variant, + previous_condition_mode=args.previous_condition_mode, + predictor_fusion_type=args.predictor_fusion_type, + fp32_trainable_params=args.fp32_trainable_params, + ) + model.enable_gradient_checkpointing(args.gradient_checkpointing) + model.train() + if args.predictor_variant == "v4" and any( + parameter.requires_grad for parameter in model.context_projector.parameters() + ): + raise RuntimeError("Frozen Context projector unexpectedly has trainable parameters") + + fusion_parameters = [ + parameter + for parameter in ( + list(model.img_fusion.parameters()) + + list(model.residual_out.parameters()) + ) + if parameter.requires_grad + ] + block_parameters = [ + parameter + for parameter in model.double_blocks.parameters() + if parameter.requires_grad + ] + optimizer = AdamW( + [ + {"params": fusion_parameters, "lr": args.fusion_lr}, + {"params": block_parameters, "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), + weight_decay=args.weight_decay, + ) + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=False, + ) + scheduler = make_scheduler(optimizer, args) + global_step = 0 + micro_step = 0 + epoch = 0 + + if args.resume: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != args.source_blocks: + raise ValueError("Resume checkpoint source blocks do not match") + saved_variant = state.get("args", {}).get("predictor_variant", "v4") + if saved_variant != args.predictor_variant: + raise ValueError("Resume checkpoint predictor variant does not match") + saved_arch = state.get("predictor_config", {}).get( + "predictor_arch", state.get("args", {}).get("predictor_arch", "baseline") + ) + saved_atc_variant = state.get("predictor_config", {}).get( + "atc_variant", state.get("args", {}).get("atc_variant", "full") + ) + if saved_arch != args.predictor_arch: + raise ValueError( + f"Resume predictor architecture mismatch: saved={saved_arch}, " + f"requested={args.predictor_arch}" + ) + if saved_arch == "atc" and saved_atc_variant != args.atc_variant: + raise ValueError( + f"Resume ATC variant mismatch: saved={saved_atc_variant}, " + f"requested={args.atc_variant}" + ) + if saved_arch == "atc": + saved_cfg = state.get("predictor_config", {}) + # Checkpoints written before the transport switch are local-window ATC. + saved_transport = saved_cfg.get("atc_transport_mode", "local") + saved_scope = saved_cfg.get("atc_previous_scope", "chunk") + if (saved_transport, saved_scope) != (args.atc_transport_mode, args.atc_previous_scope): + raise ValueError( + f"Resume ATC transport mismatch: saved={saved_transport}/{saved_scope}, " + f"requested={args.atc_transport_mode}/{args.atc_previous_scope}" + ) + saved_mode = state.get("args", {}).get( + "previous_condition_mode", + state.get("predictor_config", {}).get( + "previous_condition_mode", "same_step" + ), + ) + if saved_mode != args.previous_condition_mode: + raise ValueError( + "Resume checkpoint previous-condition mode does not match: " + f"saved={saved_mode}, requested={args.previous_condition_mode}" + ) + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + state.get("args", {}).get("predictor_fusion_type", "concat"), + ) + if saved_fusion_type != args.predictor_fusion_type: + raise ValueError( + "Resume checkpoint fusion type does not match: " + f"saved={saved_fusion_type}, requested={args.predictor_fusion_type}" + ) + saved_fp32_master = bool( + state.get("args", {}).get("fp32_trainable_params", False) + ) + if saved_fp32_master != args.fp32_trainable_params: + raise ValueError("Resume checkpoint FP32-master setting does not match") + result = model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + key + for key in result.missing_keys + if key.startswith(("img_fusion.", "double_blocks.", "residual_out.")) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Resume model mismatch; " + f"unexpected={result.unexpected_keys}, missing={missing_trainable}" + ) + optimizer.load_state_dict(state["optimizer"]) + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + micro_step = int(state.get("micro_step", global_step)) + epoch = int(state.get("epoch", 0)) + + ddp = DDP( + model, + device_ids=[local_rank], + output_device=local_rank, + broadcast_buffers=False, + gradient_as_bucket_view=True, + find_unused_parameters=False, + ) + trainable_parameters = sum( + parameter.numel() for parameter in model.parameters() if parameter.requires_grad + ) + frozen_projector_parameters = sum( + parameter.numel() for parameter in getattr(model, "context_projector", ()).parameters() + ) if hasattr(model, "context_projector") else 0 + run_config = { + **vars(args), + "manifest": str(Path(args.manifest).resolve()), + "teacher_checkpoint": str(Path(args.teacher_checkpoint).resolve()), + "output_dir": str(output_dir), + "source_blocks": list(args.source_blocks), + "predictor_version": ( + "predictor_disca" if args.predictor_variant == "disca" else "v4_prefeature" + ), + "schema_version": ( + "predictor_disca_text_kv_v1" + if args.predictor_variant == "disca" + else "predictor_context_prefeature_bf16_v1" + ), + "supervision_pairs": [[0, 1], [1, 2], [2, 3]], + "previous_condition_mode": args.previous_condition_mode, + "previous_condition_source": "offline_full_teacher_hidden", + "records": len(dataset.records), + "dataset_pairs": len(dataset), + "steps_per_epoch_per_rank": len(loader), + "world_size": world_size, + "per_rank_batch_size": args.batch_size, + "global_batch_size": ( + args.batch_size * world_size * args.gradient_accumulation_steps + ), + "trainable_parameters": trainable_parameters, + "frozen_projector_parameters": frozen_projector_parameters, + "dtype": "bfloat16", + "activation_dtype": "bfloat16", + "trainable_parameter_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + "optimizer_moment_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + "inference_checkpoint_dtype": "bfloat16", + "gpu": torch.cuda.get_device_name(device), + } + if rank == 0: + config_path = output_dir / "train_config.json" + tmp = config_path.with_suffix(".json.tmp") + tmp.write_text( + json.dumps(run_config, indent=2, ensure_ascii=False, sort_keys=True) + "\n", + encoding="utf-8", + ) + os.replace(tmp, config_path) + print(json.dumps(run_config, indent=2, default=str), flush=True) + + swanlab_run = None + if rank == 0: + if args.swanlab_mode == "cloud": + api_key = os.environ.get("SWANLAB_API_KEY") + # Prefer an explicit ephemeral key, while also supporting the cached + # machine credential installed during environment setup. + swanlab.login(api_key=api_key, save=False) + swanlab_run = swanlab.init( + project=args.swanlab_project, + workspace=args.swanlab_workspace, + experiment_name=args.swanlab_experiment, + description=( + f"Predictor-{args.predictor_variant} training; three adjacent " + "denoise transitions; 8xH100 DDP" + ), + tags=[ + f"predictor-{args.predictor_variant}", + ( + "text-kv-only" + if args.predictor_variant == "disca" + else "context-prefeature" + ), + f"blocks-{args.source_blocks[0]}-{args.source_blocks[1]}", + f"fusion-{args.predictor_fusion_type}"[:20], + f"cond-{args.previous_condition_mode.replace('_step', 'step')}"[:20], + "batch16x8", + ], + config=run_config, + logdir=str(output_dir / "swanlab"), + mode=args.swanlab_mode, + ) + print( + json.dumps( + {"swanlab_run_id": getattr(swanlab_run, "id", None)}, + sort_keys=True, + ), + flush=True, + ) + + optimizer.zero_grad(set_to_none=True) + running = torch.zeros(18, device=device, dtype=torch.float64) + running_count = 0 + # loss, hidden, velocity, batch_time, context_frames + try: + while global_step < args.max_steps: + epoch += 1 + if bucket_sampler is not None: + bucket_sampler.set_epoch(epoch) + if distributed_sampler is not None: + distributed_sampler.set_epoch(epoch) + for item in loader: + if global_step >= args.max_steps: + break + inputs, target_hidden, target_velocity = move_batch( + item, device=device, dtype=dtype + ) + started = time.perf_counter() + sync_gradients = ( + (micro_step + 1) % args.gradient_accumulation_steps == 0 + ) + sync_context = ( + torch.enable_grad() + if sync_gradients + else ddp.no_sync() + ) + with sync_context: + with torch.autocast(device_type="cuda", dtype=dtype): + output = ddp(**inputs) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), target_hidden.float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), target_velocity.float() + ) + loss = ( + args.hidden_weight * hidden_loss + + args.velocity_weight * velocity_loss + ) + scaled_loss = loss / args.gradient_accumulation_steps + scaled_loss.backward() + micro_step += 1 + gate_stats = output.get("gate_stats", {}) + base_values = torch.tensor( + [ + float(loss.detach()), + float(hidden_loss.detach()), + float(velocity_loss.detach()), + time.perf_counter() - started, + sum(item["context_frames"]) / len(item["context_frames"]), + ], + device=device, + dtype=torch.float64, + ) + gate_names = ( + "gate_p_mean", + "gate_p_std", + "gate_p_p10", + "gate_p_p50", + "gate_p_p90", + "gate_p_near_zero", + "gate_p_near_one", + ) + gate_values = ( + torch.stack([gate_stats[name] for name in gate_names]).to( + device=device, dtype=torch.float64 + ) + if gate_stats + else torch.zeros(7, device=device, dtype=torch.float64) + ) + atc_stats = output.get("atc_stats", {}) + atc_names = ( + "delta_h_d_norm", + "delta_h_c_norm", + "gated_delta_h_c_norm", + "transport_residual_norm", + "anchor_distance_mean", + "transport_attention_entropy", + ) + atc_values = ( + torch.stack([atc_stats[name] for name in atc_names]).to( + device=device, dtype=torch.float64 + ) + if atc_stats + else torch.zeros(6, device=device, dtype=torch.float64) + ) + running += torch.cat([base_values, gate_values, atc_values]) + running_count += 1 + del output, loss, scaled_loss, hidden_loss, velocity_loss + del inputs, target_hidden, target_velocity, item + if not sync_gradients: + continue + + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), args.grad_clip + ) + optimizer.step() + if global_step == 0 or ( + args.resume is not None and global_step == int(state["global_step"]) + ): + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + scheduler.step() + optimizer.zero_grad(set_to_none=True) + global_step += 1 + + if global_step % args.log_every == 0 or global_step == 1: + torch.cuda.synchronize(device) + averaged = distributed_mean(running.clone()) + averaged /= running_count + grad_value = torch.tensor(float(grad_norm), device=device) + dist.all_reduce(grad_value, op=dist.ReduceOp.MAX) + if rank == 0: + record = { + "global_step": global_step, + "epoch": epoch, + "loss": float(averaged[0]), + "hidden_mse": float(averaged[1]), + "velocity_mse": float(averaged[2]), + "avg_micro_time_s": float(averaged[3]), + "avg_context_frames": float(averaged[4]), + "grad_norm_max": float(grad_value), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated(device) + / 2**30, + } + if args.predictor_arch == "atc" or args.predictor_fusion_type in { + "token_gate_prev", + "concat_residual_gate_prev", + }: + record.update( + { + "gate_p_mean": float(averaged[5]), + "gate_p_std": float(averaged[6]), + "gate_p_p10": float(averaged[7]), + "gate_p_p50": float(averaged[8]), + "gate_p_p90": float(averaged[9]), + "gate_p_near_zero": float(averaged[10]), + "gate_p_near_one": float(averaged[11]), + } + ) + if args.predictor_arch == "atc": + record.update( + { + name: float(averaged[12 + index]) + for index, name in enumerate(atc_names) + } + ) + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + swanlab.log(record, step=global_step) + print(json.dumps(record, sort_keys=True), flush=True) + running.zero_() + running_count = 0 + + should_save = ( + global_step % args.save_every == 0 + or global_step == args.max_steps + ) + if should_save: + dist.barrier() + if rank == 0: + weights_path = ( + output_dir + / f"predictor_step_{global_step:05d}.safetensors" + ) + save_predictor_weights( + model, + weights_path, + floating_dtype=torch.bfloat16, + ) + if args.save_optimizer_state: + save_training_state( + output_dir / "training_latest.pt", + { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "global_step": global_step, + "micro_step": micro_step, + "epoch": epoch, + "args": vars(args), + "predictor_config": model.config_dict, + "weights_path": str(weights_path), + }, + ) + prune_snapshots(output_dir, args.keep_snapshots) + dist.barrier() + except BaseException as error: + if rank == 0 and swanlab_run is not None: + swanlab.finish(error=str(error)) + raise + else: + if rank == 0 and swanlab_run is not None: + swanlab.finish() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/train_predictor_v4_grouped_teacher_forcing.py b/train_predictor_v4_grouped_teacher_forcing.py new file mode 100644 index 0000000000000000000000000000000000000000..c7d6a725003eb2219d7583b45d0a645023db4a3f --- /dev/null +++ b/train_predictor_v4_grouped_teacher_forcing.py @@ -0,0 +1,825 @@ +#!/usr/bin/env python3 +"""Eight-GPU grouped teacher forcing using only the offline v4 dataset.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import random +import time +from pathlib import Path + +import swanlab +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW + +from hyvideo.commons.parallel_states import initialize_parallel_state +from models.predictor_v2 import PredictorV2Config +from models.predictor_v4 import HYWorldPlayPredictorV4 +from predictor_training.checkpoint import ( + TRAINABLE_PREFIXES, + load_concat_fusion_initialization, + load_concat_residual_gate_initialization, + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.timestep_rollout_dataset import ( + load_case_ids, + load_timestep_rollout_record, + load_timestep_rollout_records, + move_timestep_rollout_record, +) +from train_predictor_v4 import ( + configure_model_precision, + initialize_runtime_state, + make_scheduler, + prune_snapshots, + validate_optimizer_precision, +) + + +DEFAULT_DATA_ROOT = Path("/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4") +DEFAULT_CASES = Path("configs/rollout_overfit8_seed20260721_cases.jsonl") +DEFAULT_TEACHER = Path( + "/mnt/s3files/s3-us-west2-default/zoubin/cz/checkpoints/hy_worldplay/" + "huggingface/hub/models--tencent--HY-WorldPlay/snapshots/" + "f4c29235647707b571479a69b569e4166f9f5bf8/ar_distilled_action_model/" + "diffusion_pytorch_model.safetensors" +) +DEFAULT_INIT_STATE = Path( + "checkpoints/" + "predictor_v4_prefeature_blocks1-52_fp32master_8gpu_" + "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/" + "training_latest.pt" +) +SOURCE_BLOCKS = (1, 52) +EXPECTED_DENOISE_TIMESTEPS = (1000.0, 937.5, 833.3333129882812, 625.0) +EXPECTED_DENOISE_SIGMAS = (1.0, 0.9375, 0.8333333134651184, 0.625, 0.0) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--training_mode", + choices=( + "grouped_teacher_forcing_offline_prefeature", + "detached_timestep_rollout_offline_ffff_targets", + ), + default="grouped_teacher_forcing_offline_prefeature", + ) + parser.add_argument("--data_root", type=Path, default=DEFAULT_DATA_ROOT) + parser.add_argument("--cases_path", type=Path, default=DEFAULT_CASES) + parser.add_argument("--teacher_checkpoint", type=Path, default=DEFAULT_TEACHER) + parser.add_argument("--init_training_state", type=Path, default=DEFAULT_INIT_STATE) + parser.add_argument("--output_dir", type=Path, required=True) + parser.add_argument("--max_steps", type=int, default=1000) + parser.add_argument("--seed", type=int, default=20260721) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument( + "--predictor_fusion_type", + choices=( + "concat", + "projection_add", + "token_gate_prev", + "concat_residual_gate_prev", + ), + default="concat", + ) + parser.add_argument("--gate_hidden_dim", type=int, default=256) + parser.add_argument("--gate_initial_probability", type=float, default=0.8) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--scheduler_warmup_steps", type=int, default=8) + parser.add_argument("--block_warmup_steps", type=int, default=8) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument( + "--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--fp32_trainable_params", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument("--save_every", type=int, default=100) + parser.add_argument("--keep_snapshots", type=int, default=3) + parser.add_argument("--log_every", type=int, default=1) + parser.add_argument("--resume", type=Path, default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v4") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument("--swanlab_experiment", required=True) + parser.add_argument( + "--swanlab_mode", + choices=("cloud", "offline", "local", "disabled"), + default="cloud", + ) + return parser.parse_args() + + +def distributed_average(values: torch.Tensor) -> torch.Tensor: + dist.all_reduce(values, op=dist.ReduceOp.SUM) + values /= dist.get_world_size() + return values + + +def load_stage1_initialization( + model: HYWorldPlayPredictorV4, + path: Path, +) -> None: + """Load the Stage-1 FP32 master model while resetting optimizer/LR state.""" + state = torch.load(path, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != SOURCE_BLOCKS: + raise ValueError("Stage-1 initialization source blocks do not match") + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + state.get("args", {}).get("predictor_fusion_type", "concat"), + ) + target_fusion_type = model.predictor_config.predictor_fusion_type + if ( + saved_fusion_type == "concat" + and target_fusion_type == "concat_residual_gate_prev" + ): + report = load_concat_residual_gate_initialization(model, state["model"]) + print( + json.dumps( + { + "fusion_initialization": "concat_residual_gate", + "fresh_residual_key_count": len( + report["fresh_residual_keys"] + ), + }, + sort_keys=True, + ), + flush=True, + ) + return + if saved_fusion_type == "concat" and target_fusion_type != "concat": + report = load_concat_fusion_initialization(model, state["model"]) + print( + json.dumps( + { + "fusion_initialization": "concat_conversion", + "discarded_concat_key_count": len( + report["discarded_concat_keys"] + ), + "fresh_fusion_key_count": len(report["fresh_fusion_keys"]), + }, + sort_keys=True, + ), + flush=True, + ) + return + if saved_fusion_type != target_fusion_type: + raise ValueError( + "Stage-1 initialization fusion type does not match: " + f"saved={saved_fusion_type}, requested={target_fusion_type}" + ) + result = model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + key for key in result.missing_keys if key.startswith(TRAINABLE_PREFIXES) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Stage-1 initialization mismatch; " + f"unexpected={result.unexpected_keys}, missing={missing_trainable}" + ) + + +def assert_offline_timesteps(item: dict, device: torch.device) -> None: + expected = torch.tensor( + EXPECTED_DENOISE_TIMESTEPS, device=device, dtype=torch.float32 + ) + values = [] + for step_id in range(4): + value = item["timesteps"][step_id] + if value.dtype != torch.float32 or value.numel() != 1: + raise TypeError(f"Offline timestep {step_id} is not one FP32 scalar") + values.append(value.reshape(1)) + actual = torch.cat(values) + if not torch.equal(actual, expected): + raise RuntimeError( + "Offline denoise timesteps do not match the fixed four-step grid: " + f"actual={actual.detach().cpu().tolist()}" + ) + + +@torch.no_grad() +def rebuild_offline_context_kv( + model: HYWorldPlayPredictorV4, + item: dict, +) -> dict[int, dict[str, torch.Tensor]]: + """Run only v4's frozen projectors on saved BF16 Context prefeatures.""" + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + rebuilt = model.context_projector(dict(item["context_prefeature"])) + result = {} + for block_id in SOURCE_BLOCKS: + result[block_id] = { + **dict(item["text_kv"][block_id]), + **rebuilt[block_id], + "context_valid_mask": item["context_prefeature"][block_id][ + "context_valid_mask" + ], + } + return result + + +def model_inputs( + *, + pair: dict, + item: dict, + context_kv: dict[int, dict[str, torch.Tensor]], +) -> dict: + return { + "target_model_input": pair["model_input"], + "anchor_hidden": pair["anchor_hidden"], + "target_frame_condition": pair["frame_condition"], + "target_viewmats": item["target_viewmats"], + "target_Ks": item["target_Ks"], + "rope_temporal_size": item["rope_temporal_size"], + "start_rope_start_idx": item["start_rope_start_idx"], + "context_kv": context_kv, + "previous_chunk_hidden": pair["previous_chunk_hidden"], + } + + +def empty_drift_metrics() -> dict[str, float]: + return { + "latent_drift_step2": 0.0, + "latent_drift_step3": 0.0, + "anchor_drift_step2": 0.0, + "anchor_drift_step3": 0.0, + } + + +def detached_scheduler_step( + sample: torch.Tensor, + velocity: torch.Tensor, + *, + step_id: int, +) -> torch.Tensor: + """Apply the exact four-step Euler update and truncate its graph.""" + if step_id not in (1, 2): + raise ValueError(f"Detached rollout only advances after steps 1/2, got {step_id}") + sigmas = torch.tensor( + EXPECTED_DENOISE_SIGMAS, + device=sample.device, + dtype=torch.float32, + ) + with torch.no_grad(): + next_sample = sample.detach().float() + velocity.detach().float() * ( + sigmas[step_id + 1] - sigmas[step_id] + ) + if next_sample.requires_grad: + raise RuntimeError("Detached scheduler output unexpectedly requires gradients") + return next_sample + + +def assert_offline_ffff_transitions(item: dict) -> None: + """Account for BF16 storage while checking the saved nominal trajectory.""" + for step_id in (1, 2): + pair = item["teacher_forcing"][step_id] + expected = detached_scheduler_step( + pair["model_input"][:, :32], + pair["target_velocity"], + step_id=step_id, + ) + target = item["teacher_forcing"][step_id + 1]["model_input"][:, :32] + error = (expected - target.float()).abs() + if float(error.max()) > 0.02 or float(error.mean()) > 0.002: + raise RuntimeError( + f"Offline FFFF transition {step_id}->{step_id + 1} does not " + f"match the fixed scheduler: max={float(error.max())}, " + f"mean={float(error.mean())}" + ) + + +def train_grouped_record( + *, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + optimizer: torch.optim.Optimizer, + item: dict, + args: argparse.Namespace, + device: torch.device, +) -> dict[str, float]: + started = time.perf_counter() + assert_offline_timesteps(item, device) + if tuple(sorted(item["teacher_forcing"])) != (1, 2, 3): + raise ValueError("Grouped record does not contain all three supervision pairs") + context_kv = rebuild_offline_context_kv(predictor_model, item) + + optimizer.zero_grad(set_to_none=True) + losses: list[torch.Tensor] = [] + hidden_losses: list[torch.Tensor] = [] + velocity_losses: list[torch.Tensor] = [] + for step_id in (1, 2, 3): + pair = item["teacher_forcing"][step_id] + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = predictor( + **model_inputs(pair=pair, item=item, context_kv=context_kv) + ) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), pair["target_hidden"].float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), pair["target_velocity"].float() + ) + step_loss = ( + args.hidden_weight * hidden_loss + + args.velocity_weight * velocity_loss + ) + losses.append(step_loss) + hidden_losses.append(hidden_loss.detach()) + velocity_losses.append(velocity_loss.detach()) + del output, hidden_loss, velocity_loss + + objective_loss = torch.stack(losses).mean() + objective_loss.backward() + if any(parameter.grad is not None for parameter in predictor_model.context_projector.parameters()): + raise RuntimeError("Frozen Context projector unexpectedly received gradients") + grad_norm = torch.nn.utils.clip_grad_norm_( + predictor_model.parameters(), args.grad_clip + ) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite grad norm: {float(grad_norm)}") + optimizer.step() + result = { + "loss": float(objective_loss.detach()), + "hidden_mse_step1": float(hidden_losses[0]), + "hidden_mse_step2": float(hidden_losses[1]), + "hidden_mse_step3": float(hidden_losses[2]), + "velocity_mse_step1": float(velocity_losses[0]), + "velocity_mse_step2": float(velocity_losses[1]), + "velocity_mse_step3": float(velocity_losses[2]), + "grad_norm": float(grad_norm), + "record_time_s": time.perf_counter() - started, + **empty_drift_metrics(), + } + del context_kv, losses, hidden_losses, velocity_losses, objective_loss + return result + + +def train_detached_rollout_record( + *, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + optimizer: torch.optim.Optimizer, + item: dict, + args: argparse.Namespace, + device: torch.device, +) -> dict[str, float]: + """Roll out Predictor inputs while supervising against fixed offline FFFF outputs.""" + started = time.perf_counter() + assert_offline_timesteps(item, device) + if tuple(sorted(item["teacher_forcing"])) != (1, 2, 3): + raise ValueError("Detached rollout record lacks one of steps 1/2/3") + assert_offline_ffff_transitions(item) + context_kv = rebuild_offline_context_kv(predictor_model, item) + + optimizer.zero_grad(set_to_none=True) + losses: list[torch.Tensor] = [] + hidden_losses: list[torch.Tensor] = [] + velocity_losses: list[torch.Tensor] = [] + latent_drifts: dict[int, torch.Tensor] = {} + anchor_drifts: dict[int, torch.Tensor] = {} + + first_pair = item["teacher_forcing"][1] + rollout_latent = first_pair["model_input"][:, :32] + fixed_condition = first_pair["model_input"][:, 32:] + rollout_anchor = first_pair["anchor_hidden"] + for step_id in (1, 2, 3): + target_pair = item["teacher_forcing"][step_id] + if step_id > 1: + if rollout_latent.requires_grad or rollout_anchor.requires_grad: + raise RuntimeError( + f"Step {step_id} received a gradient-connected rollout state" + ) + latent_drifts[step_id] = F.mse_loss( + rollout_latent.float(), + target_pair["model_input"][:, :32].float(), + ).detach() + anchor_drifts[step_id] = F.mse_loss( + rollout_anchor.float(), target_pair["anchor_hidden"].float() + ).detach() + model_input = ( + target_pair["model_input"] + if step_id == 1 + else torch.cat([rollout_latent, fixed_condition], dim=1) + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = predictor( + **model_inputs( + pair={ + **target_pair, + "model_input": model_input, + "anchor_hidden": rollout_anchor, + }, + item=item, + context_kv=context_kv, + ) + ) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), target_pair["target_hidden"].float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), target_pair["target_velocity"].float() + ) + step_loss = ( + args.hidden_weight * hidden_loss + + args.velocity_weight * velocity_loss + ) + losses.append(step_loss) + hidden_losses.append(hidden_loss.detach()) + velocity_losses.append(velocity_loss.detach()) + if step_id < 3: + rollout_anchor = output["pred_hidden"].detach() + rollout_latent = detached_scheduler_step( + rollout_latent, output["pred_velocity"], step_id=step_id + ) + del output, hidden_loss, velocity_loss, model_input + + objective_loss = torch.stack(losses).mean() + objective_loss.backward() + if any(parameter.grad is not None for parameter in predictor_model.context_projector.parameters()): + raise RuntimeError("Frozen Context projector unexpectedly received gradients") + grad_norm = torch.nn.utils.clip_grad_norm_( + predictor_model.parameters(), args.grad_clip + ) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite grad norm: {float(grad_norm)}") + optimizer.step() + result = { + "loss": float(objective_loss.detach()), + "hidden_mse_step1": float(hidden_losses[0]), + "hidden_mse_step2": float(hidden_losses[1]), + "hidden_mse_step3": float(hidden_losses[2]), + "velocity_mse_step1": float(velocity_losses[0]), + "velocity_mse_step2": float(velocity_losses[1]), + "velocity_mse_step3": float(velocity_losses[2]), + "latent_drift_step2": float(latent_drifts[2]), + "latent_drift_step3": float(latent_drifts[3]), + "anchor_drift_step2": float(anchor_drifts[2]), + "anchor_drift_step3": float(anchor_drifts[3]), + "grad_norm": float(grad_norm), + "record_time_s": time.perf_counter() - started, + } + del context_kv, losses, hidden_losses, velocity_losses + del latent_drifts, anchor_drifts, objective_loss, rollout_latent, rollout_anchor + return result + + +def main() -> None: + args = parse_args() + if args.max_steps <= 0: + raise ValueError("max_steps must be positive") + if not 0.0 < args.gate_initial_probability < 1.0: + raise ValueError("gate_initial_probability must be in (0, 1)") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + raise RuntimeError("Launch grouped teacher forcing with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + if world_size != 8: + raise ValueError(f"The overfit8 experiment requires 8 ranks, got {world_size}") + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="nccl", device_id=device) + initialize_parallel_state(sp=1) + initialize_runtime_state() + random.seed(args.seed + rank) + torch.manual_seed(args.seed + rank) + torch.cuda.manual_seed_all(args.seed + rank) + + data_root = args.data_root.resolve() + cases_path = args.cases_path.resolve() + manifest_path = data_root / "train_eval_manifest.jsonl" + case_ids = load_case_ids(cases_path) + if len(case_ids) != world_size: + raise ValueError("Exactly one selected case per rank is required") + case_id = case_ids[rank] + record_root, records = load_timestep_rollout_records( + manifest_path, case_id=case_id + ) + required = [ + data_root, + cases_path, + manifest_path, + args.teacher_checkpoint.resolve(), + args.init_training_state.resolve(), + ] + if args.resume is not None: + required.append(args.resume.resolve()) + for path in required: + if not path.exists(): + raise FileNotFoundError(path) + + output_dir = args.output_dir.resolve() + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + dist.barrier() + + model = HYWorldPlayPredictorV4( + PredictorV2Config( + source_block_ids=SOURCE_BLOCKS, + predictor_fusion_type=args.predictor_fusion_type, + gate_hidden_dim=args.gate_hidden_dim, + gate_initial_probability=args.gate_initial_probability, + ) + ) + # This selectively reads the two source blocks and fixed input/output modules + # from safetensors. It does not instantiate the Full Teacher model. + model.load_teacher_initialization(args.teacher_checkpoint.resolve()) + configure_model_precision( + model, + device=device, + activation_dtype=torch.bfloat16, + fp32_trainable_params=args.fp32_trainable_params, + ) + if args.resume is None: + load_stage1_initialization(model, args.init_training_state.resolve()) + model.enable_gradient_checkpointing(args.gradient_checkpointing) + model.train() + if any(parameter.requires_grad for parameter in model.context_projector.parameters()): + raise RuntimeError("Context projector must remain frozen") + + fusion_parameters = list(model.img_fusion.parameters()) + list( + model.residual_out.parameters() + ) + block_parameters = [ + parameter + for parameter in model.double_blocks.parameters() + if parameter.requires_grad + ] + optimizer = AdamW( + [ + {"params": fusion_parameters, "lr": args.fusion_lr}, + {"params": block_parameters, "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), + weight_decay=args.weight_decay, + ) + scheduler = make_scheduler(optimizer, args) + global_step = 0 + epoch = 0 + batch_in_epoch = 0 + if args.resume is not None: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + saved_args = state.get("args", {}) + if saved_args.get("training_mode") != args.training_mode: + raise ValueError("Resume checkpoint training mode does not match") + if bool(saved_args.get("fp32_trainable_params")) != args.fp32_trainable_params: + raise ValueError("Resume precision mode does not match") + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + saved_args.get("predictor_fusion_type", "concat"), + ) + if saved_fusion_type != args.predictor_fusion_type: + raise ValueError("Resume fusion type does not match") + result = model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + key for key in result.missing_keys if key.startswith(TRAINABLE_PREFIXES) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Resume model mismatch; " + f"unexpected={result.unexpected_keys}, missing={missing_trainable}" + ) + optimizer.load_state_dict(state["optimizer"]) + scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + epoch = int(state.get("epoch", 0)) + batch_in_epoch = int(state.get("batch_in_epoch", 0)) + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=global_step > 0, + ) + predictor = DDP( + model, + device_ids=[local_rank], + output_device=local_rank, + broadcast_buffers=False, + gradient_as_bucket_view=True, + find_unused_parameters=False, + ) + + detached_rollout = ( + args.training_mode == "detached_timestep_rollout_offline_ffff_targets" + ) + run_args = vars(args) + run_config = { + **run_args, + "data_root": str(data_root), + "cases_path": str(cases_path), + "output_dir": str(output_dir), + "world_size": world_size, + "case_ids": case_ids, + "records_per_rank": len(records), + "groups_per_optimizer_step_per_rank": 1, + "pairs_per_group": 3, + "source_blocks": list(SOURCE_BLOCKS), + "initialization": "stage1_step2000_model_only_optimizer_reset", + "runtime_full_teacher_loaded": False, + "teacher_checkpoint_usage": "selective_predictor_initialization_only", + "rollout_sidecar_loaded": False, + "context_source": "offline_bf16_prefeature_plus_offline_text_kv", + "context_projection": "frozen_v4_projector_once_per_group", + "targets": "offline_full_teacher_steps_1_2_3", + "input_state_source": ( + "step1_offline_then_predictor_detached_rollout" + if detached_rollout else "offline_full_teacher_each_step" + ), + "anchor_source": ( + "step1_offline_then_predictor_hidden_detached" + if detached_rollout else "offline_full_teacher_each_step" + ), + "timestep_graph": ( + "P1|P2|P3_detached_state_rollout_one_mean_loss" + if detached_rollout + else "three_independent_offline_pairs_one_mean_loss" + ), + "offline_target_state_matches_rollout_input": not detached_rollout, + "offline_denoise_timesteps": list(EXPECTED_DENOISE_TIMESTEPS), + "offline_denoise_sigmas": list(EXPECTED_DENOISE_SIGMAS), + "offline_denoise_timestep_dtype": "float32", + "activation_dtype": "bfloat16", + "trainable_parameter_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + } + if rank == 0: + (output_dir / "train_config.json").write_text( + json.dumps(run_config, indent=2, ensure_ascii=False, default=str, sort_keys=True) + + "\n", + encoding="utf-8", + ) + print(json.dumps(run_config, indent=2, ensure_ascii=False, default=str), flush=True) + + swanlab_run = None + if rank == 0: + if args.swanlab_mode == "cloud": + swanlab.login(api_key=os.environ.get("SWANLAB_API_KEY"), save=False) + swanlab_run = swanlab.init( + project=args.swanlab_project, + workspace=args.swanlab_workspace, + experiment_name=args.swanlab_experiment, + description=( + "v4 detached timestep rollout with fixed offline FFFF targets; " + "offline prefeature path; no runtime Full Teacher" + if detached_rollout + else ( + "v4 grouped teacher forcing on overfit8; fully offline Stage-1 " + "prefeature path; no runtime Full Teacher" + ) + ), + tags=[ + "predictor-v4", + "detached-rollout" if detached_rollout else "grouped-teacher", + "offline-prefeature", + "overfit8", + "8gpu", + f"fusion-{args.predictor_fusion_type}"[:20], + ], + config=run_config, + logdir=str(output_dir / "swanlab"), + mode=args.swanlab_mode, + ) + + log_path = output_dir / "train_log.jsonl" + metric_names = ( + "loss", + "hidden_mse_step1", + "hidden_mse_step2", + "hidden_mse_step3", + "velocity_mse_step1", + "velocity_mse_step2", + "velocity_mse_step3", + "latent_drift_step2", + "latent_drift_step3", + "anchor_drift_step2", + "anchor_drift_step3", + "record_time_s", + ) + try: + resume_epoch = epoch + resume_batch = batch_in_epoch + while global_step < args.max_steps: + order = list(range(len(records))) + random.Random(args.seed + rank + epoch * 100003).shuffle(order) + for loader_index, record_index in enumerate(order): + if epoch == resume_epoch and loader_index < resume_batch: + continue + if global_step >= args.max_steps: + break + cpu_item = load_timestep_rollout_record( + record_root, + records[record_index], + replay_step=None, + grouped_teacher_forcing=True, + source_block_ids=SOURCE_BLOCKS, + ) + item = move_timestep_rollout_record(cpu_item, device) + torch.cuda.reset_peak_memory_stats(device) + train_function = ( + train_detached_rollout_record + if detached_rollout else train_grouped_record + ) + metrics = train_function( + predictor=predictor, predictor_model=model, + optimizer=optimizer, item=item, args=args, device=device, + ) + if global_step == 0: + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + scheduler.step() + global_step += 1 + batch_in_epoch = loader_index + 1 + + values = torch.tensor( + [metrics[name] for name in metric_names], + device=device, + dtype=torch.float64, + ) + averaged = distributed_average(values) + grad_max = torch.tensor(metrics["grad_norm"], device=device) + dist.all_reduce(grad_max, op=dist.ReduceOp.MAX) + if rank == 0 and ( + global_step % args.log_every == 0 or global_step == 1 + ): + record_log = { + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + **{ + name: float(averaged[index]) + for index, name in enumerate(metric_names) + }, + "grad_norm_max": float(grad_max), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated(device) / 2**30, + } + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record_log, sort_keys=True) + "\n") + swanlab.log(record_log, step=global_step) + print(json.dumps(record_log, sort_keys=True), flush=True) + + should_save = ( + global_step % args.save_every == 0 + or global_step == args.max_steps + ) + if should_save: + dist.barrier() + if rank == 0: + weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" + save_predictor_weights( + model, weights_path, floating_dtype=torch.bfloat16 + ) + save_training_state( + output_dir / "training_latest.pt", + { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + "args": run_args, + "predictor_config": model.config_dict, + "weights_path": str(weights_path), + }, + ) + prune_snapshots(output_dir, args.keep_snapshots) + dist.barrier() + del cpu_item, item, metrics, values, averaged, grad_max + gc.collect() + epoch += 1 + batch_in_epoch = 0 + resume_batch = 0 + except BaseException as error: + if rank == 0 and swanlab_run is not None: + try: + swanlab.finish(error=str(error)) + except Exception: + pass + raise + else: + if rank == 0 and swanlab_run is not None: + swanlab.finish() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/train_predictor_v4_random_exit_dmd.py b/train_predictor_v4_random_exit_dmd.py new file mode 100644 index 0000000000000000000000000000000000000000..ef335818fb6d1612a090b72006690d45c90b649a --- /dev/null +++ b/train_predictor_v4_random_exit_dmd.py @@ -0,0 +1,1242 @@ +#!/usr/bin/env python3 +"""Eight-GPU Predictor-only random-exit DMD for HY-WorldPlay Predictor-v4. + +The four-step distilled autoregressive model constructs F/P rollout states. +An undistilled 50-step bidirectional model is the frozen real score and an +independent copy of the same model is trained as the fake score. Only the +Predictor call at the sampled P1/P2/P3 exit retains a graph. +""" + +from __future__ import annotations + +import argparse +import copy +import gc +import json +import math +import os +import random +import shutil +import time +from functools import partial +from pathlib import Path +from typing import Any + +import numpy as np +import swanlab +import torch +import torch.distributed as dist +import torch.nn.functional as F +from PIL import Image +from safetensors.torch import load_file +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + apply_activation_checkpointing, + checkpoint_wrapper, +) +from torch.distributed.fsdp import ( + CPUOffload, + FullyShardedDataParallel as FSDP, + MixedPrecision, + ShardingStrategy, +) +from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from hyvideo.commons import auto_offload_model +from hyvideo.commons.infer_state import initialize_infer_state +from hyvideo.commons.parallel_states import initialize_parallel_state +from hyvideo.models.transformers.worldplay_1_5_transformer import ( + HunyuanVideo_1_5_DiffusionTransformer, +) +from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline +from hyvideo.utils.retrieval_context import ( + generate_points_in_sphere, + select_aligned_memory_frames, +) +from models.predictor_v2 import PredictorV2Config +from models.predictor_v4 import HYWorldPlayPredictorV4 +from predictor_training.checkpoint import ( + TRAINABLE_PREFIXES, + save_predictor_weights, + trainable_state_dict, +) +from predictor_training.rollout_utils import configure_exact_teacher, pose_to_input +from predictor_training.trajectory_dataset import ( + PredictorV4TrajectoryDataset, + trajectory_collate, +) +from train_predictor_v4_rollout import ( + configure_model_precision, + differentiable_euler_step, + history_prefill, + load_image_condition, + load_initial_training_state, + teacher_denoise, +) + + +CHUNK_FRAMES = 4 +NUM_CHUNKS = 8 +NUM_STEPS = 4 +TRAIN_EXITS = (1, 2, 3) +LATENT_CHANNELS = 32 +LATENT_HEIGHT = 30 +LATENT_WIDTH = 52 + +LOCAL_MODEL_ROOT = Path( + "/mnt/local_nvme/cz/checkpoints/hy_worldplay/huggingface/hub" +) +DEFAULT_BASE = LOCAL_MODEL_ROOT / ( + "models--tencent--HunyuanVideo-1.5/snapshots/" + "9b49404b3f5df2a8f0b31df27a0c7ab872e7b038" +) +REPAIRED_BASE = Path("outputs/model_cache_repaired/HunyuanVideo-1.5-9b49404-v2") +if REPAIRED_BASE.is_dir(): + DEFAULT_BASE = REPAIRED_BASE +DEFAULT_WORLDPLAY = LOCAL_MODEL_ROOT / ( + "models--tencent--HY-WorldPlay/snapshots/" + "f4c29235647707b571479a69b569e4166f9f5bf8" +) +DEFAULT_ROLLOUT_TEACHER = ( + DEFAULT_WORLDPLAY + / "ar_distilled_action_model/diffusion_pytorch_model.safetensors" +) +DEFAULT_SCORE_TEACHER = ( + DEFAULT_WORLDPLAY / "bidirectional_model/diffusion_pytorch_model.safetensors" +) +DEFAULT_DATA = Path("datasets/hyworldplay_predictor_prefeature_v4") +DEFAULT_STAGE1 = Path( + "checkpoints/" + "predictor_v4_prefeature_blocks1-52_fp32master_8gpu_" + "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/" + "training_latest.pt" +) +DEFAULT_OUTPUT = Path( + "checkpoints/" + "predictor_v4_random-exit-dmd-predictor-only-stage1step2000_" + "blocks1-52_fp32master_8gpu" +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base_model", type=Path, default=DEFAULT_BASE) + parser.add_argument( + "--rollout_teacher_checkpoint", type=Path, default=DEFAULT_ROLLOUT_TEACHER + ) + parser.add_argument( + "--score_teacher_checkpoint", type=Path, default=DEFAULT_SCORE_TEACHER + ) + parser.add_argument("--data_root", type=Path, default=DEFAULT_DATA) + parser.add_argument("--cases_path", type=Path, default=None) + parser.add_argument("--stage1_training_state", type=Path, default=DEFAULT_STAGE1) + parser.add_argument("--output_dir", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--max_steps", type=int, default=2000) + parser.add_argument("--max_trajectories", type=int, default=None) + parser.add_argument("--seed", type=int, default=20260818) + parser.add_argument("--latent_seed", type=int, default=0) + parser.add_argument("--forced_exit_step", type=int, choices=TRAIN_EXITS, default=None) + parser.add_argument("--critic_updates_per_student", type=int, default=5) + parser.add_argument("--fusion_lr", type=float, default=5e-6) + parser.add_argument("--blocks_lr", type=float, default=5e-7) + parser.add_argument("--fake_score_lr", type=float, default=2e-7) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--warmup_steps", type=int, default=100) + parser.add_argument("--predictor_grad_clip", type=float, default=1.0) + parser.add_argument("--fake_score_grad_clip", type=float, default=10.0) + parser.add_argument("--real_guidance_scale", type=float, default=6.0) + parser.add_argument("--score_timestep_min", type=float, default=0.02) + parser.add_argument("--score_timestep_max", type=float, default=0.98) + parser.add_argument("--flow_shift", type=float, default=5.0) + parser.add_argument("--save_every", type=int, default=200) + parser.add_argument("--keep_checkpoints", type=int, default=2) + parser.add_argument("--no_save", action="store_true") + parser.add_argument("--log_every", type=int, default=1) + parser.add_argument("--resume", type=Path, default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v4") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument( + "--swanlab_experiment", + default="predictor-v4-random-exit-dmd-predictor-only-stage1step2000-8gpu", + ) + parser.add_argument( + "--swanlab_mode", + choices=("cloud", "offline", "local", "disabled"), + default="cloud", + ) + parser.add_argument("--validate_config_only", action="store_true") + return parser.parse_args() + + +def initialize_runtime_state() -> None: + class Args: + sage_blocks_range = "0-53" + use_sageattn = False + enable_torch_compile = False + use_fp8_gemm = False + quant_type = "fp8-per-block" + include_patterns = "double_blocks" + use_vae_parallel = False + + initialize_infer_state(Args()) + + +def _unwrap(module: torch.nn.Module) -> torch.nn.Module: + while hasattr(module, "module"): + module = module.module + return module + + +def _json_dump(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True, default=str) + + "\n", + encoding="utf-8", + ) + os.replace(temporary, path) + + +def _resolve_hf_pointer(path: Path) -> Path: + """Resolve this workspace's text materialization of a HF cache symlink.""" + path = path.resolve() + if path.stat().st_size > 4096: + return path + candidate = (path.parent / path.read_text(encoding="utf-8").strip()).resolve() + if not candidate.is_file(): + raise FileNotFoundError(f"Broken Hugging Face pointer {path} -> {candidate}") + return candidate + + +def _cosine_factor(step: int, warmup: int, total: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = min(1.0, float(step - warmup) / max(1, total - warmup)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + +def _make_scheduler( + optimizer: torch.optim.Optimizer, *, warmup: int, total: int +) -> LambdaLR: + return LambdaLR(optimizer, lambda step: _cosine_factor(step, warmup, total)) + + +def _fsdp_wrap(module: torch.nn.Module, device: torch.device) -> FSDP: + policy = partial(size_based_auto_wrap_policy, min_num_params=int(5e7)) + mixed = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False, + ) + return FSDP( + module, + auto_wrap_policy=policy, + sharding_strategy=ShardingStrategy.FULL_SHARD, + mixed_precision=mixed, + device_id=device, + limit_all_gathers=True, + use_orig_params=True, + cpu_offload=CPUOffload(offload_params=False), + sync_module_states=False, + ) + + +def _enable_activation_checkpointing(module: torch.nn.Module) -> None: + block_ids = {id(block) for block in module.double_blocks} + wrapper = partial( + checkpoint_wrapper, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + preserve_rng_state=False, + ) + apply_activation_checkpointing( + module, + checkpoint_wrapper_fn=wrapper, + check_fn=lambda child: id(child) in block_ids, + ) + + +def _load_score_model( + *, + base_model: Path, + checkpoint: Path, + trainable: bool, + activation_checkpointing: bool, +) -> HunyuanVideo_1_5_DiffusionTransformer: + model = HunyuanVideo_1_5_DiffusionTransformer.from_pretrained( + base_model / "transformer/480p_i2v", + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + model.add_action_parameters() + checkpoint = _resolve_hf_pointer(checkpoint) + state = load_file(str(checkpoint), device="cpu") + model.load_state_dict(state, strict=True) + del state + configure_exact_teacher(model) + model.to(dtype=torch.bfloat16) + model.requires_grad_(trainable) + if activation_checkpointing: + _enable_activation_checkpointing(model) + model.train(trainable) + return model + + +def _predictor_groups( + model: torch.nn.Module, +) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter]]: + fusion: list[torch.nn.Parameter] = [] + blocks: list[torch.nn.Parameter] = [] + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + continue + if name.startswith("double_blocks."): + blocks.append(parameter) + else: + fusion.append(parameter) + if not fusion or not blocks: + raise RuntimeError( + f"Invalid Predictor groups: fusion={len(fusion)}, blocks={len(blocks)}" + ) + return fusion, blocks + + +def _load_exact_stage1( + model: torch.nn.Module, path: Path, source_blocks: tuple[int, int] +) -> None: + state = torch.load(path, map_location="cpu", weights_only=False) + actual_step = int(state.get("global_step", -1)) + if actual_step != 2000: + raise ValueError(f"Stage-1 initialization must be step 2000, got {actual_step}") + del state + load_initial_training_state(model, path, source_blocks) + + +@torch.no_grad() +def _prepare_conditions( + pipe: HunyuanVideo_1_5_Pipeline, + record: dict[str, Any], + device: torch.device, +) -> dict[str, Any]: + """Prepare shared rollout inputs and positive/negative score conditions.""" + image = Image.open(record["image_path"]).convert("RGB") + image_np = np.asarray(image) + pipe._guidance_scale = 6.0 + pipe._guidance_rescale = 0.0 + pipe._clip_skip = None + with auto_offload_model(pipe.text_encoder, device, enabled=True): + positive, negative, positive_mask, negative_mask = pipe.encode_prompt( + record["caption"], + device, + 1, + True, + "", + clip_skip=None, + data_type="video", + ) + with auto_offload_model(pipe.byt5_model, device, enabled=True): + combined_extra = pipe._prepare_byt5_embeddings(record["caption"], device) + with auto_offload_model(pipe.vision_encoder, device, enabled=True): + combined_vision = pipe._prepare_vision_states( + image_np, pipe.ideal_resolution, torch.empty(1, device=device), device + ) + if combined_vision is None or combined_vision.shape[0] != 2: + raise RuntimeError("Expected duplicated I2V vision states for CFG") + + scheduler = pipe._create_scheduler(pipe.config.flow_shift) + scheduler.set_timesteps(NUM_STEPS, device=device) + generator = torch.Generator(device=device).manual_seed(int(record["seed"])) + latents = pipe.prepare_latents( + 1, + LATENT_CHANNELS, + LATENT_HEIGHT, + LATENT_WIDTH, + NUM_CHUNKS * CHUNK_FRAMES, + torch.bfloat16, + device, + generator, + ) + image_condition = load_image_condition(record["case_tensor_path"], device) + cond_latents = pipe._prepare_cond_latents( + "i2v", + image_condition, + latents, + pipe.get_task_mask("i2v", NUM_CHUNKS * CHUNK_FRAMES), + ).to(dtype=torch.bfloat16) + viewmats, Ks, action = pose_to_input( + record["pose"], NUM_CHUNKS * CHUNK_FRAMES + ) + + extra_negative = { + key: value[0:1] for key, value in combined_extra.items() + } + extra_positive = { + key: value[1:2] for key, value in combined_extra.items() + } + return { + "latents": latents, + "cond_latents": cond_latents, + "timesteps": scheduler.timesteps, + "sigmas": scheduler.sigmas.to(device=device), + "viewmats_cpu": viewmats.unsqueeze(0), + "viewmats": viewmats.unsqueeze(0).to(device=device, dtype=torch.bfloat16), + "Ks": Ks.unsqueeze(0).to(device=device, dtype=torch.bfloat16), + "action": action.unsqueeze(0).to(device=device), + "positive": { + "text": positive, + "mask": positive_mask, + "vision": combined_vision[1:2], + "extra": extra_positive, + }, + "negative": { + "text": negative, + "mask": negative_mask, + "vision": combined_vision[0:1], + "extra": extra_negative, + }, + } + + +@torch.no_grad() +def _text_prefill( + teacher: torch.nn.Module, + condition: dict[str, Any], + device: torch.device, +) -> list[dict[str, torch.Tensor | None]]: + module = _unwrap(teacher) + cache = [ + {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None} + for _ in module.double_blocks + ] + with torch.autocast("cuda", dtype=torch.bfloat16): + return teacher( + bi_inference=False, + ar_txt_inference=True, + ar_vision_inference=False, + timestep_txt=torch.tensor([0], device=device, dtype=torch.bfloat16), + text_states=condition["text"][0, None, ...], + encoder_attention_mask=condition["mask"][0, None, ...], + vision_states=condition["vision"][0, None, ...], + mask_type="i2v", + extra_kwargs={ + "byt5_text_states": condition["extra"]["byt5_text_states"][0, None, ...], + "byt5_text_mask": condition["extra"]["byt5_text_mask"][0, None, ...], + }, + kv_cache=cache, + cache_txt=True, + ) + + +def _select_context_indices( + viewmats_cpu: torch.Tensor, + chunk_id: int, + points_local: torch.Tensor, + device: torch.device, +) -> list[int]: + if chunk_id == 0: + return [] + current = chunk_id * CHUNK_FRAMES + selected = select_aligned_memory_frames( + viewmats_cpu[0].numpy(), + current, + memory_frames=20, + temporal_context_size=12, + pred_latent_size=CHUNK_FRAMES, + points_local=points_local, + device=device, + ) + return sorted(set(index for index in selected if index < current)) + + +def _sample_exit(device: torch.device, forced: int | None) -> int: + if dist.get_rank() == 0: + value = torch.tensor( + [forced if forced is not None else random.choice(TRAIN_EXITS)], + device=device, + dtype=torch.long, + ) + else: + value = torch.empty(1, device=device, dtype=torch.long) + dist.broadcast(value, src=0) + return int(value.item()) + + +def _x0_from_flow( + latent: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor +) -> torch.Tensor: + return latent.float() - sigma.to(device=latent.device, dtype=torch.float32) * velocity.float() + + +def _generate_random_exit_rollout( + *, + pipe: HunyuanVideo_1_5_Pipeline, + teacher: torch.nn.Module, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + condition_time_in: torch.nn.Module, + condition_action_in: torch.nn.Module, + record: dict[str, Any], + points_local: torch.Tensor, + exit_step: int, + retain_predictor_graph: bool, + device: torch.device, +) -> tuple[torch.Tensor, dict[str, Any]]: + static = _prepare_conditions(pipe, record, device) + latents = static["latents"] + cond_latents = static["cond_latents"] + timesteps = static["timesteps"] + sigmas = static["sigmas"] + viewmats = static["viewmats"] + Ks = static["Ks"] + action = static["action"] + cache = _text_prefill(teacher, static["positive"], device) + previous_hidden: dict[int, torch.Tensor] = {} + clean_chunks: list[torch.Tensor] = [] + + for chunk_id in range(NUM_CHUNKS): + start = chunk_id * CHUNK_FRAMES + end = start + CHUNK_FRAMES + selected = _select_context_indices( + static["viewmats_cpu"], chunk_id, points_local, device + ) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + cache = history_prefill( + teacher, + cache, + latents, + cond_latents, + selected, + viewmats, + Ks, + action, + timesteps.dtype, + ) + chunk_latent = latents[:, :, start:end].detach().clone() + chunk_cond = cond_latents[:, :, start:end] + chunk_viewmats = viewmats[:, start:end] + chunk_Ks = Ks[:, start:end] + chunk_action = action[:, start:end] + current_hidden: dict[int, torch.Tensor] = {} + + if chunk_id == 0: + clean = None + for step_id in range(exit_step + 1): + timestep_input = torch.full( + (CHUNK_FRAMES,), + timesteps[step_id], + device=device, + dtype=timesteps.dtype, + ) + model_input = torch.cat([chunk_latent, chunk_cond], dim=1) + velocity, hidden, _ = teacher_denoise( + teacher, + model_input, + timestep_input, + chunk_viewmats, + chunk_Ks, + chunk_action, + cache, + len(selected), + ) + if step_id > 0: + current_hidden[step_id] = hidden.detach() + clean = _x0_from_flow( + chunk_latent, velocity, sigmas[step_id] + ).to(dtype=torch.bfloat16) + if step_id < exit_step: + chunk_latent = differentiable_euler_step( + chunk_latent, velocity, sigmas, step_id + ).detach().to(dtype=torch.bfloat16) + if clean is None: + raise RuntimeError("Chunk 0 produced no x0 estimate") + previous_hidden = current_hidden + else: + timestep0 = torch.full( + (CHUNK_FRAMES,), timesteps[0], device=device, dtype=timesteps.dtype + ) + model_input = torch.cat([chunk_latent, chunk_cond], dim=1) + velocity0, anchor_hidden, _ = teacher_denoise( + teacher, + model_input, + timestep0, + chunk_viewmats, + chunk_Ks, + chunk_action, + cache, + len(selected), + ) + chunk_latent = differentiable_euler_step( + chunk_latent, velocity0, sigmas, 0 + ).detach().to(dtype=torch.bfloat16) + clean = None + for step_id in range(1, exit_step + 1): + if step_id not in previous_hidden: + raise RuntimeError( + f"Missing previous hidden for chunk={chunk_id}, step={step_id}" + ) + timestep_input = torch.full( + (CHUNK_FRAMES,), + timesteps[step_id], + device=device, + dtype=timesteps.dtype, + ) + model_input = torch.cat([chunk_latent, chunk_cond], dim=1) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + frame_condition = condition_time_in( + timestep_input.reshape(-1) + ) + condition_action_in(chunk_action.reshape(-1)) + frame_condition = frame_condition.reshape(1, CHUNK_FRAMES, -1) + keep_graph = retain_predictor_graph and step_id == exit_step + grad_context = torch.enable_grad() if keep_graph else torch.no_grad() + with grad_context, torch.autocast("cuda", dtype=torch.bfloat16): + output = predictor( + target_model_input=model_input, + anchor_hidden=anchor_hidden.detach(), + target_frame_condition=frame_condition, + target_viewmats=chunk_viewmats, + target_Ks=chunk_Ks, + rope_temporal_size=CHUNK_FRAMES + len(selected), + start_rope_start_idx=len(selected), + context_kv={ + block_id: cache[block_id] + for block_id in predictor_model.predictor_config.source_block_ids + }, + previous_chunk_hidden=previous_hidden[step_id], + ) + velocity = output["pred_velocity"] + clean = _x0_from_flow( + chunk_latent, velocity, sigmas[step_id] + ) + current_hidden[step_id] = output["pred_hidden"].detach() + anchor_hidden = output["pred_hidden"].detach() + if step_id < exit_step: + chunk_latent = differentiable_euler_step( + chunk_latent, velocity.detach(), sigmas, step_id + ).detach().to(dtype=torch.bfloat16) + if clean is None: + raise RuntimeError(f"Chunk {chunk_id} produced no x0 estimate") + previous_hidden = current_hidden + + clean_chunks.append(clean) + latents[:, :, start:end] = clean.detach().to(dtype=latents.dtype) + + generated = torch.cat(clean_chunks, dim=2) + score_conditions = { + "cond_latents": cond_latents, + "viewmats": viewmats, + "Ks": Ks, + "action": action, + "positive": static["positive"], + "negative": static["negative"], + } + return generated, score_conditions + + +def _score_forward( + model: torch.nn.Module, + *, + noisy: torch.Tensor, + timestep: torch.Tensor, + conditions: dict[str, Any], + branch: str, +) -> torch.Tensor: + condition = conditions[branch] + batch, _, frames = noisy.shape[:3] + model_input = torch.cat([noisy, conditions["cond_latents"]], dim=1) + timestep_frames = timestep[:, None].expand(batch, frames).reshape(-1) + action = conditions["action"].expand(batch, -1).reshape(-1) + viewmats = conditions["viewmats"].expand(batch, -1, -1, -1) + Ks = conditions["Ks"].expand(batch, -1, -1, -1) + with torch.autocast("cuda", dtype=torch.bfloat16): + output = model( + bi_inference=True, + ar_txt_inference=False, + ar_vision_inference=False, + hidden_states=model_input, + timestep=timestep_frames, + timestep_txt=timestep, + text_states=condition["text"], + text_states_2=None, + encoder_attention_mask=condition["mask"], + timestep_r=None, + vision_states=condition["vision"], + mask_type="i2v", + guidance=None, + return_dict=False, + extra_kwargs=condition["extra"], + viewmats=viewmats, + Ks=Ks, + action=action, + rope_temporal_size=frames, + start_rope_start_idx=0, + ) + return output[0] + + +def _sample_score_noise( + clean: torch.Tensor, + args: argparse.Namespace, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch = clean.shape[0] + raw = torch.rand(batch, device=clean.device, dtype=torch.float32) + raw = raw * (args.score_timestep_max - args.score_timestep_min) + args.score_timestep_min + sigma = args.flow_shift * raw / (1.0 + (args.flow_shift - 1.0) * raw) + timestep = sigma * 1000.0 + sigma_view = sigma.reshape(batch, 1, 1, 1, 1) + noise = torch.randn_like(clean) + noisy = (1.0 - sigma_view) * clean + sigma_view * noise + return noisy.to(dtype=torch.bfloat16), noise, timestep, sigma_view + + +def _critic_loss( + fake_score: torch.nn.Module, + clean: torch.Tensor, + conditions: dict[str, Any], + args: argparse.Namespace, +) -> tuple[torch.Tensor, torch.Tensor]: + clean = clean.detach().to(dtype=torch.bfloat16) + noisy, noise, timestep, _ = _sample_score_noise(clean, args) + predicted_flow = _score_forward( + fake_score, + noisy=noisy, + timestep=timestep, + conditions=conditions, + branch="positive", + ) + target_flow = noise.float() - clean.float() + return F.mse_loss(predicted_flow.float(), target_flow), timestep.detach() + + +def _dmd_loss( + real_score: torch.nn.Module, + fake_score: torch.nn.Module, + clean: torch.Tensor, + conditions: dict[str, Any], + args: argparse.Namespace, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + noisy, _, timestep, sigma = _sample_score_noise(clean.detach().to(torch.bfloat16), args) + with torch.no_grad(): + fake_flow = _score_forward( + fake_score, + noisy=noisy, + timestep=timestep, + conditions=conditions, + branch="positive", + ) + real_cond = _score_forward( + real_score, + noisy=noisy, + timestep=timestep, + conditions=conditions, + branch="positive", + ) + real_uncond = _score_forward( + real_score, + noisy=noisy, + timestep=timestep, + conditions=conditions, + branch="negative", + ) + real_flow = real_uncond + args.real_guidance_scale * (real_cond - real_uncond) + fake_x0 = noisy.float() - sigma * fake_flow.float() + real_x0 = noisy.float() - sigma * real_flow.float() + gradient = fake_x0 - real_x0 + normalizer = (clean.detach().float() - real_x0).abs().mean( + dim=(1, 2, 3, 4), keepdim=True + ).clamp_min(1e-6) + gradient = torch.nan_to_num(gradient / normalizer) + target = (clean.detach().float() - gradient).detach() + loss = 0.5 * F.mse_loss(clean.float(), target) + return loss, { + "gradient_norm": gradient.abs().mean().detach(), + "score_timestep": timestep.mean().detach(), + "normalizer": normalizer.mean().detach(), + } + + +class Trainer: + def __init__(self, args: argparse.Namespace) -> None: + self.args = args + self.rank = int(os.environ["RANK"]) + self.world_size = int(os.environ["WORLD_SIZE"]) + self.local_rank = int(os.environ["LOCAL_RANK"]) + if self.world_size != 8: + raise ValueError(f"Random-exit DMD requires 8 ranks, got {self.world_size}") + torch.cuda.set_device(self.local_rank) + self.device = torch.device("cuda", self.local_rank) + dist.init_process_group("nccl", device_id=self.device) + initialize_parallel_state(sp=1) + initialize_runtime_state() + self.is_main = self.rank == 0 + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision("high") + random.seed(args.seed + self.rank) + np.random.seed(args.seed + self.rank) + torch.manual_seed(args.seed + self.rank) + torch.cuda.manual_seed_all(args.seed + self.rank) + + self.output_dir = args.output_dir.resolve() + if self.is_main: + self.output_dir.mkdir(parents=True, exist_ok=True) + dist.barrier() + self.log_path = self.output_dir / "train_log.jsonl" + + cases_path = ( + args.cases_path.resolve() + if args.cases_path is not None + else args.data_root.resolve() / "cases.jsonl" + ) + dataset = PredictorV4TrajectoryDataset( + cases_path, + tensor_root=args.data_root.resolve(), + seed=args.latent_seed, + max_trajectories=args.max_trajectories, + ) + self.sampler = DistributedSampler( + dataset, + num_replicas=self.world_size, + rank=self.rank, + shuffle=True, + seed=args.seed, + drop_last=True, + ) + self.loader = DataLoader( + dataset, + batch_size=1, + sampler=self.sampler, + drop_last=True, + num_workers=0, + collate_fn=trajectory_collate, + ) + self.iterator = None + self.epoch = 0 + + self.pipe = HunyuanVideo_1_5_Pipeline.create_pipeline( + pretrained_model_name_or_path=str(args.base_model.resolve()), + transformer_version="480p_i2v", + enable_offloading=True, + enable_group_offloading=False, + create_sr_pipeline=False, + force_sparse_attn=False, + transformer_dtype=torch.bfloat16, + action_ckpt=str(args.rollout_teacher_checkpoint.resolve()), + ) + rollout_teacher = self.pipe.transformer + configure_exact_teacher(rollout_teacher) + rollout_teacher.eval().requires_grad_(False) + self.condition_time_in = copy.deepcopy(rollout_teacher.time_in).to( + device=self.device, dtype=torch.bfloat16 + ).eval().requires_grad_(False) + self.condition_action_in = copy.deepcopy(rollout_teacher.action_in).to( + device=self.device, dtype=torch.bfloat16 + ).eval().requires_grad_(False) + + predictor_model = HYWorldPlayPredictorV4( + PredictorV2Config(source_block_ids=(1, 52)) + ) + predictor_model.load_teacher_initialization(args.rollout_teacher_checkpoint) + configure_model_precision( + predictor_model, + device=self.device, + fp32_trainable_params=True, + ) + _load_exact_stage1( + predictor_model, + args.stage1_training_state.resolve(), + (1, 52), + ) + predictor_model.train() + self.predictor_model = predictor_model + self.predictor = DDP( + predictor_model, + device_ids=[self.local_rank], + output_device=self.local_rank, + broadcast_buffers=False, + gradient_as_bucket_view=True, + find_unused_parameters=False, + ) + + rollout_teacher = _fsdp_wrap(rollout_teacher, self.device) + self.pipe.transformer = rollout_teacher + self.rollout_teacher = rollout_teacher + + real = _load_score_model( + base_model=args.base_model.resolve(), + checkpoint=args.score_teacher_checkpoint.resolve(), + trainable=False, + activation_checkpointing=False, + ) + fake = _load_score_model( + base_model=args.base_model.resolve(), + checkpoint=args.score_teacher_checkpoint.resolve(), + trainable=True, + activation_checkpointing=True, + ) + self.real_score = _fsdp_wrap(real, self.device) + self.fake_score = _fsdp_wrap(fake, self.device) + self.real_score.eval() + self.fake_score.train() + + fusion, blocks = _predictor_groups(self.predictor_model) + self.predictor_optimizer = AdamW( + [ + {"params": fusion, "lr": args.fusion_lr, "name": "fusion_residual"}, + {"params": blocks, "lr": args.blocks_lr, "name": "blocks"}, + ], + betas=(0.0, 0.999), + weight_decay=args.weight_decay, + ) + self.predictor_scheduler = _make_scheduler( + self.predictor_optimizer, + warmup=args.warmup_steps, + total=args.max_steps, + ) + self.fake_optimizer = AdamW( + [parameter for parameter in self.fake_score.parameters() if parameter.requires_grad], + lr=args.fake_score_lr, + betas=(0.0, 0.999), + weight_decay=args.weight_decay, + ) + self.points_local = generate_points_in_sphere(50000, 8.0).to(self.device) + self.student_step = 0 + self.critic_step = 0 + self.exit_counts = torch.zeros(4, device=self.device, dtype=torch.long) + + if args.resume is not None: + self._load_resume(args.resume.resolve()) + + run_config = { + **vars(args), + "base_model": str(args.base_model.resolve()), + "rollout_teacher_checkpoint": str(args.rollout_teacher_checkpoint.resolve()), + "score_teacher_checkpoint": str(args.score_teacher_checkpoint.resolve()), + "stage1_training_state": str(args.stage1_training_state.resolve()), + "data_root": str(args.data_root.resolve()), + "cases_path": str(cases_path), + "output_dir": str(self.output_dir), + "world_size": self.world_size, + "training_mode": "predictor_only_random_exit_dmd", + "random_exits": list(TRAIN_EXITS), + "regression_loss": False, + "rollout_model": "4step_ar_distilled_action_model", + "real_score_model": "50step_bidirectional_model_frozen", + "fake_score_model": "50step_bidirectional_model_trainable", + "stage1_expected_step": 2000, + "predictor_trainable_parameters": sum( + p.numel() for p in predictor_model.parameters() if p.requires_grad + ), + } + if self.is_main: + _json_dump(self.output_dir / "train_config.json", run_config) + self.swanlab_run = self._init_swanlab(run_config) + torch.cuda.empty_cache() + + def _init_swanlab(self, config: dict[str, Any]): + if not self.is_main or self.args.swanlab_mode == "disabled": + return None + if self.args.swanlab_mode == "cloud": + api_key = os.environ.get("SWANLAB_API_KEY") + if api_key: + swanlab.login(api_key=api_key, save=False) + else: + swanlab.login() + return swanlab.init( + project=self.args.swanlab_project, + workspace=self.args.swanlab_workspace, + experiment_name=self.args.swanlab_experiment, + description=( + "Stage-1 step-2000 initialized Predictor-only DMD; random shared " + "P1/P2/P3 exits; no hidden/velocity regression" + ), + tags=["predictor-v4", "random-exit", "dmd", "predictor-only", "8gpu"], + config=config, + logdir=str(self.output_dir / "swanlab"), + mode=self.args.swanlab_mode, + ) + + def _next_batch(self) -> dict[str, Any]: + if self.iterator is None: + self.sampler.set_epoch(self.epoch) + self.iterator = iter(self.loader) + try: + return next(self.iterator) + except StopIteration: + self.epoch += 1 + self.sampler.set_epoch(self.epoch) + self.iterator = iter(self.loader) + return next(self.iterator) + + def _load_resume(self, checkpoint: Path) -> None: + predictor_path = checkpoint / "predictor_training.pt" + fake_path = checkpoint / f"fake_score_rank_{self.rank:02d}.pt" + if not predictor_path.is_file() or not fake_path.is_file(): + raise FileNotFoundError(f"Incomplete DMD resume checkpoint: {checkpoint}") + state = torch.load(predictor_path, map_location="cpu", weights_only=False) + result = self.predictor_model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + name for name in result.missing_keys if name.startswith(TRAINABLE_PREFIXES) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Predictor resume mismatch: " + f"unexpected={result.unexpected_keys}, missing={missing_trainable}" + ) + self.predictor_optimizer.load_state_dict(state["optimizer"]) + self.predictor_scheduler.load_state_dict(state["scheduler"]) + self.student_step = int(state["student_step"]) + self.critic_step = int(state["critic_step"]) + self.epoch = int(state.get("epoch", 0)) + self.exit_counts.copy_(state["exit_counts"].to(self.device)) + fake_state = torch.load(fake_path, map_location="cpu", weights_only=False) + self.fake_score.load_state_dict(fake_state["model"], strict=True) + self.fake_optimizer.load_state_dict(fake_state["optimizer"]) + dist.barrier() + + def _rollout(self, *, exit_step: int, retain_graph: bool): + return _generate_random_exit_rollout( + pipe=self.pipe, + teacher=self.rollout_teacher, + predictor=self.predictor, + predictor_model=self.predictor_model, + condition_time_in=self.condition_time_in, + condition_action_in=self.condition_action_in, + record=self._next_batch(), + points_local=self.points_local, + exit_step=exit_step, + retain_predictor_graph=retain_graph, + device=self.device, + ) + + def critic_update(self, exit_step: int) -> dict[str, float]: + self.fake_optimizer.zero_grad(set_to_none=True) + with torch.no_grad(): + clean, conditions = self._rollout(exit_step=exit_step, retain_graph=False) + loss, timestep = _critic_loss(self.fake_score, clean, conditions, self.args) + loss.backward() + grad_norm = self.fake_score.clip_grad_norm_(self.args.fake_score_grad_clip) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite fake-score grad norm: {grad_norm}") + self.fake_optimizer.step() + self.critic_step += 1 + return { + "critic_loss": float(loss.detach()), + "critic_grad_norm": float(grad_norm), + "critic_timestep": float(timestep.mean()), + } + + def student_update(self, exit_step: int) -> dict[str, float]: + self.predictor_optimizer.zero_grad(set_to_none=True) + clean, conditions = self._rollout(exit_step=exit_step, retain_graph=True) + loss, logs = _dmd_loss( + self.real_score, self.fake_score, clean, conditions, self.args + ) + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_( + self.predictor_model.parameters(), self.args.predictor_grad_clip + ) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite Predictor grad norm: {grad_norm}") + self.predictor_optimizer.step() + self.predictor_scheduler.step() + self.student_step += 1 + self.exit_counts[exit_step] += 1 + return { + "dmd_loss": float(loss.detach()), + "dmd_gradient_norm": float(logs["gradient_norm"]), + "dmd_score_timestep": float(logs["score_timestep"]), + "dmd_normalizer": float(logs["normalizer"]), + "predictor_grad_norm": float(grad_norm), + } + + def _save(self, final: bool = False) -> None: + dist.barrier() + suffix = "final" if final else f"step_{self.student_step:05d}" + checkpoint = self.output_dir / f"checkpoint_{suffix}" + checkpoint.mkdir(parents=True, exist_ok=True) + if self.is_main: + save_predictor_weights( + self.predictor_model, + checkpoint / "predictor.safetensors", + floating_dtype=torch.bfloat16, + ) + torch.save( + { + "model": trainable_state_dict(self.predictor_model), + "optimizer": self.predictor_optimizer.state_dict(), + "scheduler": self.predictor_scheduler.state_dict(), + "student_step": self.student_step, + "critic_step": self.critic_step, + "epoch": self.epoch, + "exit_counts": self.exit_counts.cpu(), + "predictor_config": self.predictor_model.predictor_config.__dict__, + }, + checkpoint / "predictor_training.pt", + ) + torch.save( + { + "model": self.fake_score.state_dict(), + "optimizer": self.fake_optimizer.state_dict(), + "rank": self.rank, + "world_size": self.world_size, + }, + checkpoint / f"fake_score_rank_{self.rank:02d}.pt", + ) + dist.barrier() + if self.is_main: + _json_dump( + checkpoint / "state.json", + { + "student_step": self.student_step, + "critic_step": self.critic_step, + "exit_counts": self.exit_counts.cpu().tolist(), + "final": final, + }, + ) + latest = self.output_dir / "latest" + temporary = self.output_dir / f".latest.{os.getpid()}" + if temporary.exists() or temporary.is_symlink(): + temporary.unlink() + temporary.symlink_to(checkpoint.name) + os.replace(temporary, latest) + snapshots = sorted( + path for path in self.output_dir.glob("checkpoint_step_*") if path.is_dir() + ) + for old in snapshots[: -self.args.keep_checkpoints]: + shutil.rmtree(old) + dist.barrier() + + def train(self) -> None: + try: + while self.student_step < self.args.max_steps: + started = time.perf_counter() + critic_sums = {"critic_loss": 0.0, "critic_grad_norm": 0.0, "critic_timestep": 0.0} + for _ in range(self.args.critic_updates_per_student): + critic_exit = _sample_exit( + self.device, self.args.forced_exit_step + ) + metrics = self.critic_update(critic_exit) + for key in critic_sums: + critic_sums[key] += metrics[key] + exit_step = _sample_exit(self.device, self.args.forced_exit_step) + student = self.student_update(exit_step) + values = torch.tensor( + [ + student["dmd_loss"], + student["dmd_gradient_norm"], + student["dmd_score_timestep"], + student["dmd_normalizer"], + student["predictor_grad_norm"], + critic_sums["critic_loss"] / self.args.critic_updates_per_student, + critic_sums["critic_grad_norm"] / self.args.critic_updates_per_student, + critic_sums["critic_timestep"] / self.args.critic_updates_per_student, + time.perf_counter() - started, + ], + device=self.device, + dtype=torch.float64, + ) + dist.all_reduce(values) + values /= self.world_size + counts = self.exit_counts.clone() + dist.all_reduce(counts) + counts //= self.world_size + record = { + "student_step": self.student_step, + "critic_step": self.critic_step, + "exit_step": exit_step, + "dmd_loss": float(values[0]), + "dmd_gradient_norm": float(values[1]), + "dmd_score_timestep": float(values[2]), + "dmd_normalizer": float(values[3]), + "predictor_grad_norm": float(values[4]), + "critic_loss": float(values[5]), + "critic_grad_norm": float(values[6]), + "critic_timestep": float(values[7]), + "step_time_s": float(values[8]), + "lr_predictor_fusion": self.predictor_optimizer.param_groups[0]["lr"], + "lr_predictor_blocks": self.predictor_optimizer.param_groups[1]["lr"], + "lr_fake_score": self.fake_optimizer.param_groups[0]["lr"], + "exit_count_p1": int(counts[1]), + "exit_count_p2": int(counts[2]), + "exit_count_p3": int(counts[3]), + "peak_memory_gib": torch.cuda.max_memory_allocated(self.device) / 2**30, + } + if self.is_main and ( + self.student_step == 1 or self.student_step % self.args.log_every == 0 + ): + with self.log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + print(json.dumps(record, sort_keys=True), flush=True) + if self.swanlab_run is not None: + swanlab.log(record, step=self.student_step) + if not self.args.no_save and ( + self.student_step % self.args.save_every == 0 + or self.student_step == self.args.max_steps + ): + self._save(final=self.student_step == self.args.max_steps) + del values, counts, student, critic_sums + gc.collect() + torch.cuda.empty_cache() + except BaseException as error: + if self.is_main and self.swanlab_run is not None: + try: + swanlab.finish(error=str(error)) + except Exception: + pass + raise + else: + if self.is_main and self.swanlab_run is not None: + swanlab.finish() + finally: + dist.destroy_process_group() + + +def validate_config(args: argparse.Namespace) -> dict[str, Any]: + cases = args.cases_path or args.data_root / "cases.jsonl" + required = [ + args.base_model, + args.rollout_teacher_checkpoint, + args.score_teacher_checkpoint, + args.data_root, + cases, + args.stage1_training_state, + ] + missing = [str(path) for path in required if not path.resolve().exists()] + if missing: + raise FileNotFoundError(f"Missing required paths: {missing}") + if not 0.0 < args.score_timestep_min < args.score_timestep_max < 1.0: + raise ValueError("score timestep bounds must satisfy 0 < min < max < 1") + if args.critic_updates_per_student < 1: + raise ValueError("critic_updates_per_student must be positive") + return { + "valid": True, + "training_mode": "predictor_only_random_exit_dmd", + "random_exits": list(TRAIN_EXITS), + "regression_loss": False, + "stage1_training_state": str(args.stage1_training_state.resolve()), + "rollout_teacher": str(args.rollout_teacher_checkpoint.resolve()), + "score_teacher": str(args.score_teacher_checkpoint.resolve()), + } + + +def main() -> None: + args = parse_args() + result = validate_config(args) + if args.validate_config_only: + print(json.dumps(result, ensure_ascii=False, sort_keys=True)) + return + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + raise RuntimeError("Launch random-exit DMD with torchrun") + Trainer(args).train() + + +if __name__ == "__main__": + main() diff --git a/train_predictor_v4_rollout.py b/train_predictor_v4_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..3137afaadb8f464359968f54a08d710648405009 --- /dev/null +++ b/train_predictor_v4_rollout.py @@ -0,0 +1,1753 @@ +#!/usr/bin/env python3 +"""Eight-GPU self-rollout fine-tuning for Predictor-v4.""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import os +import random +import time +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import swanlab +import torch +import torch.distributed as dist +import torch.nn.functional as F +from PIL import Image +from safetensors import safe_open +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from hyvideo.commons import auto_offload_model +from hyvideo.commons.infer_state import initialize_infer_state +from hyvideo.commons.parallel_states import initialize_parallel_state +from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline +from hyvideo.utils.retrieval_context import ( + generate_points_in_sphere, + select_aligned_memory_frames, +) +from models.predictor_v2 import PredictorV2Config +from models.predictor_v4 import HYWorldPlayPredictorV4 +from predictor_data.text_kv_schema import ( + TEXT_KV_HEADS, + TEXT_KV_HEAD_DIM, + TEXT_KV_LAYER_COUNT, + TEXT_KV_TOKENS, + expected_text_kv_keys, + text_kv_key, +) +from predictor_training.checkpoint import ( + TRAINABLE_PREFIXES, + load_concat_fusion_initialization, + load_concat_residual_gate_initialization, + load_predictor_weights, + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.trajectory_dataset import ( + PredictorV4TrajectoryDataset, + trajectory_collate, +) +from predictor_training.rollout_utils import configure_exact_teacher, pose_to_input + + +DEFAULT_NVME_ROOT = Path("/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4") +DEFAULT_TEXT_KV_ROOT = DEFAULT_NVME_ROOT / "text_kv_exact_all54_bf16" +DEFAULT_BASE = Path( + "/mnt/s3files/s3-us-west2-default/zoubin/cz/checkpoints/hy_worldplay/huggingface/hub/" + "models--tencent--HunyuanVideo-1.5/snapshots/9b49404b3f5df2a8f0b31df27a0c7ab872e7b038" +) +DEFAULT_TEACHER = Path( + "/mnt/s3files/s3-us-west2-default/zoubin/cz/checkpoints/hy_worldplay/huggingface/hub/" + "models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/" + "ar_distilled_action_model/diffusion_pytorch_model.safetensors" +) +DEFAULT_INIT = Path( + "checkpoints/predictor_v4_prefeature_blocks1-52_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/" + "predictor_step_02000.safetensors" +) +DEFAULT_OUTPUT = Path( + "checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_" + "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab" +) + + +def parse_source_blocks(value: str) -> tuple[int, int]: + result = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + if result not in {(0, 53), (1, 52)}: + raise argparse.ArgumentTypeError("source_blocks must be 0,53 or 1,52") + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data_root", type=Path, default=DEFAULT_NVME_ROOT) + parser.add_argument( + "--cases_path", + type=Path, + default=None, + help=( + "Optional case manifest; tensors still come from --data_root. " + "Useful for reproducible overfit subsets." + ), + ) + parser.add_argument("--base_model", type=Path, default=DEFAULT_BASE) + parser.add_argument("--teacher_checkpoint", type=Path, default=DEFAULT_TEACHER) + parser.add_argument( + "--text_kv_root", + type=Path, + default=None, + help=( + "Optional exact 54-layer BF16 Text K/V directory. When set, skip " + "prompt, vision, ByT5 encoding and online text prefill." + ), + ) + parser.add_argument("--init_weights", type=Path, default=DEFAULT_INIT) + parser.add_argument( + "--init_training_state", + type=Path, + default=None, + help=( + "Optional offline training_latest.pt whose FP32 model state initializes " + "a fresh rollout optimizer; takes precedence over --init_weights" + ), + ) + parser.add_argument("--output_dir", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--source_blocks", type=parse_source_blocks, default=(1, 52)) + parser.add_argument( + "--predictor_fusion_type", + choices=( + "concat", + "projection_add", + "token_gate_prev", + "concat_residual_gate_prev", + ), + default="concat", + ) + parser.add_argument("--gate_hidden_dim", type=int, default=256) + parser.add_argument("--gate_initial_probability", type=float, default=0.8) + parser.add_argument("--max_steps", type=int, default=2000) + parser.add_argument( + "--gradient_accumulation_steps", type=int, default=1, + help="Trajectories per rank accumulated into one optimizer step " + "(global trajectories per step = world_size x this).", + ) + parser.add_argument( + "--expected_world_size", type=int, default=8, + help="Number of ranks the launch must provide (the original recipe used 8).", + ) + parser.add_argument( + "--predictor_arch", choices=("auto", "baseline", "atc"), default="auto", + help="auto rebuilds the architecture (incl. ATC transport settings) from the " + "predictor_config saved in --init_training_state / --resume.", + ) + parser.add_argument("--num_chunks", type=int, default=8) + parser.add_argument( + "--rollout_schedule", + choices=("fppp", "fppf"), + default="fppp", + help=( + "Student latent rollout for chunks after chunk 0. Both schedules train " + "Predictor steps 1, 2, and 3; fppf uses the online Full-Teacher step-3 " + "velocity and hidden as the next-chunk state." + ), + ) + parser.add_argument( + "--supervision_mode", + choices=("onpolicy", "offline_ffff"), + default="onpolicy", + help=( + "Query the Full Teacher at each current student state, or supervise " + "student self-rollout with fixed hidden/velocity targets from the " + "offline FFFF trajectory. offline_ffff also restores Chunk 0 from " + "the saved FFFF tensors instead of recomputing it online." + ), + ) + parser.add_argument("--max_trajectories", type=int, default=None) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--latent_seed", type=int, default=0, + help="Initial 32-frame noise seed, independent of sampler shuffle seed.", + ) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--scheduler_warmup_steps", type=int, default=100) + parser.add_argument("--block_warmup_steps", type=int, default=100) + parser.add_argument( + "--block_warmup_mode", choices=("delayed", "linear"), default="delayed" + ) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--fp32_trainable_params", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Keep requires_grad=True Predictor parameters and AdamW moments in " + "FP32 while frozen parameters and forward activations remain BF16" + ), + ) + parser.add_argument( + "--detach_between_timesteps", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Feed Predictor hidden/velocity into the next timestep while detaching " + "both paths, so later losses cannot update earlier Predictor calls." + ), + ) + parser.add_argument("--save_every", type=int, default=100) + parser.add_argument("--keep_snapshots", type=int, default=20) + parser.add_argument("--log_every", type=int, default=1) + parser.add_argument("--resume", type=Path, default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v4") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument( + "--swanlab_experiment", + default=( + "predictor-v4-rollout-fppp-blocks1-52-fp32master-8gpu-" + "flr1e-4-blr1e-5-schedsteps2000-swarmup100-bwarmup100" + ), + ) + parser.add_argument( + "--swanlab_mode", choices=("cloud", "offline", "local", "disabled"), default="cloud" + ) + parser.add_argument("--validate_scheduler_only", action="store_true") + return parser.parse_args() + + +def initialize_runtime_state() -> None: + class Args: + sage_blocks_range = "0-53" + use_sageattn = False + enable_torch_compile = False + use_fp8_gemm = False + quant_type = "fp8-per-block" + include_patterns = "double_blocks" + use_vae_parallel = False + + initialize_infer_state(Args()) + + +def cosine_factor(step: int, warmup: int, total: int) -> float: + if step < warmup: + return max(1e-8, float(step + 1) / max(1, warmup)) + progress = min(1.0, float(step - warmup) / max(1, total - warmup)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + +def make_lr_scheduler(optimizer: torch.optim.Optimizer, args: argparse.Namespace) -> LambdaLR: + block_warmup_mode = getattr(args, "block_warmup_mode", "delayed") + if block_warmup_mode not in {"delayed", "linear"}: + raise ValueError(f"Unsupported block warmup mode: {block_warmup_mode}") + + def block_factor(step: int) -> float: + if block_warmup_mode == "linear": + return cosine_factor(step, args.block_warmup_steps, args.max_steps) + return 0.0 if step < args.block_warmup_steps else cosine_factor( + step, args.scheduler_warmup_steps, args.max_steps + ) + + return LambdaLR( + optimizer, + [ + lambda step: cosine_factor(step, args.scheduler_warmup_steps, args.max_steps), + block_factor, + ], + ) + + +def configure_model_precision( + model: torch.nn.Module, + *, + device: torch.device, + fp32_trainable_params: bool, +) -> None: + """Keep frozen state in BF16 and optionally retain FP32 master parameters.""" + model.to(device=device, dtype=torch.bfloat16) + if fp32_trainable_params: + with torch.no_grad(): + for parameter in model.parameters(): + if parameter.requires_grad: + parameter.data = parameter.data.to(dtype=torch.float32) + + +def validate_optimizer_precision( + optimizer: torch.optim.Optimizer, + *, + fp32_trainable_params: bool, + require_state: bool, +) -> None: + expected = torch.float32 if fp32_trainable_params else torch.bfloat16 + for group in optimizer.param_groups: + for parameter in group["params"]: + if parameter.dtype != expected: + raise RuntimeError( + f"Optimizer parameter dtype {parameter.dtype} != expected {expected}" + ) + state = optimizer.state.get(parameter, {}) + for name in ("exp_avg", "exp_avg_sq"): + value = state.get(name) + if value is None: + if require_state: + raise RuntimeError(f"Optimizer state {name} has not been initialized") + continue + if value.dtype != expected: + raise RuntimeError( + f"Optimizer state {name} dtype {value.dtype} != expected {expected}" + ) + + +def load_initial_training_state( + model: torch.nn.Module, + path: Path, + source_blocks: tuple[int, int], +) -> None: + """Load only the exact model master state and intentionally reset optimizer state.""" + state = torch.load(path, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != source_blocks: + raise ValueError("Initialization training-state source blocks do not match") + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + state.get("args", {}).get("predictor_fusion_type", "concat"), + ) + target_fusion_type = model.predictor_config.predictor_fusion_type + if ( + saved_fusion_type == "concat" + and target_fusion_type == "concat_residual_gate_prev" + ): + report = load_concat_residual_gate_initialization(model, state["model"]) + print( + json.dumps( + { + "fusion_initialization": "concat_residual_gate", + "fresh_residual_key_count": len( + report["fresh_residual_keys"] + ), + }, + sort_keys=True, + ), + flush=True, + ) + return + if saved_fusion_type == "concat" and target_fusion_type != "concat": + report = load_concat_fusion_initialization(model, state["model"]) + print( + json.dumps( + { + "fusion_initialization": "concat_conversion", + "discarded_concat_key_count": len( + report["discarded_concat_keys"] + ), + "fresh_fusion_key_count": len(report["fresh_fusion_keys"]), + }, + sort_keys=True, + ), + flush=True, + ) + return + if saved_fusion_type != target_fusion_type: + raise ValueError( + "Initialization training-state fusion type does not match: " + f"saved={saved_fusion_type}, requested={target_fusion_type}" + ) + result = model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + key for key in result.missing_keys if key.startswith(TRAINABLE_PREFIXES) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Initialization training-state mismatch; " + f"unexpected={result.unexpected_keys}, missing_trainable={missing_trainable}" + ) + + +def differentiable_euler_step( + sample: torch.Tensor, + model_output: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, +) -> torch.Tensor: + """Stateless equivalent of the project's FlowMatch Euler scheduler step.""" + dt = (sigmas[step_index + 1] - sigmas[step_index]).to( + device=sample.device, dtype=torch.float32 + ) + return sample.float() + model_output.float() * dt + + +def offline_step_tensor_path( + data_root: Path, + record: dict, + chunk_id: int, +) -> Path: + path = ( + data_root + / "steps" + / f"case_{int(record['case_id']):04d}" + / f"action_{int(record['action_id']):02d}" + / f"chunk_{chunk_id:02d}.safetensors" + ) + if not path.is_file(): + raise FileNotFoundError(path) + return path + + +def validate_offline_timestep( + saved: torch.Tensor, + expected: torch.Tensor, + *, + path: Path, + step_id: int, +) -> None: + if saved.dtype != torch.float32 or tuple(saved.shape) != (1,): + raise TypeError( + f"{path}: step {step_id} timestep must be FP32 [1], " + f"got {saved.dtype} {tuple(saved.shape)}" + ) + expected_cpu = expected.detach().float().cpu().reshape(1) + if not torch.equal(saved, expected_cpu): + raise ValueError( + f"{path}: step {step_id} timestep {saved.tolist()} does not match " + f"current scheduler {expected_cpu.tolist()}" + ) + + +def validate_offline_tensor( + value: torch.Tensor, + expected_shape: tuple[int, ...], + *, + path: Path, + name: str, +) -> None: + if value.dtype != torch.bfloat16 or tuple(value.shape) != expected_shape: + raise TypeError( + f"{path}: {name} must be BF16 {expected_shape}, " + f"got {value.dtype} {tuple(value.shape)}" + ) + + +@torch.no_grad() +def load_offline_chunk0_state( + path: Path, + *, + device: torch.device, + sigmas: torch.Tensor, + timesteps: torch.Tensor, +) -> tuple[torch.Tensor, dict[int, torch.Tensor]]: + """Restore Chunk-0 FFFF final latent and per-step hidden from disk.""" + names = { + "step_3_noisy_sample", + "step_3_velocity", + "step_3_timestep", + *(f"step_{step_id}_final_hidden" for step_id in (1, 2, 3)), + } + with safe_open(str(path), framework="pt", device="cpu") as handle: + missing = names.difference(handle.keys()) + if missing: + raise KeyError(f"{path}: missing offline Chunk-0 tensors {sorted(missing)}") + noisy_sample = handle.get_tensor("step_3_noisy_sample") + velocity = handle.get_tensor("step_3_velocity") + saved_timestep = handle.get_tensor("step_3_timestep") + hidden_by_step = { + step_id: handle.get_tensor(f"step_{step_id}_final_hidden") + for step_id in (1, 2, 3) + } + latent_shape = (1, 32, 4, 30, 52) + hidden_shape = (1, 6240, 2048) + validate_offline_tensor( + noisy_sample, latent_shape, path=path, name="step_3_noisy_sample" + ) + validate_offline_tensor( + velocity, latent_shape, path=path, name="step_3_velocity" + ) + validate_offline_timestep( + saved_timestep, timesteps[3], path=path, step_id=3 + ) + for step_id, hidden in hidden_by_step.items(): + validate_offline_tensor( + hidden, + hidden_shape, + path=path, + name=f"step_{step_id}_final_hidden", + ) + final_latent = differentiable_euler_step( + noisy_sample.to(device=device), + velocity.to(device=device), + sigmas, + 3, + ).detach() + return final_latent, { + step_id: hidden.to(device=device) + for step_id, hidden in hidden_by_step.items() + } + + +def load_offline_ffff_targets( + path: Path, + *, + timesteps: torch.Tensor, +) -> dict[int, dict[str, torch.Tensor]]: + """Load fixed FFFF P1/P2/P3 labels and latent drift references on CPU.""" + targets: dict[int, dict[str, torch.Tensor]] = {} + latent_shape = (1, 32, 4, 30, 52) + hidden_shape = (1, 6240, 2048) + frame_condition_shape = (1, 4, 2048) + with safe_open(str(path), framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + for step_id in (1, 2, 3): + names = { + "noisy_sample": f"step_{step_id}_noisy_sample", + "hidden": f"step_{step_id}_final_hidden", + "velocity": f"step_{step_id}_velocity", + "frame_condition": f"step_{step_id}_frame_condition", + "timestep": f"step_{step_id}_timestep", + } + missing = set(names.values()).difference(keys) + if missing: + raise KeyError( + f"{path}: missing offline step-{step_id} tensors " + f"{sorted(missing)}" + ) + values = {name: handle.get_tensor(key) for name, key in names.items()} + validate_offline_tensor( + values["noisy_sample"], + latent_shape, + path=path, + name=names["noisy_sample"], + ) + validate_offline_tensor( + values["hidden"], + hidden_shape, + path=path, + name=names["hidden"], + ) + validate_offline_tensor( + values["velocity"], + latent_shape, + path=path, + name=names["velocity"], + ) + validate_offline_tensor( + values["frame_condition"], + frame_condition_shape, + path=path, + name=names["frame_condition"], + ) + validate_offline_timestep( + values["timestep"], + timesteps[step_id], + path=path, + step_id=step_id, + ) + targets[step_id] = values + return targets + + +def validate_differentiable_scheduler() -> dict[str, float]: + scheduler = HunyuanVideo_1_5_Pipeline._create_scheduler(5.0) + scheduler.set_timesteps(4, device=torch.device("cpu")) + generator = torch.Generator().manual_seed(123) + outputs = [torch.randn((1, 2, 2), generator=generator) for _ in range(4)] + expected = torch.randn((1, 2, 2), generator=generator) + actual = expected.clone() + scheduler.set_timesteps(4, device=torch.device("cpu")) + for index, timestep in enumerate(scheduler.timesteps): + expected = scheduler.step(outputs[index], timestep, expected, return_dict=False)[0] + actual = differentiable_euler_step(actual, outputs[index], scheduler.sigmas, index) + error = float((actual - expected).abs().max()) + if error > 1e-7: + raise RuntimeError(f"Differentiable scheduler mismatch: {error}") + sample = torch.randn((1, 2, 2), requires_grad=True) + velocity = torch.randn((1, 2, 2), requires_grad=True) + differentiable_euler_step(sample, velocity, scheduler.sigmas, 0).square().mean().backward() + grad = float(velocity.grad.abs().sum()) + if not math.isfinite(grad) or grad == 0.0: + raise RuntimeError("Differentiable scheduler did not propagate velocity gradients") + return {"scheduler_max_abs_error": error, "scheduler_velocity_grad_l1": grad} + + +def load_image_condition(path: str | Path, device: torch.device) -> torch.Tensor: + with safe_open(str(path), framework="pt", device="cpu") as handle: + if "image_condition_latent" not in handle.keys(): + raise KeyError(f"Missing image_condition_latent in {path}") + value = handle.get_tensor("image_condition_latent") + if tuple(value.shape) != (1, 32, 1, 30, 52): + raise ValueError(f"Unexpected image condition shape: {tuple(value.shape)}") + return value.to(device=device, dtype=torch.bfloat16) + + +@torch.no_grad() +def prepare_rollout_inputs( + pipe: HunyuanVideo_1_5_Pipeline, + record: dict, + device: torch.device, +) -> dict: + """Build only policy-independent latent, condition, camera, and action inputs.""" + scheduler = pipe._create_scheduler(pipe.config.flow_shift) + scheduler.set_timesteps(4, device=device) + generator = torch.Generator(device=device).manual_seed(int(record["seed"])) + latents = pipe.prepare_latents( + 1, 32, 30, 52, 32, torch.bfloat16, device, generator + ) + image_condition = load_image_condition(record["case_tensor_path"], device) + cond_latents = pipe._prepare_cond_latents( + "i2v", image_condition, latents, pipe.get_task_mask("i2v", 32) + ).to(dtype=torch.bfloat16) + viewmats, Ks, action = pose_to_input(record["pose"], 32) + return { + "latents": latents, + "cond_latents": cond_latents, + "timesteps": scheduler.timesteps, + "sigmas": scheduler.sigmas, + "viewmats_cpu": viewmats.unsqueeze(0), + "Ks_cpu": Ks.unsqueeze(0), + "action_cpu": action.unsqueeze(0), + } + + +def load_precomputed_text_cache( + cache_path: Path, + device: torch.device, +) -> list[dict[str, torch.Tensor | None]]: + """Load exact all-layer Text K/V and strip the fixed 903-token padding.""" + with safe_open(str(cache_path), framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + missing = expected_text_kv_keys().difference(keys) + unexpected = keys.difference(expected_text_kv_keys()) + if missing or unexpected: + raise ValueError( + f"Invalid Text K/V cache {cache_path}: missing={sorted(missing)}, " + f"unexpected={sorted(unexpected)}" + ) + mask = handle.get_tensor("text_valid_mask") + if tuple(mask.shape) != (1, TEXT_KV_TOKENS) or mask.dtype != torch.bool: + raise ValueError(f"Invalid text_valid_mask in {cache_path}") + valid_tokens = int(mask.sum().item()) + expected_mask = torch.arange(TEXT_KV_TOKENS)[None] < valid_tokens + if not torch.equal(mask, expected_mask): + raise ValueError(f"Non-prefix text_valid_mask in {cache_path}") + expected_shape = (1, TEXT_KV_HEADS, TEXT_KV_TOKENS, TEXT_KV_HEAD_DIM) + cache: list[dict[str, torch.Tensor | None]] = [] + for block_id in range(TEXT_KV_LAYER_COUNT): + layer: dict[str, torch.Tensor | None] = { + "k_vision": None, + "v_vision": None, + } + for kind in ("k_txt", "v_txt"): + name = text_kv_key(block_id, kind) + value = handle.get_tensor(name) + if tuple(value.shape) != expected_shape or value.dtype != torch.bfloat16: + raise ValueError( + f"Invalid {name} in {cache_path}: " + f"shape={tuple(value.shape)}, dtype={value.dtype}" + ) + layer[kind] = value[:, :, :valid_tokens].to( + device=device, dtype=torch.bfloat16 + ) + cache.append(layer) + return cache + + +@torch.no_grad() +def prepare_online_conditions( + pipe: HunyuanVideo_1_5_Pipeline, + record: dict, + device: torch.device, +) -> dict: + """Build static I2V conditions online while reusing the stored VAE latent.""" + image = Image.open(record["image_path"]).convert("RGB") + image_np = np.array(image) + pipe._guidance_scale = 1.0 + pipe._guidance_rescale = 0.0 + pipe._clip_skip = None + + with auto_offload_model(pipe.text_encoder, device, enabled=pipe.enable_offloading): + prompt_embeds, _, prompt_mask, _ = pipe.encode_prompt( + record["caption"], device, 1, False, "", clip_skip=None, data_type="video" + ) + extra_kwargs = {} + if pipe.config.glyph_byT5_v2: + with auto_offload_model(pipe.byt5_model, device, enabled=pipe.enable_offloading): + extra_kwargs = pipe._prepare_byt5_embeddings(record["caption"], device) + + scheduler = pipe._create_scheduler(pipe.config.flow_shift) + scheduler.set_timesteps(4, device=device) + generator = torch.Generator(device=device).manual_seed(int(record["seed"])) + latents = pipe.prepare_latents( + 1, 32, 30, 52, 32, torch.bfloat16, device, generator + ) + image_condition = load_image_condition(record["case_tensor_path"], device) + cond_latents = pipe._prepare_cond_latents( + "i2v", image_condition, latents, pipe.get_task_mask("i2v", 32) + ).to(dtype=torch.bfloat16) + with auto_offload_model(pipe.vision_encoder, device, enabled=pipe.enable_offloading): + vision_states = pipe._prepare_vision_states( + image_np, pipe.ideal_resolution, latents, device + ) + viewmats, Ks, action = pose_to_input(record["pose"], 32) + return { + "latents": latents, + "cond_latents": cond_latents, + "timesteps": scheduler.timesteps, + "sigmas": scheduler.sigmas, + "prompt_embeds": prompt_embeds, + "prompt_mask": prompt_mask, + "vision_states": vision_states, + "extra_kwargs": extra_kwargs, + "viewmats_cpu": viewmats.unsqueeze(0), + "Ks_cpu": Ks.unsqueeze(0), + "action_cpu": action.unsqueeze(0), + } + + +@torch.no_grad() +def text_prefill( + teacher: torch.nn.Module, + prompt_embeds: torch.Tensor, + prompt_mask: torch.Tensor, + vision_states: torch.Tensor, + extra_kwargs: dict, + device: torch.device, + dtype: torch.dtype, +) -> list[dict]: + cache = [ + {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None} + for _ in teacher.double_blocks + ] + return teacher( + bi_inference=False, + ar_txt_inference=True, + ar_vision_inference=False, + timestep_txt=torch.tensor([0], device=device, dtype=dtype), + text_states=prompt_embeds[0, None, ...], + encoder_attention_mask=prompt_mask[0, None, ...], + vision_states=vision_states[0, None, ...], + mask_type="i2v", + extra_kwargs={ + "byt5_text_states": extra_kwargs["byt5_text_states"][0, None, ...], + "byt5_text_mask": extra_kwargs["byt5_text_mask"][0, None, ...], + }, + kv_cache=cache, + cache_txt=True, + ) + + +def select_context_indices( + viewmats_cpu: torch.Tensor, + chunk_id: int, + points_local: torch.Tensor, + device: torch.device, +) -> list[int]: + if chunk_id == 0: + return [] + current_frame_idx = chunk_id * 4 + selected = select_aligned_memory_frames( + viewmats_cpu[0].numpy(), + current_frame_idx, + memory_frames=20, + temporal_context_size=12, + pred_latent_size=4, + points_local=points_local, + device=device, + ) + return sorted(set(index for index in selected if index < current_frame_idx)) + + +@torch.no_grad() +def history_prefill( + teacher: torch.nn.Module, + cache: list[dict], + latents: torch.Tensor, + cond_latents: torch.Tensor, + selected: list[int], + viewmats: torch.Tensor, + Ks: torch.Tensor, + action: torch.Tensor, + timestep_dtype: torch.dtype, +) -> list[dict]: + if not selected: + return cache + # ``cache_vision=True`` constructs a fresh list and only carries Text K/V + # from the input cache. Drop the previous chunk's large Vision K/V first; + # otherwise two 20-frame, 54-layer caches coexist during prefill. + for layer_cache in cache: + layer_cache["k_vision"] = None + layer_cache["v_vision"] = None + context_input = torch.cat( + [latents[:, :, selected], cond_latents[:, :, selected]], dim=1 + ) + return teacher( + bi_inference=False, + ar_txt_inference=False, + ar_vision_inference=True, + hidden_states=context_input, + timestep=torch.full( + (len(selected),), 14, device=latents.device, dtype=timestep_dtype + ), + timestep_r=None, + mask_type="i2v", + return_dict=False, + viewmats=viewmats[:, selected], + Ks=Ks[:, selected], + action=action[:, selected], + kv_cache=cache, + cache_vision=True, + rope_temporal_size=len(selected), + start_rope_start_idx=0, + context_frame_indices=selected, + ) + + +def teacher_denoise( + teacher: torch.nn.Module, + model_input: torch.Tensor, + timestep_input: torch.Tensor, + viewmats: torch.Tensor, + Ks: torch.Tensor, + action: torch.Tensor, + cache: list[dict], + context_frames: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = teacher( + bi_inference=False, + ar_txt_inference=False, + ar_vision_inference=True, + hidden_states=model_input, + timestep=timestep_input, + timestep_r=None, + mask_type="i2v", + return_dict=False, + viewmats=viewmats, + Ks=Ks, + action=action, + kv_cache=cache, + cache_vision=False, + rope_temporal_size=4 + context_frames, + start_rope_start_idx=context_frames, + return_final_hidden=True, + ) + return output[0], output[2], output[3] + + +def prune_snapshots(output_dir: Path, keep: int) -> None: + snapshots = sorted(output_dir.glob("predictor_step_*.safetensors")) + for path in snapshots[:-keep] if keep > 0 else snapshots: + path.unlink() + + +def distributed_average(values: torch.Tensor) -> torch.Tensor: + dist.all_reduce(values, op=dist.ReduceOp.SUM) + return values / dist.get_world_size() + + +def _average_micro_metrics(metrics_list: list[dict]) -> dict: + """Mean of per-trajectory metrics over one accumulation group; grad_norm + is only defined on the last micro-step, so it is taken from there.""" + if len(metrics_list) == 1: + return metrics_list[0] + out: dict = {} + for key in metrics_list[0]: + values = [m[key] for m in metrics_list] + if key == "grad_norm": + out[key] = values[-1] + elif isinstance(values[0], list): + out[key] = [sum(v[i] for v in values) / len(values) for i in range(len(values[0]))] + else: + out[key] = sum(values) / len(values) + return out + + +def run_trajectory( + *, + pipe: HunyuanVideo_1_5_Pipeline, + teacher: torch.nn.Module, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + optimizer: torch.optim.Optimizer, + record: dict, + points_local: torch.Tensor, + args: argparse.Namespace, + device: torch.device, + accumulation_index: int = 0, + accumulation_steps: int = 1, +) -> dict[str, float | list[float]]: + """One trajectory. Gradients are accumulated over ``accumulation_steps`` + trajectories; the optimizer only steps on the last one.""" + started = time.perf_counter() + is_last_micro_step = accumulation_index == accumulation_steps - 1 + if args.text_kv_root is None: + static = prepare_online_conditions(pipe, record, device) + else: + static = prepare_rollout_inputs(pipe, record, device) + latents = static["latents"] + cond_latents = static["cond_latents"] + timesteps = static["timesteps"] + sigmas = static["sigmas"] + viewmats_cpu = static["viewmats_cpu"] + viewmats = viewmats_cpu.to(device=device, dtype=torch.bfloat16) + Ks = static["Ks_cpu"].to(device=device, dtype=torch.bfloat16) + action = static["action_cpu"].to(device=device, dtype=torch.bfloat16) + + if args.text_kv_root is None: + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + cache = text_prefill( + teacher, + static["prompt_embeds"], + static["prompt_mask"], + static["vision_states"], + static["extra_kwargs"], + device, + torch.bfloat16, + ) + for key in ("prompt_embeds", "prompt_mask", "vision_states", "extra_kwargs"): + del static[key] + else: + cache_path = ( + args.text_kv_root + / f"case_{int(record['case_id']):04d}_text_kv_exact_all54_bf16.safetensors" + ) + cache = load_precomputed_text_cache(cache_path, device) + + if accumulation_index == 0: + optimizer.zero_grad(set_to_none=True) + previous_hidden_by_step: dict[int, torch.Tensor] = {} + metric_sums = torch.zeros(7, device=device, dtype=torch.float64) + # total, h1, h2, h3, v1, v2, v3 + gate_stat_names = ( + "gate_p_mean", + "gate_p_std", + "gate_p_p10", + "gate_p_p50", + "gate_p_p90", + "gate_p_near_zero", + "gate_p_near_one", + ) + gate_metric_sums = torch.zeros(7, device=device, dtype=torch.float64) + gate_metric_count = 0 + atc_stat_names = ( + "delta_h_d_norm", + "gated_delta_h_c_norm", + "transport_attention_entropy", + "transport_argmax_displacement", + ) + atc_metric_sums = torch.zeros(len(atc_stat_names), device=device, dtype=torch.float64) + atc_metric_count = 0 + trajectory_step_losses: list[torch.Tensor] = [] + rollout_state_requires_grad_count = 0 + train_chunks = args.num_chunks - 1 + offline_latent_mse_by_chunk_step: list[torch.Tensor] = [] + + for chunk_id in range(args.num_chunks): + start = chunk_id * 4 + end = start + 4 + if args.supervision_mode == "offline_ffff" and chunk_id == 0: + chunk_latent, previous_hidden_by_step = load_offline_chunk0_state( + offline_step_tensor_path(args.data_root, record, chunk_id), + device=device, + sigmas=sigmas, + timesteps=timesteps, + ) + latents[:, :, start:end] = chunk_latent.to(dtype=latents.dtype) + del chunk_latent + continue + + selected = select_context_indices(viewmats_cpu, chunk_id, points_local, device) + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + cache = history_prefill( + teacher, cache, latents, cond_latents, selected, + viewmats, Ks, action, timesteps.dtype, + ) + chunk_latent = latents[:, :, start:end].detach().clone() + chunk_cond = cond_latents[:, :, start:end] + chunk_viewmats = viewmats[:, start:end] + chunk_Ks = Ks[:, start:end] + chunk_action = action[:, start:end] + current_hidden_by_step: dict[int, torch.Tensor] = {} + full_anchor_hidden = None + + # Chunk 0 is all-Full. Later chunks use Full only at denoise step 0. + full_steps = range(4) if chunk_id == 0 else range(1) + for step_id in full_steps: + timestep_input = torch.full( + (4,), timesteps[step_id], device=device, dtype=timesteps.dtype + ) + model_input = torch.cat([chunk_latent, chunk_cond], dim=1) + velocity, hidden, _ = teacher_denoise( + teacher, model_input, timestep_input, + chunk_viewmats, chunk_Ks, chunk_action, cache, len(selected), + ) + if chunk_id == 0 and step_id > 0: + current_hidden_by_step[step_id] = hidden.detach() + if chunk_id > 0 and step_id == 0: + full_anchor_hidden = hidden.detach() + chunk_latent = differentiable_euler_step( + chunk_latent, velocity, sigmas, step_id + ).detach() + del velocity, hidden, model_input + + if chunk_id == 0: + latents[:, :, start:end] = chunk_latent.detach() + previous_hidden_by_step = current_hidden_by_step + continue + + # DDP all-reduce only on the very last backward of the accumulation group. + sync_context = ( + nullcontext() + if (chunk_id == args.num_chunks - 1 and is_last_micro_step) + else predictor.no_sync() + ) + with sync_context: + step_losses = [] + hidden_losses = [] + velocity_losses = [] + intermediate_latents = [] + full_final_latent = None + if full_anchor_hidden is None: + raise RuntimeError(f"Missing Full step-0 anchor for chunk {chunk_id}") + anchor_hidden = full_anchor_hidden + offline_targets = ( + load_offline_ffff_targets( + offline_step_tensor_path(args.data_root, record, chunk_id), + timesteps=timesteps, + ) + if args.supervision_mode == "offline_ffff" + else None + ) + + for step_id in (1, 2, 3): + timestep_input = torch.full( + (4,), timesteps[step_id], device=device, dtype=timesteps.dtype + ) + model_input = torch.cat([chunk_latent, chunk_cond], dim=1) + if offline_targets is None: + target_velocity, target_hidden, frame_condition = teacher_denoise( + teacher, model_input.detach(), timestep_input, + chunk_viewmats, chunk_Ks, chunk_action, cache, len(selected), + ) + else: + offline_step = offline_targets[step_id] + target_velocity = offline_step["velocity"].to(device=device) + target_hidden = offline_step["hidden"].to(device=device) + frame_condition = offline_step["frame_condition"].to(device=device) + offline_noisy_sample = offline_step["noisy_sample"].to(device=device) + offline_latent_mse_by_chunk_step.append( + F.mse_loss( + chunk_latent.detach().float(), + offline_noisy_sample.float(), + ).detach() + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = predictor( + target_model_input=model_input, + anchor_hidden=anchor_hidden, + target_frame_condition=frame_condition, + target_viewmats=chunk_viewmats, + target_Ks=chunk_Ks, + rope_temporal_size=4 + len(selected), + start_rope_start_idx=len(selected), + context_kv={ + block_id: cache[block_id] + for block_id in predictor_model.predictor_config.source_block_ids + }, + previous_chunk_hidden=previous_hidden_by_step[step_id], + # ATC anchor distance: |t_target - t_anchor|, the anchor + # being the previous denoising step (Full at step 0). + anchor_distance=( + timesteps[step_id].float() - timesteps[step_id - 1].float() + ).abs().reshape(1), + ) + hidden_loss = F.mse_loss( + output["pred_hidden"].float(), target_hidden.float() + ) + velocity_loss = F.mse_loss( + output["pred_velocity"].float(), target_velocity.float() + ) + step_loss = ( + args.hidden_weight * hidden_loss + + args.velocity_weight * velocity_loss + ) + step_losses.append(step_loss) + hidden_losses.append(hidden_loss.detach()) + velocity_losses.append(velocity_loss.detach()) + if "gate_stats" in output: + gate_metric_sums += torch.stack( + [output["gate_stats"][name] for name in gate_stat_names] + ).to(dtype=torch.float64) + gate_metric_count += 1 + if "atc_stats" in output: + atc_metric_sums += torch.stack( + [ + output["atc_stats"][name].float() + if name in output["atc_stats"] + else torch.zeros((), device=device) + for name in atc_stat_names + ] + ).to(dtype=torch.float64) + atc_metric_count += 1 + if args.rollout_schedule == "fppf" and step_id == 3: + # P3 remains a fully supervised auxiliary branch, but an + # F-P-P-F trajectory must expose Full step-3 state to the + # next chunk. This mirrors inference: neither P3 hidden + # nor P3 velocity becomes rollout state. + current_hidden_by_step[step_id] = target_hidden.detach() + full_final_latent = differentiable_euler_step( + chunk_latent.detach(), + target_velocity.detach(), + sigmas, + step_id, + ).detach() + else: + current_hidden_by_step[step_id] = ( + output["pred_hidden"].detach() + if args.detach_between_timesteps + else output["pred_hidden"] + ) + anchor_hidden = ( + output["pred_hidden"].detach() + if args.detach_between_timesteps + else output["pred_hidden"] + ) + if not (args.rollout_schedule == "fppf" and step_id == 3): + if args.detach_between_timesteps: + chunk_latent = differentiable_euler_step( + chunk_latent.detach(), + output["pred_velocity"].detach(), + sigmas, + step_id, + ).detach() + else: + chunk_latent = differentiable_euler_step( + chunk_latent, output["pred_velocity"], sigmas, step_id + ) + if step_id in (1, 2) and not args.detach_between_timesteps: + chunk_latent.retain_grad() + intermediate_latents.append(chunk_latent) + if chunk_latent.requires_grad or anchor_hidden.requires_grad: + rollout_state_requires_grad_count += 1 + del output, target_velocity, target_hidden, frame_condition, model_input + if offline_targets is not None: + del offline_noisy_sample + + if args.rollout_schedule == "fppf": + if full_final_latent is None: + raise RuntimeError( + f"Missing Full step-3 rollout latent for chunk {chunk_id}" + ) + chunk_latent = full_final_latent + chunk_loss = torch.stack(step_losses).mean() + (chunk_loss / train_chunks / accumulation_steps).backward() + + if args.detach_between_timesteps: + latent_grad_values = [0.0, 0.0] + if rollout_state_requires_grad_count != 0: + raise RuntimeError( + "Detached timestep rollout unexpectedly retained a state graph" + ) + else: + latent_grad_values = [ + float(value.grad.float().abs().sum()) + if value.grad is not None else 0.0 + for value in intermediate_latents + ] + if len(latent_grad_values) != 2 or min(latent_grad_values) <= 0.0: + raise RuntimeError( + f"Latent BPTT is disconnected in chunk {chunk_id}: " + f"{latent_grad_values}" + ) + + metric_sums[0] += chunk_loss.detach() + trajectory_step_losses.extend(loss.detach() for loss in step_losses) + for index in range(3): + metric_sums[1 + index] += hidden_losses[index] + metric_sums[4 + index] += velocity_losses[index] + latents[:, :, start:end] = chunk_latent.detach() + previous_hidden_by_step = { + step: hidden.detach() for step, hidden in current_hidden_by_step.items() + } + del chunk_loss, step_losses, hidden_losses, velocity_losses + del current_hidden_by_step, anchor_hidden, chunk_latent, intermediate_latents + if offline_targets is not None: + del offline_targets + + if len(trajectory_step_losses) != train_chunks * 3: + raise RuntimeError( + f"Expected {train_chunks * 3} chunk/step losses, got " + f"{len(trajectory_step_losses)}" + ) + expected_offline_drift_records = ( + train_chunks * 3 if args.supervision_mode == "offline_ffff" else 0 + ) + if len(offline_latent_mse_by_chunk_step) != expected_offline_drift_records: + raise RuntimeError( + "Unexpected offline latent drift record count: " + f"{len(offline_latent_mse_by_chunk_step)} != " + f"{expected_offline_drift_records}" + ) + if is_last_micro_step: + grad_norm = torch.nn.utils.clip_grad_norm_(predictor_model.parameters(), args.grad_clip) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite grad norm: {float(grad_norm)}") + optimizer.step() + else: + grad_norm = torch.tensor(float("nan"), device=device) + metric_sums /= train_chunks + if gate_metric_count: + gate_metric_sums /= gate_metric_count + if atc_metric_count: + atc_metric_sums /= atc_metric_count + elapsed = time.perf_counter() - started + return { + "loss": float(metric_sums[0]), + "hidden_mse_step1": float(metric_sums[1]), + "hidden_mse_step2": float(metric_sums[2]), + "hidden_mse_step3": float(metric_sums[3]), + "velocity_mse_step1": float(metric_sums[4]), + "velocity_mse_step2": float(metric_sums[5]), + "velocity_mse_step3": float(metric_sums[6]), + "grad_norm": float(grad_norm), + "trajectory_time_s": elapsed, + "context_frames_last": float(len(selected)), + "latent_grad_l1_min": min(latent_grad_values), + "rollout_state_requires_grad_count": float(rollout_state_requires_grad_count), + "chunk_step_losses": [float(value) for value in trajectory_step_losses], + "offline_latent_mse_by_chunk_step": [ + float(value) for value in offline_latent_mse_by_chunk_step + ], + "gate_stats": [float(value) for value in gate_metric_sums], + "atc_stats": [float(value) for value in atc_metric_sums], + } + + +def main() -> None: + args = parse_args() + scheduler_validation = validate_differentiable_scheduler() + if args.validate_scheduler_only: + print(json.dumps(scheduler_validation, sort_keys=True)) + return + if not 2 <= args.num_chunks <= 8: + raise ValueError("num_chunks must be in [2, 8]") + if args.supervision_mode == "offline_ffff" and args.rollout_schedule != "fppp": + raise ValueError("offline_ffff supervision currently requires FPPP rollout") + if not 0.0 < args.gate_initial_probability < 1.0: + raise ValueError("gate_initial_probability must be in (0, 1)") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + raise RuntimeError("Launch rollout training with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + if world_size != args.expected_world_size: + raise ValueError( + f"Trajectory rollout training expects {args.expected_world_size} ranks " + f"(--expected_world_size), got {world_size}" + ) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="nccl", device_id=device) + initialize_parallel_state(sp=1) + initialize_runtime_state() + random.seed(args.seed + rank) + np.random.seed(args.seed + rank) + torch.manual_seed(args.seed + rank) + torch.cuda.manual_seed_all(args.seed + rank) + + data_root = args.data_root.resolve() + if args.text_kv_root is not None: + args.text_kv_root = args.text_kv_root.resolve() + cases_path = ( + args.cases_path.resolve() + if args.cases_path is not None + else data_root / "cases.jsonl" + ) + required_paths = [ + cases_path, + args.base_model.resolve(), + args.teacher_checkpoint.resolve(), + ] + if args.text_kv_root is not None: + required_paths.append(args.text_kv_root) + if args.resume is not None: + required_paths.append(args.resume.resolve()) + elif args.init_training_state is not None: + required_paths.append(args.init_training_state.resolve()) + else: + required_paths.append(args.init_weights.resolve()) + for path in required_paths: + if not path.exists(): + raise FileNotFoundError(path) + output_dir = args.output_dir.resolve() + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + dist.barrier() + + dataset = PredictorV4TrajectoryDataset( + cases_path, + tensor_root=data_root, + seed=args.latent_seed, + max_trajectories=args.max_trajectories, + ) + sampler = DistributedSampler( + dataset, num_replicas=world_size, rank=rank, + shuffle=True, seed=args.seed, drop_last=True, + ) + loader = DataLoader( + dataset, batch_size=1, sampler=sampler, drop_last=True, + num_workers=0, collate_fn=trajectory_collate, + ) + if len(loader) == 0: + raise ValueError("Trajectory loader has zero batches") + + pipe = HunyuanVideo_1_5_Pipeline.create_pipeline( + pretrained_model_name_or_path=str(args.base_model.resolve()), + transformer_version="480p_i2v", + enable_offloading=True, + enable_group_offloading=False, + create_sr_pipeline=False, + force_sparse_attn=False, + transformer_dtype=torch.bfloat16, + action_ckpt=str(args.teacher_checkpoint.resolve()), + ) + teacher = pipe.transformer + configure_exact_teacher(teacher) + teacher.eval().requires_grad_(False) + + config_kwargs = dict( + source_block_ids=args.source_blocks, + predictor_fusion_type=args.predictor_fusion_type, + gate_hidden_dim=args.gate_hidden_dim, + gate_initial_probability=args.gate_initial_probability, + ) + saved_predictor_config = None + config_source = args.resume if args.resume is not None else args.init_training_state + if args.predictor_arch != "baseline" and config_source is not None: + saved_predictor_config = torch.load( + config_source.resolve(), map_location="cpu", weights_only=False, mmap=True + ).get("predictor_config") + if args.predictor_arch == "atc" and ( + saved_predictor_config is None or saved_predictor_config.get("predictor_arch") != "atc" + ): + raise ValueError("--predictor_arch atc needs an ATC init/resume state with predictor_config") + if saved_predictor_config is not None and saved_predictor_config.get("predictor_arch") == "atc": + # Rebuild the exact Stage-1 ATC architecture. States written before the + # transport switch carry no atc_transport_mode and are local-window ATC. + atc_kwargs = { + key: value for key, value in saved_predictor_config.items() + if key.startswith("atc_") or key in ("predictor_arch", "previous_condition_mode") + } + atc_kwargs.setdefault("atc_transport_mode", "local") + config_kwargs.update(atc_kwargs) + if rank == 0: + print(json.dumps({"predictor_arch": "atc", **{k: v for k, v in atc_kwargs.items() if k.startswith("atc_")}}, sort_keys=True), flush=True) + model = HYWorldPlayPredictorV4(PredictorV2Config(**config_kwargs)) + model.load_teacher_initialization(args.teacher_checkpoint) + configure_model_precision( + model, + device=device, + fp32_trainable_params=args.fp32_trainable_params, + ) + if args.resume is None: + if args.init_training_state is not None: + load_initial_training_state( + model, + args.init_training_state.resolve(), + args.source_blocks, + ) + else: + load_predictor_weights(model, args.init_weights.resolve()) + model.enable_gradient_checkpointing(args.gradient_checkpointing) + model.train() + fusion_parameters = list(model.img_fusion.parameters()) + list(model.residual_out.parameters()) + block_parameters = [ + parameter for parameter in model.double_blocks.parameters() if parameter.requires_grad + ] + optimizer = AdamW( + [ + {"params": fusion_parameters, "lr": args.fusion_lr}, + {"params": block_parameters, "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), weight_decay=args.weight_decay, + ) + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=False, + ) + lr_scheduler = make_lr_scheduler(optimizer, args) + global_step = 0 + epoch = 0 + batch_in_epoch = 0 + if args.resume is not None: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + if tuple(state["predictor_config"]["source_block_ids"]) != args.source_blocks: + raise ValueError("Resume source blocks do not match") + saved_fp32_master = bool( + state.get("args", {}).get("fp32_trainable_params", False) + ) + if saved_fp32_master != args.fp32_trainable_params: + raise ValueError("Resume checkpoint FP32-master setting does not match") + saved_rollout_schedule = state.get("args", {}).get( + "rollout_schedule", "fppp" + ) + if saved_rollout_schedule != args.rollout_schedule: + raise ValueError( + "Resume checkpoint rollout schedule does not match: " + f"saved={saved_rollout_schedule}, requested={args.rollout_schedule}" + ) + saved_detach = bool( + state.get("args", {}).get("detach_between_timesteps", False) + ) + if saved_detach != args.detach_between_timesteps: + raise ValueError("Resume timestep-detach setting does not match") + saved_supervision_mode = state.get("args", {}).get( + "supervision_mode", "onpolicy" + ) + if saved_supervision_mode != args.supervision_mode: + raise ValueError( + "Resume supervision mode does not match: " + f"saved={saved_supervision_mode}, " + f"requested={args.supervision_mode}" + ) + saved_fusion_type = state.get("predictor_config", {}).get( + "predictor_fusion_type", + state.get("args", {}).get("predictor_fusion_type", "concat"), + ) + if saved_fusion_type != args.predictor_fusion_type: + raise ValueError("Resume checkpoint fusion type does not match") + result = model.load_state_dict(state["model"], strict=False) + missing_trainable = [ + key for key in result.missing_keys if key.startswith(TRAINABLE_PREFIXES) + ] + if result.unexpected_keys or missing_trainable: + raise RuntimeError( + "Resume model mismatch; " + f"unexpected={result.unexpected_keys}, missing={missing_trainable}" + ) + optimizer.load_state_dict(state["optimizer"]) + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + lr_scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + epoch = int(state.get("epoch", 0)) + batch_in_epoch = int(state.get("batch_in_epoch", 0)) + + predictor = DDP( + model, device_ids=[local_rank], output_device=local_rank, + broadcast_buffers=False, gradient_as_bucket_view=True, + find_unused_parameters=False, + ) + # Keep the Teacher resident during rollout; other condition encoders are + # temporarily onloaded one at a time by prepare_online_conditions. + teacher.to(device=device, dtype=torch.bfloat16) + points_local = generate_points_in_sphere(50000, 8.0).to(device) + torch.cuda.empty_cache() + + run_config = { + **vars(args), + "data_root": str(data_root), + "cases_path": str(cases_path), + "base_model": str(args.base_model.resolve()), + "teacher_checkpoint": str(args.teacher_checkpoint.resolve()), + "init_weights": str(args.init_weights.resolve()), + "init_training_state": ( + str(args.init_training_state.resolve()) + if args.init_training_state is not None + else None + ), + "output_dir": str(output_dir), + "source_blocks": list(args.source_blocks), + "text_kv_root": ( + str(args.text_kv_root) if args.text_kv_root is not None else None + ), + "world_size": world_size, + "trajectories": len(dataset), + "trajectories_per_rank_per_epoch": len(loader), + "gradient_accumulation_steps": args.gradient_accumulation_steps, + "global_trajectories_per_step": world_size * args.gradient_accumulation_steps, + "predictor_config": model.config_dict, + "schedule": ( + f"chunk0=FFFF,chunks1-{args.num_chunks - 1}=" + f"{args.rollout_schedule.upper()}" + ), + "trained_predictor_steps": [1, 2, 3], + "previous_chunk_hidden_sources": { + "step1": "predictor", + "step2": "predictor", + "step3": ( + "full_teacher" if args.rollout_schedule == "fppf" else "predictor" + ), + }, + "supervision_mode": args.supervision_mode, + "online_teacher_targets": args.supervision_mode == "onpolicy", + "offline_ffff_targets": args.supervision_mode == "offline_ffff", + "chunk0_source": ( + "offline_ffff_tensors" + if args.supervision_mode == "offline_ffff" + else "online_full_teacher" + ), + "teacher_denoise_calls_per_trajectory": ( + args.num_chunks - 1 + if args.supervision_mode == "offline_ffff" + else 4 + (args.num_chunks - 1) * 4 + ), + "online_text_kv": args.text_kv_root is None, + "precomputed_exact_text_kv": args.text_kv_root is not None, + "reuse_image_condition_latent": True, + "detach_between_chunks": True, + "detach_between_timesteps": args.detach_between_timesteps, + "timestep_gradient_graph": ( + "P1|P2|P3_stop_gradient" + if args.detach_between_timesteps + else "P1_to_P2_to_P3_full_bptt" + ), + "scheduler_validation": scheduler_validation, + "trainable_parameters": sum(p.numel() for p in model.parameters() if p.requires_grad), + "dtype": "bfloat16", + "activation_dtype": "bfloat16", + "teacher_dtype": "bfloat16", + "trainable_parameter_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + "optimizer_moment_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + "inference_checkpoint_dtype": "bfloat16", + "gpu": torch.cuda.get_device_name(device), + } + if rank == 0: + config_path = output_dir / "train_config.json" + config_path.write_text( + json.dumps(run_config, indent=2, ensure_ascii=False, default=str, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(run_config, indent=2, ensure_ascii=False, default=str), flush=True) + + swanlab_run = None + if rank == 0: + if args.swanlab_mode == "cloud": + swanlab.login(api_key=os.environ.get("SWANLAB_API_KEY"), save=False) + swanlab_run = swanlab.init( + project=args.swanlab_project, + workspace=args.swanlab_workspace, + experiment_name=args.swanlab_experiment, + description=( + f"v4 {args.supervision_mode} supervision recursive " + f"{args.rollout_schedule.upper()} self-rollout; 8-GPU DDP; " + "Predictor steps 1/2/3 all supervised; " + f"timestep_detach={args.detach_between_timesteps}; chunk_detach=True" + ), + tags=[ + "predictor-v4", + args.supervision_mode, + args.rollout_schedule.upper(), + ( + "detached-timestep" + if args.detach_between_timesteps else "timestep-bptt" + ), + "detached-chunk", + "8gpu", + f"blocks-{args.source_blocks[0]}-{args.source_blocks[1]}", + ], + config=run_config, + logdir=str(output_dir / "swanlab"), + mode=args.swanlab_mode, + ) + + log_path = output_dir / "train_log.jsonl" + try: + resume_epoch = epoch + resume_batch = batch_in_epoch + micro_metrics: list[dict] = [] + while global_step < args.max_steps: + sampler.set_epoch(epoch) + # A resumed epoch restarts at a group boundary (states are only saved + # after a full group); an incomplete trailing group is dropped. + micro_metrics = [] + optimizer.zero_grad(set_to_none=True) + for loader_index, record in enumerate(loader): + if epoch == resume_epoch and loader_index < resume_batch: + continue + if global_step >= args.max_steps: + break + torch.cuda.reset_peak_memory_stats(device) + trajectory_key = torch.tensor( + int(record["case_id"]) * 4 + int(record["action_id"]), + device=device, + dtype=torch.int64, + ) + gathered_keys = [torch.zeros_like(trajectory_key) for _ in range(world_size)] + dist.all_gather(gathered_keys, trajectory_key) + key_values = [int(value) for value in gathered_keys] + if len(set(key_values)) != world_size: + raise RuntimeError(f"Duplicate trajectories across ranks: {key_values}") + accumulation_index = loader_index % args.gradient_accumulation_steps + metrics = run_trajectory( + pipe=pipe, teacher=teacher, predictor=predictor, + predictor_model=model, optimizer=optimizer, record=record, + points_local=points_local, args=args, device=device, + accumulation_index=accumulation_index, + accumulation_steps=args.gradient_accumulation_steps, + ) + micro_metrics.append(metrics) + if accumulation_index + 1 < args.gradient_accumulation_steps: + continue # gradients kept; optimizer steps on the last micro-batch + metrics = _average_micro_metrics(micro_metrics) + micro_metrics = [] + if global_step == 0: + validate_optimizer_precision( + optimizer, + fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + lr_scheduler.step() + global_step += 1 + batch_in_epoch = loader_index + 1 + + values = torch.tensor( + [ + metrics["loss"], metrics["hidden_mse_step1"], + metrics["hidden_mse_step2"], metrics["hidden_mse_step3"], + metrics["velocity_mse_step1"], metrics["velocity_mse_step2"], + metrics["velocity_mse_step3"], metrics["trajectory_time_s"], + metrics["latent_grad_l1_min"], + metrics["rollout_state_requires_grad_count"], + *metrics["chunk_step_losses"], + *metrics["offline_latent_mse_by_chunk_step"], + *metrics["gate_stats"], + *metrics["atc_stats"], + ], device=device, dtype=torch.float64, + ) + averaged = distributed_average(values) + grad_max = torch.tensor(metrics["grad_norm"], device=device) + dist.all_reduce(grad_max, op=dist.ReduceOp.MAX) + if rank == 0 and (global_step % args.log_every == 0 or global_step == 1): + record_log = { + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + "loss": float(averaged[0]), + "hidden_mse_step1": float(averaged[1]), + "hidden_mse_step2": float(averaged[2]), + "hidden_mse_step3": float(averaged[3]), + "velocity_mse_step1": float(averaged[4]), + "velocity_mse_step2": float(averaged[5]), + "velocity_mse_step3": float(averaged[6]), + "trajectory_time_s": float(averaged[7]), + "latent_grad_l1_min": float(averaged[8]), + "rollout_state_requires_grad_count": float(averaged[9]), + "grad_norm_max": float(grad_max), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated(device) / 2**30, + } + is_atc = model.predictor_config.predictor_arch == "atc" + if is_atc: + for atc_offset, atc_name in enumerate( + ( + "delta_h_d_norm", + "gated_delta_h_c_norm", + "transport_attention_entropy", + "transport_argmax_displacement", + ) + ): + record_log[atc_name] = float( + averaged[len(averaged) - 4 + atc_offset] + ) + if is_atc or args.predictor_fusion_type in { + "token_gate_prev", + "concat_residual_gate_prev", + }: + for gate_offset, gate_name in enumerate( + ( + "gate_p_mean", + "gate_p_std", + "gate_p_p10", + "gate_p_p50", + "gate_p_p90", + "gate_p_near_zero", + "gate_p_near_one", + ) + ): + record_log[gate_name] = float( + averaged[len(averaged) - 11 + gate_offset] + ) + for chunk_offset, chunk_id in enumerate(range(1, args.num_chunks)): + for step_offset, step_id in enumerate((1, 2, 3)): + flat_index = chunk_offset * 3 + step_offset + record_log[ + f"loss_chunk{chunk_id}_step{step_id}" + ] = float(averaged[10 + flat_index]) + if args.supervision_mode == "offline_ffff": + drift_offset = 10 + (args.num_chunks - 1) * 3 + record_log[ + f"offline_latent_mse_chunk{chunk_id}_step{step_id}" + ] = float(averaged[drift_offset + flat_index]) + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record_log, sort_keys=True) + "\n") + swanlab.log(record_log, step=global_step) + print(json.dumps(record_log, sort_keys=True), flush=True) + + should_save = global_step % args.save_every == 0 or global_step == args.max_steps + if should_save: + dist.barrier() + if rank == 0: + weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" + save_predictor_weights( + model, + weights_path, + floating_dtype=torch.bfloat16, + ) + save_training_state( + output_dir / "training_latest.pt", + { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "scheduler": lr_scheduler.state_dict(), + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + "args": vars(args), + "predictor_config": model.config_dict, + "weights_path": str(weights_path), + }, + ) + prune_snapshots(output_dir, args.keep_snapshots) + dist.barrier() + del metrics, values, averaged, grad_max + gc.collect() + epoch += 1 + batch_in_epoch = 0 + resume_batch = 0 + except BaseException as error: + if rank == 0 and swanlab_run is not None: + try: + swanlab.finish(error=str(error)) + except Exception: + pass + raise + else: + if rank == 0 and swanlab_run is not None: + swanlab.finish() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/train_predictor_v4_timestep_rollout.py b/train_predictor_v4_timestep_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..40fa4d926fa2356ade73b40a0b2462f748cc0d7a --- /dev/null +++ b/train_predictor_v4_timestep_rollout.py @@ -0,0 +1,1080 @@ +#!/usr/bin/env python3 +"""Eight-GPU v4 training with offline Full chunks and on-policy timestep rollout.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import random +import time +from contextlib import nullcontext +from pathlib import Path + +import swanlab +import torch +import torch.distributed as dist +import torch.nn.functional as F +from safetensors.torch import load_file +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW + +from hyvideo.commons.parallel_states import initialize_parallel_state +from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline +from models.predictor_v2 import PredictorV2Config +from models.predictor_v4 import HYWorldPlayPredictorV4 +from predictor_data.prefeature_schema import DEFAULT_ACTIONS +from predictor_data.text_kv_schema import ( + TEXT_KV_HEADS, + TEXT_KV_HEAD_DIM, + TEXT_KV_LAYER_COUNT, + TEXT_KV_TOKENS, + expected_text_kv_keys, + text_kv_key, +) +from predictor_training.checkpoint import ( + save_predictor_weights, + save_training_state, + trainable_state_dict, +) +from predictor_training.dataset_v2 import _load_case_file +from predictor_training.rollout_utils import configure_exact_teacher, pose_to_input +from predictor_training.timestep_rollout_dataset import ( + load_case_ids, + load_timestep_rollout_record, + load_timestep_rollout_records, + move_timestep_rollout_record, +) +from train_predictor_v4_rollout import ( + DEFAULT_BASE, + DEFAULT_TEACHER, + configure_model_precision, + differentiable_euler_step, + distributed_average, + history_prefill, + initialize_runtime_state, + load_initial_training_state, + make_lr_scheduler, + prune_snapshots, + teacher_denoise, + text_prefill, + validate_differentiable_scheduler, + validate_optimizer_precision, +) + + +DEFAULT_DATA_ROOT = Path("/mnt/local_nvme/zoubin/cz/hyworldplay_predictor_prefeature_v4") +DEFAULT_CASES = Path("configs/rollout_overfit8_seed20260721_cases.jsonl") +DEFAULT_SIDECAR = Path( + "/mnt/local_nvme/zoubin/cz/" + "hyworldplay_predictor_timestep_rollout_overfit8_sidecar" +) +DEFAULT_TEXT_KV_ROOT = DEFAULT_DATA_ROOT / "text_kv_exact_all54_bf16" +DEFAULT_INIT_STATE = Path( + "checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/" + "training_latest.pt" +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--training_mode", + choices=("timestep_rollout", "grouped_teacher_forcing"), + default="timestep_rollout", + help=( + "Grouped teacher forcing trains the three offline Full-to-Predictor " + "step pairs together without feeding Predictor outputs across steps." + ), + ) + parser.add_argument("--data_root", type=Path, default=DEFAULT_DATA_ROOT) + parser.add_argument("--cases_path", type=Path, default=DEFAULT_CASES) + parser.add_argument("--sidecar_root", type=Path, default=DEFAULT_SIDECAR) + parser.add_argument("--text_kv_root", type=Path, default=DEFAULT_TEXT_KV_ROOT) + parser.add_argument("--base_model", type=Path, default=DEFAULT_BASE) + parser.add_argument("--teacher_checkpoint", type=Path, default=DEFAULT_TEACHER) + parser.add_argument("--init_training_state", type=Path, default=DEFAULT_INIT_STATE) + parser.add_argument("--output_dir", type=Path, required=True) + parser.add_argument("--max_steps", type=int, default=200) + parser.add_argument( + "--stop_after_steps", + type=int, + default=None, + help=( + "Optional earlier stopping step while retaining max_steps as the LR " + "scheduler horizon; useful for reconstructing an intermediate checkpoint." + ), + ) + parser.add_argument("--replay_weight", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=20260721) + parser.add_argument("--fusion_lr", type=float, default=1e-4) + parser.add_argument("--blocks_lr", type=float, default=1e-5) + parser.add_argument("--weight_decay", type=float, default=0.01) + parser.add_argument("--scheduler_warmup_steps", type=int, default=8) + parser.add_argument("--block_warmup_steps", type=int, default=8) + parser.add_argument( + "--block_warmup_mode", + choices=("delayed", "linear"), + default="delayed", + help=( + "Delayed keeps the legacy zero-LR block freeze. Linear warms the " + "block group from zero to its configured LR." + ), + ) + parser.add_argument( + "--batch_size", + type=int, + default=1, + help=( + "Physical records per GPU. Exact per-record Teacher history caches " + "currently require batch size 1." + ), + ) + parser.add_argument("--hidden_weight", type=float, default=0.1) + parser.add_argument("--velocity_weight", type=float, default=1.0) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument( + "--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--fp32_trainable_params", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--detach_between_timesteps", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Use Predictor outputs to form the next timestep state, but stop " + "gradients through both the hidden-anchor and scheduler paths." + ), + ) + parser.add_argument("--save_every", type=int, default=100) + parser.add_argument("--keep_snapshots", type=int, default=3) + parser.add_argument("--log_every", type=int, default=1) + parser.add_argument("--resume", type=Path, default=None) + parser.add_argument("--swanlab_project", default="HY-WorldPlay-Predictor-v4") + parser.add_argument("--swanlab_workspace", default=None) + parser.add_argument("--swanlab_experiment", required=True) + parser.add_argument( + "--swanlab_mode", choices=("cloud", "offline", "local", "disabled"), + default="cloud", + ) + parser.add_argument("--validate_scheduler_only", action="store_true") + return parser.parse_args() + + +def build_full_condition(case_tensor_path: Path, device: torch.device) -> torch.Tensor: + image_condition = _load_case_file(str(case_tensor_path))["image_condition_latent"] + condition = torch.zeros( + (1, 33, 32, 30, 52), device=device, dtype=torch.bfloat16 + ) + condition[:, :32, 0:1] = image_condition.to(device=device, dtype=torch.bfloat16) + condition[:, 32:, 0:1] = 1 + return condition + + +def load_text_cache( + teacher: torch.nn.Module, + condition_path: Path, + device: torch.device, +) -> list[dict]: + values = load_file(str(condition_path), device="cpu") + prompt_embeds = values["prompt_embeds"].to(device=device) + prompt_mask = values["prompt_mask"].to(device=device) + vision_states = values["vision_states"].to(device=device) + extra_kwargs = { + "byt5_text_states": values["byt5_text_states"].to(device=device), + "byt5_text_mask": values["byt5_text_mask"].to(device=device), + } + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + cache = text_prefill( + teacher, prompt_embeds, prompt_mask, vision_states, + extra_kwargs, device, torch.bfloat16, + ) + del values, prompt_embeds, prompt_mask, vision_states, extra_kwargs + return cache + + +def load_precomputed_text_cache( + cache_path: Path, + device: torch.device, +) -> list[dict[str, torch.Tensor]]: + """Load exact all-layer Text K/V and remove on-disk padding before use.""" + values = load_file(str(cache_path), device="cpu") + missing = expected_text_kv_keys().difference(values) + unexpected = set(values).difference(expected_text_kv_keys()) + if missing or unexpected: + raise ValueError( + f"Invalid Text K/V cache {cache_path}: missing={sorted(missing)}, " + f"unexpected={sorted(unexpected)}" + ) + mask = values["text_valid_mask"] + if tuple(mask.shape) != (1, TEXT_KV_TOKENS) or mask.dtype != torch.bool: + raise ValueError(f"Invalid text_valid_mask in {cache_path}") + valid_tokens = int(mask.sum().item()) + expected_mask = torch.arange(TEXT_KV_TOKENS)[None] < valid_tokens + if not torch.equal(mask, expected_mask): + raise ValueError(f"Non-prefix text_valid_mask in {cache_path}") + expected_shape = (1, TEXT_KV_HEADS, TEXT_KV_TOKENS, TEXT_KV_HEAD_DIM) + cache: list[dict[str, torch.Tensor]] = [] + for block_id in range(TEXT_KV_LAYER_COUNT): + layer = {} + for kind in ("k_txt", "v_txt"): + name = text_kv_key(block_id, kind) + value = values[name] + if tuple(value.shape) != expected_shape or value.dtype != torch.bfloat16: + raise ValueError( + f"Invalid {name} in {cache_path}: " + f"shape={tuple(value.shape)}, dtype={value.dtype}" + ) + layer[kind] = value[:, :, :valid_tokens].to( + device=device, dtype=torch.bfloat16, non_blocking=True + ) + cache.append(layer) + del values + return cache + + +def loss_pair( + output: dict[str, torch.Tensor], + target_hidden: torch.Tensor, + target_velocity: torch.Tensor, + args: argparse.Namespace, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden = F.mse_loss(output["pred_hidden"].float(), target_hidden.float()) + velocity = F.mse_loss(output["pred_velocity"].float(), target_velocity.float()) + return args.hidden_weight * hidden + args.velocity_weight * velocity, hidden, velocity + + +def predictor_inputs( + *, + model_input: torch.Tensor, + anchor_hidden: torch.Tensor, + frame_condition: torch.Tensor, + previous_hidden: torch.Tensor, + item: dict, + cache: list[dict], + source_blocks: tuple[int, int], +) -> dict: + return { + "target_model_input": model_input, + "anchor_hidden": anchor_hidden, + "target_frame_condition": frame_condition, + "target_viewmats": item["target_viewmats"], + "target_Ks": item["target_Ks"], + "rope_temporal_size": item["rope_temporal_size"], + "start_rope_start_idx": item["start_rope_start_idx"], + "context_kv": {block_id: cache[block_id] for block_id in source_blocks}, + "previous_chunk_hidden": previous_hidden, + } + + +def assert_exact_scheduler_timesteps( + saved_timesteps: dict[int, torch.Tensor], + scheduler_timesteps: torch.Tensor, +) -> None: + """Require the offline record to use the current scheduler's exact FP32 grid.""" + if scheduler_timesteps.dtype != torch.float32: + raise TypeError( + "Current scheduler timesteps must be FP32, got " + f"{scheduler_timesteps.dtype}" + ) + expected_steps = tuple(range(int(scheduler_timesteps.numel()))) + if tuple(sorted(saved_timesteps)) != expected_steps: + raise ValueError( + f"Saved timestep keys must be {expected_steps}, got " + f"{tuple(sorted(saved_timesteps))}" + ) + values: list[torch.Tensor] = [] + for step_id in expected_steps: + value = saved_timesteps[step_id] + if value.dtype != torch.float32: + raise TypeError( + f"Saved timestep {step_id} must be FP32, got {value.dtype}" + ) + if value.numel() != 1: + raise ValueError( + f"Saved timestep {step_id} must be scalar-like, got {value.shape}" + ) + values.append(value.reshape(1)) + saved = torch.cat(values) + current = scheduler_timesteps.reshape(-1).to(device=saved.device) + if not torch.equal(saved, current): + differences = (saved - current).abs() + raise RuntimeError( + "Saved timesteps do not exactly match the current scheduler: " + f"saved={saved.detach().cpu().tolist()}, " + f"current={current.detach().cpu().tolist()}, " + f"max_abs_error={float(differences.max())}" + ) + + +def load_records_by_case( + manifest_path: Path, + case_ids: list[int], +) -> tuple[Path, dict[int, list[dict]]]: + """Load the 4-action x 7-chunk records for every selected training case.""" + record_root: Path | None = None + records_by_case: dict[int, list[dict]] = {} + for case_id in case_ids: + current_root, records = load_timestep_rollout_records( + manifest_path, case_id=case_id + ) + if record_root is None: + record_root = current_root + elif current_root != record_root: + raise RuntimeError("All timestep rollout records must share one root") + records_by_case[case_id] = records + if record_root is None: + raise ValueError("No timestep rollout records were selected") + return record_root, records_by_case + + +def rank_epoch_records( + records_by_case: dict[int, list[dict]], + *, + seed: int, + epoch: int, + rank: int, + world_size: int, +) -> list[tuple[int, dict]]: + """Case-grouped global shuffle followed by equal contiguous rank shards.""" + case_ids = sorted(records_by_case) + random.Random(seed + epoch * 100003).shuffle(case_ids) + global_records: list[tuple[int, dict]] = [] + for case_id in case_ids: + case_records = list(records_by_case[case_id]) + random.Random(seed + epoch * 100003 + case_id * 1009).shuffle(case_records) + global_records.extend((case_id, record) for record in case_records) + if len(global_records) % world_size != 0: + raise ValueError( + f"Record count {len(global_records)} is not divisible by {world_size} ranks" + ) + per_rank = len(global_records) // world_size + start = rank * per_rank + return global_records[start : start + per_rank] + + +def next_timestep_state( + *, + pred_hidden: torch.Tensor, + chunk_latent: torch.Tensor, + pred_velocity: torch.Tensor, + sigmas: torch.Tensor, + step_id: int, + detach_between_timesteps: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the next self-rollout state, optionally truncating timestep BPTT.""" + if detach_between_timesteps: + next_anchor = pred_hidden.detach() + with torch.no_grad(): + next_latent = differentiable_euler_step( + chunk_latent.detach(), pred_velocity.detach(), sigmas, step_id + ) + if next_anchor.requires_grad or next_latent.requires_grad: + raise RuntimeError("Detached timestep state unexpectedly requires gradients") + return next_anchor, next_latent + return pred_hidden, differentiable_euler_step( + chunk_latent, pred_velocity, sigmas, step_id + ) + + +def train_record( + *, + teacher: torch.nn.Module, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + optimizer: torch.optim.Optimizer, + item: dict, + text_cache: list[dict], + history_latent: torch.Tensor, + full_condition: torch.Tensor, + full_viewmats: torch.Tensor, + full_Ks: torch.Tensor, + full_action: torch.Tensor, + sigmas: torch.Tensor, + scheduler_timesteps: torch.Tensor, + args: argparse.Namespace, + device: torch.device, +) -> dict[str, float]: + started = time.perf_counter() + assert_exact_scheduler_timesteps(item["timesteps"], scheduler_timesteps) + selected = [int(value) for value in item["selected_context_indices"].cpu().tolist()] + if len(selected) != item["context_frames"]: + raise ValueError("Stored Context indices do not match manifest context_frames") + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + cache = history_prefill( + teacher, text_cache, history_latent, full_condition, selected, + full_viewmats, full_Ks, full_action, torch.float32, + ) + + optimizer.zero_grad(set_to_none=True) + source_blocks = predictor_model.predictor_config.source_block_ids + chunk_latent = item["step1_model_input"][:, :32] + chunk_condition = item["step1_model_input"][:, 32:] + anchor_hidden = item["step0_anchor_hidden"] + rollout_losses: list[torch.Tensor] = [] + hidden_losses: list[torch.Tensor] = [] + velocity_losses: list[torch.Tensor] = [] + intermediate_latents: list[torch.Tensor] = [] + + # When replay is enabled, defer the one DDP synchronization to its backward. + rollout_sync = predictor.no_sync() if args.replay_weight > 0 else nullcontext() + with rollout_sync: + for step_id in (1, 2, 3): + if args.detach_between_timesteps and step_id > 1: + if anchor_hidden.requires_grad or chunk_latent.requires_grad: + raise RuntimeError( + f"Step {step_id} received a gradient-connected prior state" + ) + model_input = ( + item["step1_model_input"] + if step_id == 1 + else torch.cat([chunk_latent, chunk_condition], dim=1) + ) + if step_id == 1: + target_hidden = item["step1_target_hidden"] + target_velocity = item["step1_target_velocity"] + frame_condition = item["offline_frame_conditions"][1] + else: + timestep_input = item["timesteps"][step_id].reshape(1).expand(4) + target_velocity, target_hidden, frame_condition = teacher_denoise( + teacher, model_input.detach(), timestep_input, + item["target_viewmats"], item["target_Ks"], + full_action[:, item["chunk_id"] * 4 : item["chunk_id"] * 4 + 4], + cache, len(selected), + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = predictor( + **predictor_inputs( + model_input=model_input, + anchor_hidden=anchor_hidden, + frame_condition=frame_condition, + previous_hidden=item["previous_chunk_hidden"][step_id], + item=item, + cache=cache, + source_blocks=source_blocks, + ) + ) + step_loss, hidden_loss, velocity_loss = loss_pair( + output, target_hidden, target_velocity, args + ) + rollout_losses.append(step_loss) + hidden_losses.append(hidden_loss.detach()) + velocity_losses.append(velocity_loss.detach()) + if step_id < 3: + anchor_hidden, chunk_latent = next_timestep_state( + pred_hidden=output["pred_hidden"], + chunk_latent=chunk_latent, + pred_velocity=output["pred_velocity"], + sigmas=sigmas, + step_id=step_id, + detach_between_timesteps=args.detach_between_timesteps, + ) + if not args.detach_between_timesteps: + chunk_latent.retain_grad() + intermediate_latents.append(chunk_latent) + del output, target_hidden, target_velocity, frame_condition, model_input + rollout_loss = torch.stack(rollout_losses).mean() + rollout_loss.backward() + + if args.detach_between_timesteps: + cross_timestep_grad_l1 = 0.0 + else: + latent_grad_values = [ + float(value.grad.float().abs().sum()) if value.grad is not None else 0.0 + for value in intermediate_latents + ] + if len(latent_grad_values) != 2 or min(latent_grad_values) <= 0: + raise RuntimeError(f"Timestep BPTT is disconnected: {latent_grad_values}") + cross_timestep_grad_l1 = min(latent_grad_values) + + replay_loss_value = 0.0 + replay_hidden_value = 0.0 + replay_velocity_value = 0.0 + if args.replay_weight > 0: + replay = item["replay"] + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + replay_output = predictor( + **predictor_inputs( + model_input=replay["model_input"], + anchor_hidden=replay["anchor_hidden"], + frame_condition=replay["frame_condition"], + previous_hidden=replay["previous_chunk_hidden"], + item=item, + cache=cache, + source_blocks=source_blocks, + ) + ) + replay_loss, replay_hidden, replay_velocity = loss_pair( + replay_output, replay["target_hidden"], replay["target_velocity"], args + ) + (args.replay_weight * replay_loss).backward() + replay_loss_value = float(replay_loss.detach()) + replay_hidden_value = float(replay_hidden.detach()) + replay_velocity_value = float(replay_velocity.detach()) + del replay_output, replay_loss, replay_hidden, replay_velocity + + grad_norm = torch.nn.utils.clip_grad_norm_(predictor_model.parameters(), args.grad_clip) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite grad norm: {float(grad_norm)}") + optimizer.step() + total_loss = float(rollout_loss.detach()) + args.replay_weight * replay_loss_value + result = { + "loss": total_loss, + "rollout_loss": float(rollout_loss.detach()), + "replay_loss": replay_loss_value, + "replay_hidden_mse": replay_hidden_value, + "replay_velocity_mse": replay_velocity_value, + "hidden_mse_step1": float(hidden_losses[0]), + "hidden_mse_step2": float(hidden_losses[1]), + "hidden_mse_step3": float(hidden_losses[2]), + "velocity_mse_step1": float(velocity_losses[0]), + "velocity_mse_step2": float(velocity_losses[1]), + "velocity_mse_step3": float(velocity_losses[2]), + "cross_timestep_grad_l1": cross_timestep_grad_l1, + "grad_norm": float(grad_norm), + "record_time_s": time.perf_counter() - started, + } + del cache, rollout_loss, rollout_losses, hidden_losses, velocity_losses + del intermediate_latents, chunk_latent, anchor_hidden + return result + + +def train_grouped_teacher_forcing_record( + *, + teacher: torch.nn.Module, + predictor: DDP, + predictor_model: HYWorldPlayPredictorV4, + optimizer: torch.optim.Optimizer, + item: dict, + text_cache: list[dict], + history_latent: torch.Tensor, + full_condition: torch.Tensor, + full_viewmats: torch.Tensor, + full_Ks: torch.Tensor, + full_action: torch.Tensor, + scheduler_timesteps: torch.Tensor, + args: argparse.Namespace, + device: torch.device, +) -> dict[str, float]: + """Train steps 1/2/3 as one group of independent offline F-to-P pairs.""" + started = time.perf_counter() + assert_exact_scheduler_timesteps(item["timesteps"], scheduler_timesteps) + if "teacher_forcing" not in item: + raise KeyError("Grouped teacher forcing tensors were not loaded") + selected = [int(value) for value in item["selected_context_indices"].cpu().tolist()] + if len(selected) != item["context_frames"]: + raise ValueError("Stored Context indices do not match manifest context_frames") + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + cache = history_prefill( + teacher, text_cache, history_latent, full_condition, selected, + full_viewmats, full_Ks, full_action, torch.float32, + ) + + optimizer.zero_grad(set_to_none=True) + source_blocks = predictor_model.predictor_config.source_block_ids + losses: list[torch.Tensor] = [] + hidden_losses: list[torch.Tensor] = [] + velocity_losses: list[torch.Tensor] = [] + for step_id in (1, 2, 3): + pair = item["teacher_forcing"][step_id] + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = predictor( + **predictor_inputs( + model_input=pair["model_input"], + anchor_hidden=pair["anchor_hidden"], + frame_condition=pair["frame_condition"], + previous_hidden=pair["previous_chunk_hidden"], + item=item, + cache=cache, + source_blocks=source_blocks, + ) + ) + step_loss, hidden_loss, velocity_loss = loss_pair( + output, pair["target_hidden"], pair["target_velocity"], args + ) + losses.append(step_loss) + hidden_losses.append(hidden_loss.detach()) + velocity_losses.append(velocity_loss.detach()) + del output, hidden_loss, velocity_loss + + objective_loss = torch.stack(losses).mean() + objective_loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(predictor_model.parameters(), args.grad_clip) + if not torch.isfinite(grad_norm): + raise FloatingPointError(f"Non-finite grad norm: {float(grad_norm)}") + optimizer.step() + result = { + "loss": float(objective_loss.detach()), + # Preserve the shared logging schema; this field is the grouped objective here. + "rollout_loss": float(objective_loss.detach()), + "replay_loss": 0.0, + "replay_hidden_mse": 0.0, + "replay_velocity_mse": 0.0, + "hidden_mse_step1": float(hidden_losses[0]), + "hidden_mse_step2": float(hidden_losses[1]), + "hidden_mse_step3": float(hidden_losses[2]), + "velocity_mse_step1": float(velocity_losses[0]), + "velocity_mse_step2": float(velocity_losses[1]), + "velocity_mse_step3": float(velocity_losses[2]), + "cross_timestep_grad_l1": 0.0, + "grad_norm": float(grad_norm), + "record_time_s": time.perf_counter() - started, + } + del cache, losses, hidden_losses, velocity_losses, objective_loss + return result + + +def main() -> None: + args = parse_args() + scheduler_validation = validate_differentiable_scheduler() + if args.validate_scheduler_only: + print(json.dumps(scheduler_validation, sort_keys=True)) + return + if args.replay_weight < 0: + raise ValueError("replay_weight must be non-negative") + if args.batch_size != 1: + raise ValueError( + "Physical batch_size > 1 is unsupported because each record owns an " + "exact variable-length 54-layer Teacher history cache; use batch_size=1" + ) + if args.training_mode == "grouped_teacher_forcing": + if args.replay_weight != 0: + raise ValueError("Grouped teacher forcing does not use replay loss") + if args.detach_between_timesteps: + raise ValueError( + "Grouped teacher forcing has no cross-timestep graph; " + "detach_between_timesteps must be disabled" + ) + stop_after_steps = ( + args.max_steps if args.stop_after_steps is None else args.stop_after_steps + ) + if not 0 < stop_after_steps <= args.max_steps: + raise ValueError( + "stop_after_steps must satisfy 0 < stop_after_steps <= max_steps, got " + f"{stop_after_steps} and {args.max_steps}" + ) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ): + raise RuntimeError("Launch timestep rollout training with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + if world_size != 8: + raise ValueError(f"Timestep rollout training requires 8 ranks, got {world_size}") + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="nccl", device_id=device) + initialize_parallel_state(sp=1) + initialize_runtime_state() + random.seed(args.seed + rank) + torch.manual_seed(args.seed + rank) + torch.cuda.manual_seed_all(args.seed + rank) + + data_root = args.data_root.resolve() + sidecar_root = args.sidecar_root.resolve() + text_kv_root = args.text_kv_root.resolve() + cases_path = args.cases_path.resolve() + case_ids = load_case_ids(cases_path) + manifest_path = data_root / "train_eval_manifest.jsonl" + record_root, records_by_case = load_records_by_case(manifest_path, case_ids) + total_records = sum(len(items) for items in records_by_case.values()) + if total_records % world_size != 0: + raise ValueError( + f"Selected record count {total_records} is not divisible by {world_size}" + ) + records_per_rank = total_records // world_size + required = [ + data_root, manifest_path, cases_path, sidecar_root, text_kv_root, + args.base_model.resolve(), args.teacher_checkpoint.resolve(), + args.init_training_state.resolve(), + ] + if args.resume is not None: + required.append(args.resume.resolve()) + for path in required: + if not path.exists(): + raise FileNotFoundError(path) + + output_dir = args.output_dir.resolve() + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + dist.barrier() + + pipe = HunyuanVideo_1_5_Pipeline.create_pipeline( + pretrained_model_name_or_path=str(args.base_model.resolve()), + transformer_version="480p_i2v", + enable_offloading=True, + enable_group_offloading=False, + create_sr_pipeline=False, + force_sparse_attn=False, + transformer_dtype=torch.bfloat16, + action_ckpt=str(args.teacher_checkpoint.resolve()), + ) + teacher = pipe.transformer + configure_exact_teacher(teacher) + teacher.eval().requires_grad_(False).to(device=device, dtype=torch.bfloat16) + + predictor_model = HYWorldPlayPredictorV4(PredictorV2Config(source_block_ids=(1, 52))) + predictor_model.load_teacher_initialization(args.teacher_checkpoint.resolve()) + configure_model_precision( + predictor_model, device=device, + fp32_trainable_params=args.fp32_trainable_params, + ) + if args.resume is None: + load_initial_training_state( + predictor_model, args.init_training_state.resolve(), (1, 52) + ) + predictor_model.enable_gradient_checkpointing(args.gradient_checkpointing) + predictor_model.train() + fusion_parameters = list(predictor_model.img_fusion.parameters()) + list( + predictor_model.residual_out.parameters() + ) + block_parameters = [ + parameter for parameter in predictor_model.double_blocks.parameters() + if parameter.requires_grad + ] + optimizer = AdamW( + [ + {"params": fusion_parameters, "lr": args.fusion_lr}, + {"params": block_parameters, "lr": args.blocks_lr}, + ], + betas=(0.9, 0.95), weight_decay=args.weight_decay, + ) + lr_scheduler = make_lr_scheduler(optimizer, args) + global_step = 0 + epoch = 0 + batch_in_epoch = 0 + if args.resume is not None: + state = torch.load(args.resume, map_location="cpu", weights_only=False) + saved_args = state.get("args", {}) + if saved_args.get("training_mode", "timestep_rollout") != args.training_mode: + raise ValueError("Resume training_mode does not match") + if float(saved_args.get("replay_weight", -1)) != args.replay_weight: + raise ValueError("Resume replay_weight does not match") + if bool(saved_args.get("fp32_trainable_params", False)) != args.fp32_trainable_params: + raise ValueError("Resume precision mode does not match") + predictor_model.load_state_dict(state["model"], strict=False) + optimizer.load_state_dict(state["optimizer"]) + lr_scheduler.load_state_dict(state["scheduler"]) + global_step = int(state["global_step"]) + epoch = int(state.get("epoch", 0)) + batch_in_epoch = int(state.get("batch_in_epoch", 0)) + validate_optimizer_precision( + optimizer, fp32_trainable_params=args.fp32_trainable_params, + require_state=global_step > 0, + ) + predictor = DDP( + predictor_model, device_ids=[local_rank], output_device=local_rank, + broadcast_buffers=False, gradient_as_bucket_view=True, + find_unused_parameters=False, + ) + + poses = {} + for action_id, (_, pose) in enumerate(DEFAULT_ACTIONS): + viewmats, Ks, action = pose_to_input(pose, 32) + poses[action_id] = ( + viewmats.unsqueeze(0).to(device=device, dtype=torch.bfloat16), + Ks.unsqueeze(0).to(device=device, dtype=torch.bfloat16), + action.unsqueeze(0).to(device=device, dtype=torch.bfloat16), + ) + flow_shift = float(pipe.config.flow_shift) + scheduler = HunyuanVideo_1_5_Pipeline._create_scheduler(flow_shift) + scheduler.set_timesteps(4, device=device) + if scheduler.timesteps.dtype != torch.float32: + raise TypeError( + f"Scheduler timesteps must be FP32, got {scheduler.timesteps.dtype}" + ) + scheduler_timesteps = scheduler.timesteps + sigmas = scheduler.sigmas + # Static encoder outputs live in the sidecar; only the frozen Teacher remains. + pipe.transformer = None + del pipe + gc.collect() + torch.cuda.empty_cache() + + active_case_id: int | None = None + text_cache: list[dict] | None = None + full_condition: torch.Tensor | None = None + histories: dict[int, torch.Tensor] = {} + + grouped_teacher_forcing = args.training_mode == "grouped_teacher_forcing" + run_config = { + **vars(args), + "data_root": str(data_root), + "sidecar_root": str(sidecar_root), + "text_kv_root": str(text_kv_root), + "cases_path": str(cases_path), + "output_dir": str(output_dir), + "world_size": world_size, + "effective_stop_step": stop_after_steps, + "case_ids": case_ids, + "total_records_per_epoch": total_records, + "records_per_rank_per_epoch": records_per_rank, + "source_blocks": [1, 52], + "training_objective": ( + "grouped_offline_teacher_forcing" + if grouped_teacher_forcing else "timestep_self_rollout" + ), + "rollout_axis": "none" if grouped_teacher_forcing else "timestep_only", + "chunk_state": "offline_full_teacher", + "step1_target": "offline_full_teacher", + "step2_step3_targets": ( + "offline_full_teacher" + if grouped_teacher_forcing + else "online_full_teacher_on_student_latent" + ), + "timestep_graph": ( + "F1_to_P1|F2_to_P2|F3_to_P3_independent" + if grouped_teacher_forcing else "P1_to_P2_to_P3" + ), + "detach_between_timesteps": args.detach_between_timesteps, + "timestep_gradient_graph": ( + "three_independent_offline_pairs_one_mean_loss" + if grouped_teacher_forcing + else ( + "P1|P2|P3_stop_gradient" + if args.detach_between_timesteps + else "P1_to_P2_to_P3_full_bptt" + ) + ), + "flow_shift": flow_shift, + "scheduler_timesteps": scheduler_timesteps.detach().cpu().tolist(), + "scheduler_timestep_dtype": str(scheduler_timesteps.dtype), + "scheduler_timestep_assertion": "exact_torch_equal_per_record", + "scheduler_validation": scheduler_validation, + "activation_dtype": "bfloat16", + "trainable_parameter_dtype": ( + "float32" if args.fp32_trainable_params else "bfloat16" + ), + } + if rank == 0: + (output_dir / "train_config.json").write_text( + json.dumps(run_config, indent=2, ensure_ascii=False, default=str, sort_keys=True) + + "\n", + encoding="utf-8", + ) + print(json.dumps(run_config, indent=2, ensure_ascii=False, default=str), flush=True) + + swanlab_run = None + if rank == 0: + if args.swanlab_mode == "cloud": + swanlab.login(api_key=os.environ.get("SWANLAB_API_KEY"), save=False) + swanlab_run = swanlab.init( + project=args.swanlab_project, + workspace=args.swanlab_workspace, + experiment_name=args.swanlab_experiment, + description=( + f"v4 grouped offline teacher forcing on {len(case_ids)} cases; steps 1/2/3 " + "share one optimizer update" + if grouped_teacher_forcing + else ( + f"v4 detached timestep-only on-policy rollout on {len(case_ids)} cases; " + "offline Full chunk " + f"state; replay_weight={args.replay_weight}" + ) + ), + tags=[ + "predictor-v4", + ( + "grouped-teacher" + if grouped_teacher_forcing else "timestep-rollout" + ), + f"fulltrain-{len(case_ids)}cases", + f"replay-{args.replay_weight:g}", "8gpu", + ], + config=run_config, + logdir=str(output_dir / "swanlab"), + mode=args.swanlab_mode, + ) + + log_path = output_dir / "train_log.jsonl" + try: + resume_epoch = epoch + resume_batch = batch_in_epoch + while global_step < stop_after_steps: + epoch_records = rank_epoch_records( + records_by_case, + seed=args.seed, + epoch=epoch, + rank=rank, + world_size=world_size, + ) + if len(epoch_records) != records_per_rank: + raise RuntimeError( + f"Rank {rank} received {len(epoch_records)} records, " + f"expected {records_per_rank}" + ) + for loader_index, (case_id, record) in enumerate(epoch_records): + if epoch == resume_epoch and loader_index < resume_batch: + continue + if global_step >= stop_after_steps: + break + if case_id != active_case_id: + text_cache = None + full_condition = None + histories = {} + gc.collect() + torch.cuda.empty_cache() + text_kv_path = ( + text_kv_root + / f"case_{case_id:04d}_text_kv_exact_all54_bf16.safetensors" + ) + history_paths = { + action_id: sidecar_root / "histories" + / f"case_{case_id:04d}" / f"action_{action_id:02d}.safetensors" + for action_id in range(4) + } + for path in (text_kv_path, *history_paths.values()): + if not path.is_file(): + raise FileNotFoundError(path) + text_cache = load_precomputed_text_cache(text_kv_path, device) + full_condition = build_full_condition( + data_root / "cases" / f"case_{case_id:04d}.safetensors", + device, + ) + histories = { + action_id: load_file(str(path), device="cpu")[ + "chunk_final_latents" + ] + for action_id, path in history_paths.items() + } + active_case_id = case_id + if text_cache is None or full_condition is None: + raise RuntimeError("Active case resources were not initialized") + replay_step = ( + 1 + ((global_step + rank) % 3) if args.replay_weight > 0 else None + ) + cpu_item = load_timestep_rollout_record( + record_root, + record, + replay_step=replay_step, + grouped_teacher_forcing=grouped_teacher_forcing, + ) + item = move_timestep_rollout_record(cpu_item, device) + action_id = item["action_id"] + viewmats, Ks, action = poses[action_id] + history_latent = histories[action_id].to( + device=device, dtype=torch.bfloat16 + ) + torch.cuda.reset_peak_memory_stats(device) + if grouped_teacher_forcing: + metrics = train_grouped_teacher_forcing_record( + teacher=teacher, predictor=predictor, + predictor_model=predictor_model, optimizer=optimizer, + item=item, text_cache=text_cache, + history_latent=history_latent, + full_condition=full_condition, + full_viewmats=viewmats, full_Ks=Ks, full_action=action, + scheduler_timesteps=scheduler_timesteps, + args=args, device=device, + ) + else: + metrics = train_record( + teacher=teacher, predictor=predictor, + predictor_model=predictor_model, optimizer=optimizer, + item=item, text_cache=text_cache, + history_latent=history_latent, + full_condition=full_condition, + full_viewmats=viewmats, full_Ks=Ks, full_action=action, + sigmas=sigmas, scheduler_timesteps=scheduler_timesteps, + args=args, device=device, + ) + if global_step == 0: + validate_optimizer_precision( + optimizer, fp32_trainable_params=args.fp32_trainable_params, + require_state=True, + ) + lr_scheduler.step() + global_step += 1 + batch_in_epoch = loader_index + 1 + + metric_names = ( + "loss", "rollout_loss", "replay_loss", + "hidden_mse_step1", "hidden_mse_step2", "hidden_mse_step3", + "velocity_mse_step1", "velocity_mse_step2", "velocity_mse_step3", + "replay_hidden_mse", "replay_velocity_mse", + "cross_timestep_grad_l1", "record_time_s", + ) + values = torch.tensor( + [metrics[name] for name in metric_names], + device=device, dtype=torch.float64, + ) + averaged = distributed_average(values) + grad_max = torch.tensor(metrics["grad_norm"], device=device) + dist.all_reduce(grad_max, op=dist.ReduceOp.MAX) + if rank == 0 and (global_step % args.log_every == 0 or global_step == 1): + record_log = { + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + "rank0_case_id": case_id, + "rank0_action_id": action_id, + "rank0_chunk_id": int(item["chunk_id"]), + **{ + name: float(averaged[index]) + for index, name in enumerate(metric_names) + }, + "grad_norm_max": float(grad_max), + "lr_fusion": optimizer.param_groups[0]["lr"], + "lr_blocks": optimizer.param_groups[1]["lr"], + "peak_memory_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "replay_weight": args.replay_weight, + } + with log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record_log, sort_keys=True) + "\n") + swanlab.log(record_log, step=global_step) + print(json.dumps(record_log, sort_keys=True), flush=True) + + should_save = ( + global_step % args.save_every == 0 + or global_step == stop_after_steps + ) + if should_save: + dist.barrier() + if rank == 0: + weights_path = output_dir / f"predictor_step_{global_step:05d}.safetensors" + save_predictor_weights( + predictor_model, weights_path, floating_dtype=torch.bfloat16 + ) + save_training_state( + output_dir / "training_latest.pt", + { + "model": trainable_state_dict(predictor_model), + "optimizer": optimizer.state_dict(), + "scheduler": lr_scheduler.state_dict(), + "global_step": global_step, + "epoch": epoch, + "batch_in_epoch": batch_in_epoch, + "args": vars(args), + "predictor_config": predictor_model.config_dict, + "weights_path": str(weights_path), + }, + ) + prune_snapshots(output_dir, args.keep_snapshots) + dist.barrier() + del cpu_item, item, history_latent, metrics, values, averaged, grad_max + gc.collect() + del epoch_records + epoch += 1 + batch_in_epoch = 0 + resume_batch = 0 + except BaseException as error: + if rank == 0 and swanlab_run is not None: + try: + swanlab.finish(error=str(error)) + except Exception: + pass + raise + else: + if rank == 0 and swanlab_run is not None: + swanlab.finish() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/trainer/README.md b/trainer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d0e70cd5534dbf78aba31ccffdae2421baa30252 --- /dev/null +++ b/trainer/README.md @@ -0,0 +1,71 @@ +# HY-World 1.5 (WorldPlay) Training Guideline + +We provide a detailed guideline, including dataset structure and training scripts, to train the base model on your own dataset. Specifically, now we only support the autoregressive training with action control and memory design based on HunyuanVideo 1.5. + +## Prepare the Dataset +The json file structure is presented as below: + +``` +{ + "latent": latents, # video latent after vae encoding + "prompt_embeds": prompt_embeds, # text prompt encoding latent + "image_cond": cond_latents, # image latent + "vision_states": vision_states, # sigLip feature + "prompt_mask": attention_mask, # text prompt mask + "byt5_text_states": byt5_embeddings, # byt5 text prompt + "byt5_text_mask": byt5_masks # byt5 text prompt mask +} +``` +Specifically, we pre-process the video, txt, and image latent to accelerate the training process. The dataset processing is similar to the illustration in [Fastvideo](https://github.com/hao-ai-lab/FastVideo). + +## Install Dependencies +Same as the inference guideline, we suggest to install the dependencies: + +```bash +conda create --name worldplay python=3.10 -y +conda activate worldplay +pip install -r requirements.txt +``` + +## Running Training Scripts +We provide a training command on a single node as below: + +``` +bash scripts/training/hyvideo15/run_ar_hunyuan_action_mem.sh +``` +Specifically, to better understand our training framework, we provide a detailed illustration about our training parameters. + +| Argument | Type | Required | Default | Description | +|----------|------|----------|---------|-------------| +| `--json_path` | str | Yes | - | Path to training dataset json file | +| `--train_time_shift` | int | Yes | 3 | Timestep scheduler shift | +| `--window_frames` | int | Yes | 24 | The window length of video latents | +| `--model_path` | str | Yes | - | Path to pretrained model directory | +| `--output_dir` | str | Yes | - | Output directory for model checkpoints | +| `--max_train_steps` | int | Yes | - | Training steps | +| `--train_batch_size` | int | Yes | 1 | Data number of per rank | +| `--train_sp_batch_size` | int | Yes | 1 | Data number of Sequence Parallel Group | +| `--gradient_accumulation_steps` | int | Yes | 1 | The number of gradient accumulation steps | +| `--num_gpus` | int | Yes | - | Total GPU number | +| `--sp_size` | int | Yes | 4 | Sequence Parallel Size (GPU number for each sp group) | +| `--tp_size` | int | Yes | 1 | Tensor Parallel Size, set as 1 | +| `--hsdp_replicate_dim` | int | Yes | - | how many times that sharded setup is replicated across the cluster, set to $NODES | +| `--hsdp_shard_dim` | int | Yes | - | how many GPUs the model states (weights/gradients) are sharded across (set to $NUM_GPUS) | +| `--cls_name` | str | Yes | - | Class name for transformer | +| `--ar_action_load_from_dir` | str | No | - | Load path for autoregressive model with action control, mainly for memory training | +| `--checkpointing_steps` | int | Yes | 500 | Checkpoint saving interval step | +| `--resume_from_checkpoint` | str | No | - | Resume checkpoint path, e.g., /your_path/to/checkpoints/checkpoint-1000 | + + +We provide a detailed calculation of total batch size as follows: +``` +total_batch_size = world_size / training_args.sp_size * + training_args.gradient_accumulation_steps * + training_args.train_sp_batch_size +``` +where `world_size` represents the total number of GPUs. For example, we have one 8*H100 GPU node, training_args.sp_size = 4 with others set as default, the total batch size will be 2. + +## Practices + ++ We suggest pre-processing the dataset and outputing the json form to accelerate the dataset loading procedure. ++ Larger batch size is strongly suggested. Usually, the total batch size at least should set to 32. \ No newline at end of file diff --git a/trainer/STA_configuration.py b/trainer/STA_configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..be29f48b61c91ad33616768854c98a86bc945c3b --- /dev/null +++ b/trainer/STA_configuration.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from collections import defaultdict +from typing import Any + +import numpy as np + +from trainer.utils import dict_to_3d_list + + +def configure_sta(mode: str = 'STA_searching', + layer_num: int = 40, + time_step_num: int = 50, + head_num: int = 40, + **kwargs) -> list[list[list[Any]]]: + """ + Configure Sliding Tile Attention (STA) parameters based on the specified mode. + + Parameters: + ---------- + mode : str + The STA mode to use. Options are: + - 'STA_searching': Generate a set of mask candidates for initial search + - 'STA_tuning': Select best mask strategy based on previously saved results + - 'STA_inference': Load and use a previously tuned mask strategy + layer_num: int, number of layers + time_step_num: int, number of timesteps + head_num: int, number of heads + + **kwargs : dict + Mode-specific parameters: + + For 'STA_searching': + - mask_candidates: list of str, optional, mask candidates to use + - mask_selected: list of int, optional, indices of selected masks + + For 'STA_tuning': + - mask_search_files_path: str, required, path to mask search results + - mask_candidates: list of str, optional, mask candidates to use + - mask_selected: list of int, optional, indices of selected masks + - skip_time_steps: int, optional, number of time steps to use full attention (default 12) + - save_dir: str, optional, directory to save mask strategy (default "mask_candidates") + + For 'STA_inference': + - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json") + """ + valid_modes = [ + 'STA_searching', 'STA_tuning', 'STA_inference', 'STA_tuning_cfg' + ] + if mode not in valid_modes: + raise ValueError(f"Mode must be one of {valid_modes}, got {mode}") + + if mode == 'STA_searching': + # Get parameters with defaults + mask_candidates: list[str] | None = kwargs.get('mask_candidates') + if mask_candidates is None: + raise ValueError( + "mask_candidates is required for STA_searching mode") + mask_selected: list[int] = kwargs.get('mask_selected', + list(range(len(mask_candidates)))) + + # Parse selected masks + selected_masks: list[list[int]] = [] + for index in mask_selected: + mask = mask_candidates[index] + masks_list = [int(x) for x in mask.split(',')] + selected_masks.append(masks_list) + + # Create 3D mask structure with fixed dimensions (t=50, l=60) + masks_3d: list[list[list[list[int]]]] = [] + for i in range(time_step_num): # Fixed t dimension = 50 + row = [] + for j in range(layer_num): # Fixed l dimension = 60 + row.append(selected_masks) # Add all masks at each position + masks_3d.append(row) + + return masks_3d + + elif mode == 'STA_tuning': + # Get required parameters + mask_search_files_path: str | None = kwargs.get( + 'mask_search_files_path') + if not mask_search_files_path: + raise ValueError( + "mask_search_files_path is required for STA_tuning mode") + + # Get optional parameters with defaults + mask_candidates_tuning: list[str] | None = kwargs.get('mask_candidates') + if mask_candidates_tuning is None: + raise ValueError("mask_candidates is required for STA_tuning mode") + mask_selected_tuning: list[int] = kwargs.get( + 'mask_selected', list(range(len(mask_candidates_tuning)))) + skip_time_steps_tuning: int | None = kwargs.get('skip_time_steps') + save_dir_tuning: str | None = kwargs.get('save_dir', "mask_candidates") + + # Parse selected masks + selected_masks_tuning: list[list[int]] = [] + for index in mask_selected_tuning: + mask = mask_candidates_tuning[index] + masks_list = [int(x) for x in mask.split(',')] + selected_masks_tuning.append(masks_list) + + # Read JSON results + results = read_specific_json_files(mask_search_files_path) + averaged_results = average_head_losses(results, selected_masks_tuning) + + # Add full attention mask for specific cases + full_attention_mask_tuning: list[int] | None = kwargs.get( + 'full_attention_mask') + if full_attention_mask_tuning is not None: + selected_masks_tuning.append(full_attention_mask_tuning) + + # Select best mask strategy + timesteps_tuning: int = kwargs.get('timesteps', time_step_num) + if skip_time_steps_tuning is None: + skip_time_steps_tuning = 12 + mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( + averaged_results, selected_masks_tuning, skip_time_steps_tuning, + timesteps_tuning, head_num) + + # Save mask strategy + if save_dir_tuning is not None: + os.makedirs(save_dir_tuning, exist_ok=True) + file_path = os.path.join( + save_dir_tuning, + f'mask_strategy_s{skip_time_steps_tuning}.json') + with open(file_path, 'w') as f: + json.dump(mask_strategy, f, indent=4) + print(f"Successfully saved mask_strategy to {file_path}") + + # Print sparsity and strategy counts for information + print(f"Overall sparsity: {sparsity:.4f}") + print("\nStrategy usage counts:") + total_heads = time_step_num * layer_num * head_num # Fixed dimensions + for strategy, count in strategy_counts.items(): + print( + f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)" + ) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list(mask_strategy, + t_max=time_step_num, + l_max=layer_num, + h_max=head_num) + + return mask_strategy_3d + elif mode == 'STA_tuning_cfg': + # Get required parameters for both positive and negative paths + mask_search_files_path_pos: str | None = kwargs.get( + 'mask_search_files_path_pos') + mask_search_files_path_neg: str | None = kwargs.get( + 'mask_search_files_path_neg') + save_dir_cfg: str | None = kwargs.get('save_dir') + + if not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg: + raise ValueError( + "mask_search_files_path_pos, mask_search_files_path_neg, and save_dir are required for STA_tuning_cfg mode" + ) + + # Get optional parameters with defaults + mask_candidates_cfg: list[str] | None = kwargs.get('mask_candidates') + if mask_candidates_cfg is None: + raise ValueError( + "mask_candidates is required for STA_tuning_cfg mode") + mask_selected_cfg: list[int] = kwargs.get( + 'mask_selected', list(range(len(mask_candidates_cfg)))) + skip_time_steps_cfg: int | None = kwargs.get('skip_time_steps') + + # Parse selected masks + selected_masks_cfg: list[list[int]] = [] + for index in mask_selected_cfg: + mask = mask_candidates_cfg[index] + masks_list = [int(x) for x in mask.split(',')] + selected_masks_cfg.append(masks_list) + + # Read JSON results for both positive and negative paths + pos_results = read_specific_json_files(mask_search_files_path_pos) + neg_results = read_specific_json_files(mask_search_files_path_neg) + # Combine positive and negative results into one list + combined_results = pos_results + neg_results + + # Average the combined results + averaged_results = average_head_losses(combined_results, + selected_masks_cfg) + + # Add full attention mask for specific cases + full_attention_mask_cfg: list[int] | None = kwargs.get( + 'full_attention_mask') + if full_attention_mask_cfg is not None: + selected_masks_cfg.append(full_attention_mask_cfg) + + timesteps_cfg: int = kwargs.get('timesteps', time_step_num) + if skip_time_steps_cfg is None: + skip_time_steps_cfg = 12 + # Select best mask strategy using combined results + mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( + averaged_results, selected_masks_cfg, skip_time_steps_cfg, + timesteps_cfg, head_num) + + # Save mask strategy + os.makedirs(save_dir_cfg, exist_ok=True) + file_path = os.path.join(save_dir_cfg, + f'mask_strategy_s{skip_time_steps_cfg}.json') + with open(file_path, 'w') as f: + json.dump(mask_strategy, f, indent=4) + print(f"Successfully saved mask_strategy to {file_path}") + + # Print sparsity and strategy counts for information + print(f"Overall sparsity: {sparsity:.4f}") + print("\nStrategy usage counts:") + total_heads = time_step_num * layer_num * head_num # Fixed dimensions + for strategy, count in strategy_counts.items(): + print( + f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)" + ) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list(mask_strategy, + t_max=time_step_num, + l_max=layer_num, + h_max=head_num) + + return mask_strategy_3d + + else: # STA_inference + # Get parameters with defaults + load_path: str | None = kwargs.get( + 'load_path', "mask_candidates/mask_strategy.json") + if load_path is None: + raise ValueError("load_path is required for STA_inference mode") + + # Load previously saved mask strategy + with open(load_path) as f: + mask_strategy = json.load(f) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list(mask_strategy, + t_max=time_step_num, + l_max=layer_num, + h_max=head_num) + + return mask_strategy_3d + + +# Helper functions + + +def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]: + """Read and parse JSON files containing mask search results.""" + json_contents: list[dict[str, Any]] = [] + + # List files only in the current directory (no walk) + files = os.listdir(folder_path) + # Filter files + matching_files = [f for f in files if 'mask' in f and f.endswith('.json')] + print(f"Found {len(matching_files)} matching files: {matching_files}") + + for file_name in matching_files: + file_path = os.path.join(folder_path, file_name) + with open(file_path) as file: + data = json.load(file) + json_contents.append(data) + + return json_contents + + +def average_head_losses( + results: list[dict[str, Any]], + selected_masks: list[list[int]]) -> dict[str, dict[str, np.ndarray]]: + """Average losses across all prompts for each mask strategy.""" + # Initialize a dictionary to store the averaged results + averaged_losses: dict[str, dict[str, np.ndarray]] = {} + loss_type = 'L2_loss' + # Get all loss types (e.g., 'L2_loss') + averaged_losses[loss_type] = {} + + for mask in selected_masks: + mask_str = str(mask) + data_shape = np.array(results[0][loss_type][mask_str]).shape + accumulated_data = np.zeros(data_shape) + + # Sum across all prompts + for prompt_result in results: + accumulated_data += np.array(prompt_result[loss_type][mask_str]) + + # Average by dividing by number of prompts + averaged_data = accumulated_data / len(results) + averaged_losses[loss_type][mask_str] = averaged_data + + return averaged_losses + + +def select_best_mask_strategy( + averaged_results: dict[str, dict[str, np.ndarray]], + selected_masks: list[list[int]], + skip_time_steps: int = 12, + timesteps: int = 50, + head_num: int = 40 +) -> tuple[dict[str, list[int]], float, dict[str, int]]: + """Select the best mask strategy for each head based on loss minimization.""" + best_mask_strategy: dict[str, list[int]] = {} + loss_type = 'L2_loss' + # Get the shape of time steps and layers + layers = len(averaged_results[loss_type][str(selected_masks[0])][0]) + + # Counter for sparsity calculation + total_tokens = 0 # total number of masked tokens + total_length = 0 # total sequence length + + strategy_counts: dict[str, int] = { + str(strategy): 0 + for strategy in selected_masks + } + full_attn_strategy = selected_masks[-1] # Last strategy is full attention + print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ") + + for t in range(timesteps): + for layer_idx in range(layers): + for h in range(head_num): + if t < skip_time_steps: # First steps use full attention + strategy = full_attn_strategy + else: + # Get losses for this head across all strategies + head_losses = [] + for strategy in selected_masks[: + -1]: # Exclude full attention + head_losses.append(averaged_results[loss_type][str( + strategy)][t][layer_idx][h]) + + # Find which strategy gives minimum loss + best_strategy_idx = np.argmin(head_losses) + strategy = selected_masks[best_strategy_idx] + + best_mask_strategy[f'{t}_{layer_idx}_{h}'] = strategy + + # Calculate sparsity + nums = strategy # strategy is already a list of numbers + total_tokens += nums[0] * nums[1] * nums[ + 2] # masked tokens for chosen strategy + total_length += full_attn_strategy[0] * full_attn_strategy[ + 1] * full_attn_strategy[2] + + # Count strategy usage + strategy_counts[str(strategy)] += 1 + + overall_sparsity = 1 - total_tokens / total_length + + return best_mask_strategy, overall_sparsity, strategy_counts + + +def save_mask_search_results( + mask_search_final_result: list[dict[str, list[float]]], + prompt: str, + mask_strategies: list[str], + output_dir: str = 'output/mask_search_result/') -> str | None: + if not mask_search_final_result: + print("No mask search results to save") + return None + + # Create result dictionary with defaultdict for nested lists + mask_search_dict: dict[str, dict[str, list[list[float]]]] = { + "L2_loss": defaultdict(list), + "L1_loss": defaultdict(list) + } + + mask_selected = list(range(len(mask_strategies))) + selected_masks: list[list[int]] = [] + for index in mask_selected: + mask = mask_strategies[index] + masks_list = [int(x) for x in mask.split(',')] + selected_masks.append(masks_list) + + # Process each mask strategy + for i, mask_strategy in enumerate(selected_masks): + mask_strategy_str = str(mask_strategy) + # Process L2 loss + step_results: list[list[float]] = [] + for step_data in mask_search_final_result: + if isinstance(step_data, dict) and "L2_loss" in step_data: + layer_losses = [float(loss) for loss in step_data["L2_loss"]] + step_results.append(layer_losses) + mask_search_dict["L2_loss"][mask_strategy_str] = step_results + + step_results = [] + for step_data in mask_search_final_result: + if isinstance(step_data, dict) and "L1_loss" in step_data: + layer_losses = [float(loss) for loss in step_data["L1_loss"]] + step_results.append(layer_losses) + mask_search_dict["L1_loss"][mask_strategy_str] = step_results + + # Create the output directory if it doesn't exist + os.makedirs(output_dir, exist_ok=True) + + # Create a filename based on the first 20 characters of the prompt + filename = prompt[:50].replace(" ", "_") + filepath = os.path.join(output_dir, f'mask_search_{filename}.json') + + # Save the results to a JSON file + with open(filepath, 'w') as f: + json.dump(mask_search_dict, f, indent=4) + + print(f"Successfully saved mask research results to {filepath}") + + return filepath diff --git a/trainer/__init__.py b/trainer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f2966b40c343e0bf79366ffc9d91f5fd2ea4fe5e --- /dev/null +++ b/trainer/__init__.py @@ -0,0 +1,6 @@ +from trainer.configs.pipelines import PipelineConfig +from trainer.configs.sample import SamplingParam +from trainer.entrypoints.video_generator import VideoGenerator +from trainer.version import __version__ + +__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"] \ No newline at end of file diff --git a/trainer/attention/__init__.py b/trainer/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b43c1a467e346a62fc755543574b105d7a8987cf --- /dev/null +++ b/trainer/attention/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +from trainer.attention.backends.abstract import (AttentionBackend, + AttentionMetadata, + AttentionMetadataBuilder) +from trainer.attention.layer import (DistributedAttention, + DistributedAttention_VSA, LocalAttention) +from trainer.attention.selector import get_attn_backend + +__all__ = [ + "DistributedAttention", + "LocalAttention", + "DistributedAttention_VSA", + "AttentionBackend", + "AttentionMetadata", + "AttentionMetadataBuilder", + # "AttentionState", + "get_attn_backend", +] diff --git a/trainer/attention/backends/__init__.py b/trainer/attention/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/attention/backends/abstract.py b/trainer/attention/backends/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..34a5bb6654f4e17c1ff375978be5049606177a0c --- /dev/null +++ b/trainer/attention/backends/abstract.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py + +from abc import ABC, abstractmethod +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar + +if TYPE_CHECKING: + pass + +import torch + + +class AttentionBackend(ABC): + """Abstract class for attention backends.""" + # For some attention backends, we allocate an output tensor before + # calling the custom op. When piecewise cudagraph is enabled, this + # makes sure the output tensor is allocated inside the cudagraph. + accept_output_buffer: bool = False + + @staticmethod + @abstractmethod + def get_name() -> str: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_impl_cls() -> type["AttentionImpl"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + # @staticmethod + # @abstractmethod + # def get_state_cls() -> Type["AttentionState"]: + # raise NotImplementedError + + # @classmethod + # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": + # return cls.get_metadata_cls()(*args, **kwargs) + + @staticmethod + @abstractmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + +@dataclass +class AttentionMetadata: + """Attention metadata for prefill and decode batched together.""" + # Current step of diffusion process + current_timestep: int + + def asdict_zerocopy(self, + skip_fields: set[str] | None = None) -> dict[str, Any]: + """Similar to dataclasses.asdict, but avoids deepcopying.""" + if skip_fields is None: + skip_fields = set() + # Note that if we add dataclasses as fields, they will need + # similar handling. + return { + field.name: getattr(self, field.name) + for field in fields(self) if field.name not in skip_fields + } + + +T = TypeVar("T", bound=AttentionMetadata) + + +class AttentionMetadataBuilder(ABC, Generic[T]): + """Abstract class for attention metadata builders.""" + + @abstractmethod + def __init__(self) -> None: + """Create the builder, remember some configuration and parameters.""" + raise NotImplementedError + + @abstractmethod + def prepare(self) -> None: + """Prepare for one batch.""" + raise NotImplementedError + + @abstractmethod + def build( + self, + **kwargs: dict[str, Any], + ) -> AttentionMetadata: + """Build attention metadata with on-device tensors.""" + raise NotImplementedError + + +class AttentionLayer(Protocol): + + _k_scale: torch.Tensor + _v_scale: torch.Tensor + _k_scale_float: float + _v_scale_float: float + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + ... + + +class AttentionImpl(ABC, Generic[T]): + + @abstractmethod + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + raise NotImplementedError + + def preprocess_qkv(self, qkv: torch.Tensor, + attn_metadata: T) -> torch.Tensor: + """Preprocess QKV tensor before performing attention operation. + + Default implementation returns the tensor unchanged. + Subclasses can override this to implement custom preprocessing + like reshaping, tiling, scaling, or other transformations. + + Called AFTER all_to_all for distributed attention + + Args: + qkv: The query-key-value tensor + attn_metadata: Metadata for the attention operation + + Returns: + Processed QKV tensor + """ + return qkv + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: T, + ) -> torch.Tensor: + """Postprocess the output tensor after the attention operation. + + Default implementation returns the tensor unchanged. + Subclasses can override this to implement custom postprocessing + like untiling, scaling, or other transformations. + + Called BEFORE all_to_all for distributed attention + + Args: + output: The output tensor from the attention operation + attn_metadata: Metadata for the attention operation + + Returns: + Postprocessed output tensor + """ + + return output + + @abstractmethod + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T, + ) -> torch.Tensor: + raise NotImplementedError diff --git a/trainer/attention/backends/flash_attn.py b/trainer/attention/backends/flash_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..2a034d9134eb2ad2bd957e8dbef33f04147c2c96 --- /dev/null +++ b/trainer/attention/backends/flash_attn.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from flash_attn import flash_attn_func as flash_attn_2_func + +try: + from flash_attn_interface import flash_attn_func as flash_attn_3_func + + # flash_attn 3 no longer have a different API, see following commit: + # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7 + flash_attn_func = flash_attn_3_func +except ImportError: + flash_attn_func = flash_attn_2_func + +from trainer.attention.backends.abstract import (AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder) +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class FlashAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN" + + @staticmethod + def get_impl_cls() -> type["FlashAttentionImpl"]: + return FlashAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + +class FlashAttentionImpl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ): + output = flash_attn_func( + query, # type: ignore[no-untyped-call] + key, + value, + softmax_scale=self.softmax_scale, + causal=self.causal) + return output diff --git a/trainer/attention/backends/sage_attn.py b/trainer/attention/backends/sage_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..4b90e923c5bb51131422159cdd6858bdbba5c203 --- /dev/null +++ b/trainer/attention/backends/sage_attn.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from sageattention import sageattn + +from trainer.attention.backends.abstract import ( # FlashAttentionMetadata, + AttentionBackend, AttentionImpl, AttentionMetadata) +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class SageAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "SAGE_ATTN" + + @staticmethod + def get_impl_cls() -> type["SageAttentionImpl"]: + return SageAttentionImpl + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class SageAttentionImpl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + self.dropout = extra_impl_args.get("dropout_p", 0.0) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + output = sageattn( + query, + key, + value, + # since input is (batch_size, seq_len, head_num, head_dim) + tensor_layout="NHD", + is_causal=self.causal) + return output diff --git a/trainer/attention/backends/sdpa.py b/trainer/attention/backends/sdpa.py new file mode 100644 index 0000000000000000000000000000000000000000..afd089085b8f0f2de6459e4487f9d999684d0907 --- /dev/null +++ b/trainer/attention/backends/sdpa.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from trainer.attention.backends.abstract import ( # FlashAttentionMetadata, + AttentionBackend, AttentionImpl, AttentionMetadata) +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class SDPABackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "SDPA" + + @staticmethod + def get_impl_cls() -> type["SDPAImpl"]: + return SDPAImpl + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class SDPAImpl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + self.dropout = extra_impl_args.get("dropout_p", 0.0) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + # transpose to bs, heads, seq_len, head_dim + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn_kwargs = { + "attn_mask": None, + "dropout_p": self.dropout, + "is_causal": self.causal, + "scale": self.softmax_scale + } + if query.shape[1] != key.shape[1]: + attn_kwargs["enable_gqa"] = True + output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, **attn_kwargs) + output = output.transpose(1, 2) + return output diff --git a/trainer/attention/backends/sliding_tile_attn.py b/trainer/attention/backends/sliding_tile_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..0994fcc95dfdead1f86740383f5534e90082df70 --- /dev/null +++ b/trainer/attention/backends/sliding_tile_attn.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +from dataclasses import dataclass +from typing import Any + +import torch +from einops import rearrange +from st_attn import sliding_tile_attention + +import trainer.envs as envs +from trainer.attention.backends.abstract import (AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder) +from trainer.distributed import get_sp_group +from trainer.forward_context import ForwardContext, get_forward_context +from trainer.logger import init_logger +from trainer.utils import dict_to_3d_list + +logger = init_logger(__name__) + + +class RangeDict(dict): + + def __getitem__(self, item: int) -> str: + for key in self.keys(): + if isinstance(key, tuple): + low, high = key + if low <= item <= high: + return str(super().__getitem__(key)) + elif key == item: + return str(super().__getitem__(key)) + raise KeyError(f"seq_len {item} not supported for STA") + + +class SlidingTileAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + # TODO(will-refactor): check this + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "SLIDING_TILE_ATTN" + + @staticmethod + def get_impl_cls() -> type["SlidingTileAttentionImpl"]: + return SlidingTileAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]: + return SlidingTileAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]: + return SlidingTileAttentionMetadataBuilder + + +@dataclass +class SlidingTileAttentionMetadata(AttentionMetadata): + current_timestep: int + STA_param: list[list[ + Any]] # each timestep with one metadata, shape [num_layers, num_heads] + + +class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): + + def __init__(self): + pass + + def prepare(self): + pass + + def build( # type: ignore + self, + STA_param: list[list[Any]], + current_timestep: int, + **kwargs: dict[str, Any], + ) -> SlidingTileAttentionMetadata: + param = STA_param + if param is None: + return SlidingTileAttentionMetadata( + current_timestep=current_timestep, STA_param=[]) + return SlidingTileAttentionMetadata(current_timestep=current_timestep, + STA_param=param[current_timestep]) + + +class SlidingTileAttentionImpl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + # TODO(will-refactor): for now this is the mask strategy, but maybe we should + # have a more general config for STA? + config_file = envs.TRAINER_ATTENTION_CONFIG + if config_file is None: + raise ValueError("TRAINER_ATTENTION_CONFIG is not set") + + # TODO(kevin): get mask strategy for different STA modes + with open(config_file) as f: + mask_strategy = json.load(f) + self.mask_strategy = dict_to_3d_list(mask_strategy) + + self.prefix = prefix + sp_group = get_sp_group() + self.sp_size = sp_group.world_size + # STA config + self.STA_base_tile_size = [6, 8, 8] + self.dit_seq_shape_mapping = RangeDict({ + (115200, 115456): '30x48x80', + 82944: '36x48x48', + 69120: '18x48x80', + }) + self.full_window_mapping = { + '30x48x80': [5, 6, 10], + '36x48x48': [6, 6, 6], + '18x48x80': [3, 6, 10] + } + + def tile(self, x: torch.Tensor) -> torch.Tensor: + return rearrange( + x, + "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", + n_t=self.full_window_size[0], + n_h=self.full_window_size[1], + n_w=self.full_window_size[2], + ts_t=self.STA_base_tile_size[0], + ts_h=self.STA_base_tile_size[1], + ts_w=self.STA_base_tile_size[2]) + + def untile(self, x: torch.Tensor) -> torch.Tensor: + x = rearrange( + x, + "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", + n_t=self.full_window_size[0], + n_h=self.full_window_size[1], + n_w=self.full_window_size[2], + ts_t=self.STA_base_tile_size[0], + ts_h=self.STA_base_tile_size[1], + ts_w=self.STA_base_tile_size[2]) + return x + + def preprocess_qkv( + self, + qkv: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + img_sequence_length = qkv.shape[1] + self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length] + self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str] + self.dit_seq_shape_int = list( + map(int, self.dit_seq_shape_str.split('x'))) + self.img_seq_length = self.dit_seq_shape_int[ + 0] * self.dit_seq_shape_int[1] * self.dit_seq_shape_int[2] + return self.tile(qkv) + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: SlidingTileAttentionMetadata, + ) -> torch.Tensor: + return self.untile(output) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_metadata: SlidingTileAttentionMetadata, + ) -> torch.Tensor: + if self.mask_strategy is None: + raise ValueError( + "mask_strategy cannot be None for SlidingTileAttention") + if self.mask_strategy[0] is None: + raise ValueError( + "mask_strategy[0] cannot be None for SlidingTileAttention") + + timestep = attn_metadata.current_timestep + forward_context: ForwardContext = get_forward_context() + forward_batch = forward_context.forward_batch + if forward_batch is None: + raise ValueError("forward_batch cannot be None") + # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl' + layer_idx = int(self.prefix.split('.')[-3]) + if attn_metadata.STA_param is None or len( + attn_metadata.STA_param) <= layer_idx: + raise ValueError("Invalid STA_param") + STA_param = attn_metadata.STA_param[layer_idx] + + text_length = q.shape[1] - self.img_seq_length + has_text = text_length > 0 + + query = q.transpose(1, 2).contiguous() + key = k.transpose(1, 2).contiguous() + value = v.transpose(1, 2).contiguous() + + head_num = query.size(1) + sp_group = get_sp_group() + current_rank = sp_group.rank_in_group + start_head = current_rank * head_num + + # searching or tuning mode + if len(STA_param) < head_num * sp_group.world_size: + sparse_attn_hidden_states_all = [] + full_mask_window = STA_param[-1] + for window_size in STA_param[:-1]: + sparse_hidden_states = sliding_tile_attention( + query, key, value, [window_size] * head_num, text_length, + has_text, self.dit_seq_shape_str).transpose(1, 2) + sparse_attn_hidden_states_all.append(sparse_hidden_states) + + hidden_states = sliding_tile_attention( + query, key, value, [full_mask_window] * head_num, text_length, + has_text, self.dit_seq_shape_str).transpose(1, 2) + + attn_L2_loss = [] + attn_L1_loss = [] + # average loss across all heads + for sparse_attn_hidden_states in sparse_attn_hidden_states_all: + # L2 loss + attn_L2_loss_ = torch.mean((sparse_attn_hidden_states.float() - + hidden_states.float())**2, + dim=[0, 1, 3]).cpu().numpy() + attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_] + attn_L2_loss.append(attn_L2_loss_) + # L1 loss + attn_L1_loss_ = torch.mean( + torch.abs(sparse_attn_hidden_states.float() - + hidden_states.float()), + dim=[0, 1, 3]).cpu().numpy() + attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_] + attn_L1_loss.append(attn_L1_loss_) + + layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss} + + if forward_batch.is_cfg_negative: + if forward_batch.mask_search_final_result_neg is not None: + forward_batch.mask_search_final_result_neg[timestep].append( + layer_loss_save) + else: + if forward_batch.mask_search_final_result_pos is not None: + forward_batch.mask_search_final_result_pos[timestep].append( + layer_loss_save) + else: + windows = [ + STA_param[head_idx + start_head] for head_idx in range(head_num) + ] + + hidden_states = sliding_tile_attention( + query, key, value, windows, text_length, has_text, + self.dit_seq_shape_str).transpose(1, 2) + + return hidden_states diff --git a/trainer/attention/backends/video_sparse_attn.py b/trainer/attention/backends/video_sparse_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..4910e6fe6b24e4a01a1c01cef41bb018f3a82ca7 --- /dev/null +++ b/trainer/attention/backends/video_sparse_attn.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +import functools +import math +from dataclasses import dataclass + +import torch + +try: + from vsa import video_sparse_attn +except ImportError: + video_sparse_attn = None + +from typing import Any + +from trainer.attention.backends.abstract import (AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder) +from trainer.distributed import get_sp_group +from trainer.logger import init_logger + +logger = init_logger(__name__) +VSA_TILE_SIZE = (4, 4, 4) + + +@functools.lru_cache(maxsize=10) +def get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + T, H, W = dit_seq_shape + ts, hs, ws = tile_size + indices = torch.arange(T * H * W, device=device, + dtype=torch.long).reshape(T, H, W) + ls = [] + for t in range(math.ceil(T / ts)): + for h in range(math.ceil(H / hs)): + for w in range(math.ceil(W / ws)): + ls.append(indices[t * ts:min(t * ts + ts, T), + h * hs:min(h * hs + hs, H), + w * ws:min(w * ws + ws, W)].flatten()) + index = torch.cat(ls, dim=0) + return index + + +@functools.lru_cache(maxsize=10) +def get_reverse_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + return torch.argsort( + get_tile_partition_indices(dit_seq_shape, tile_size, device)) + + +@functools.lru_cache(maxsize=10) +def construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + num_tiles: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + """ + Compute the number of valid (non‑padded) tokens inside every + (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order + (t‑tile, h‑tile, w‑tile) that `rearrange` uses. + + Returns + ------- + torch.LongTensor # shape: [∏ full_window_size] + """ + # unpack + t, h, w = dit_seq_shape + ts_t, ts_h, ts_w = VSA_TILE_SIZE + n_t, n_h, n_w = num_tiles + + def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor: + """Vector with the size of each tile along one dimension.""" + sizes = torch.full((n_tiles, ), tile, dtype=torch.int, device=device) + # size of last (possibly partial) tile + remainder = dim_len - (n_tiles - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(t, ts_t, n_t) # [n_t] + h_sizes = _sizes(h, ts_h, n_h) # [n_h] + w_sizes = _sizes(w, ts_w, n_w) # [n_w] + + # broadcast‑multiply to get voxels per tile, then flatten + block_sizes = ( + t_sizes[:, None, None] # [n_t, 1, 1] + * h_sizes[None, :, None] # [1, n_h, 1] + * w_sizes[None, None, :] # [1, 1, n_w] + ).reshape(-1) # [n_t * n_h * n_w] + + return block_sizes + + +@functools.lru_cache(maxsize=10) +def get_non_pad_index( + variable_block_sizes: torch.LongTensor, + max_block_size: int, +): + n_win = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts_pad = torch.arange(n_win, device=device) * max_block_size + index_pad = starts_pad[:, None] + torch.arange(max_block_size, + device=device)[None, :] + index_mask = torch.arange( + max_block_size, device=device)[None, :] < variable_block_sizes[:, None] + return index_pad[index_mask] + + +class VideoSparseAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_name() -> str: + return "VIDEO_SPARSE_ATTN" + + @staticmethod + def get_impl_cls() -> type["VideoSparseAttentionImpl"]: + return VideoSparseAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]: + return VideoSparseAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]: + return VideoSparseAttentionMetadataBuilder + + +@dataclass +class VideoSparseAttentionMetadata(AttentionMetadata): + current_timestep: int + dit_seq_shape: list[int] + VSA_sparsity: float + num_tiles: list[int] + total_seq_length: int + tile_partition_indices: torch.LongTensor + reverse_tile_partition_indices: torch.LongTensor + variable_block_sizes: torch.LongTensor + non_pad_index: torch.LongTensor + + +class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): + + def __init__(self): + pass + + def prepare(self): + pass + + def build( # type: ignore + self, + current_timestep: int, + raw_latent_shape: tuple[int, int, int], + patch_size: tuple[int, int, int], + VSA_sparsity: float, + device: torch.device, + **kwargs: dict[str, Any], + ) -> VideoSparseAttentionMetadata: + patch_size = patch_size + dit_seq_shape = (raw_latent_shape[0] // patch_size[0], + raw_latent_shape[1] // patch_size[1], + raw_latent_shape[2] // patch_size[2]) + + num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), + math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]), + math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2])) + total_seq_length = math.prod(dit_seq_shape) + + tile_partition_indices = get_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device) + reverse_tile_partition_indices = get_reverse_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device) + variable_block_sizes = construct_variable_block_sizes( + dit_seq_shape, num_tiles, device) + non_pad_index = get_non_pad_index(variable_block_sizes, + math.prod(VSA_TILE_SIZE)) + + return VideoSparseAttentionMetadata( + current_timestep=current_timestep, + dit_seq_shape=dit_seq_shape, # type: ignore + VSA_sparsity=VSA_sparsity, # type: ignore + num_tiles=num_tiles, # type: ignore + total_seq_length=total_seq_length, # type: ignore + tile_partition_indices=tile_partition_indices, # type: ignore + reverse_tile_partition_indices=reverse_tile_partition_indices, + variable_block_sizes=variable_block_sizes, + non_pad_index=non_pad_index) + + +class VideoSparseAttentionImpl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.prefix = prefix + sp_group = get_sp_group() + self.sp_size = sp_group.world_size + + def tile(self, x: torch.Tensor, num_tiles: list[int], + tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor) -> torch.Tensor: + t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0] + h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1] + w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2] + + x_padded = torch.zeros( + (x.shape[0], t_padded_size * h_padded_size * w_padded_size, + x.shape[-2], x.shape[-1]), + device=x.device, + dtype=x.dtype) + x_padded[:, non_pad_index] = x[:, tile_partition_indices] + return x_padded + + def untile(self, x: torch.Tensor, + reverse_tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor) -> torch.Tensor: + x = x[:, non_pad_index][:, reverse_tile_partition_indices] + return x + + def preprocess_qkv( + self, + qkv: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + return self.tile(qkv, attn_metadata.num_tiles, + attn_metadata.tile_partition_indices, + attn_metadata.non_pad_index) + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + return self.untile(output, attn_metadata.reverse_tile_partition_indices, + attn_metadata.non_pad_index) + + def forward( # type: ignore[override] + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + gate_compress: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + gate_compress = gate_compress.transpose(1, 2).contiguous() + + VSA_sparsity = attn_metadata.VSA_sparsity + + cur_topk = math.ceil( + (1 - VSA_sparsity) * + (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))) + + if video_sparse_attn is None: + raise NotImplementedError("video_sparse_attn is not installed") + hidden_states = video_sparse_attn( + query, + key, + value, + variable_block_sizes=attn_metadata.variable_block_sizes, + topk=cur_topk, + block_size=VSA_TILE_SIZE, + compress_attn_weight=gate_compress).transpose(1, 2) + + return hidden_states diff --git a/trainer/attention/layer.py b/trainer/attention/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..7a46d5db18383a0fc699e7a0c8e00b77dcdac433 --- /dev/null +++ b/trainer/attention/layer.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn + +from trainer.attention.selector import backend_name_to_enum, get_attn_backend +from trainer.distributed.communication_op import ( + sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D) +from trainer.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) +from trainer.forward_context import ForwardContext, get_forward_context +from trainer.platforms import AttentionBackendEnum +from trainer.utils import get_compute_dtype + + +class DistributedAttention(nn.Module): + """Distributed attention layer. + """ + + def __init__(self, + num_heads: int, + head_size: int, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + **extra_impl_args) -> None: + super().__init__() + if softmax_scale is None: + self.softmax_scale = head_size**-0.5 + else: + self.softmax_scale = softmax_scale + + if num_kv_heads is None: + num_kv_heads = num_heads + + dtype = get_compute_dtype() + attn_backend = get_attn_backend( + head_size, + dtype, + supported_attention_backends=supported_attention_backends) + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls(num_heads=num_heads, + head_size=head_size, + causal=causal, + softmax_scale=self.softmax_scale, + num_kv_heads=num_kv_heads, + prefix=f"{prefix}.impl", + **extra_impl_args) + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + self.backend = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: torch.Tensor | None = None, + replicated_k: torch.Tensor | None = None, + replicated_v: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass for distributed attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens + replicated_k (Optional[torch.Tensor]): Replicated key tensor + replicated_v (Optional[torch.Tensor]): Replicated value tensor + + Returns: + Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: + - o (torch.Tensor): Output tensor after attention for the main sequence + - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided + """ + # Check input shapes + assert q.dim() == 4 and k.dim() == 4 and v.dim( + ) == 4, "Expected 4D tensors" + batch_size, seq_len, num_heads, head_dim = q.shape + local_rank = get_sp_parallel_rank() + world_size = get_sp_world_size() + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Stack QKV + qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] + + # Redistribute heads across sequence dimension + qkv = sequence_model_parallel_all_to_all_4D(qkv, + scatter_dim=2, + gather_dim=1) + # Apply backend-specific preprocess_qkv + qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) + + # Concatenate with replicated QKV if provided + if replicated_q is not None: + assert replicated_k is not None and replicated_v is not None + replicated_qkv = torch.cat( + [replicated_q, replicated_k, replicated_v], + dim=0) # [3, seq_len, num_heads, head_dim] + heads_per_rank = num_heads // world_size + replicated_qkv = replicated_qkv[:, :, local_rank * + heads_per_rank:(local_rank + 1) * + heads_per_rank] + qkv = torch.cat([qkv, replicated_qkv], dim=1) + + q, k, v = qkv.chunk(3, dim=0) + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + + # Redistribute back if using sequence parallelism + replicated_output = None + if replicated_q is not None: + replicated_output = output[:, seq_len * world_size:] + output = output[:, :seq_len * world_size] + # TODO: make this asynchronous + replicated_output = sequence_model_parallel_all_gather( + replicated_output.contiguous(), dim=2) + # Apply backend-specific postprocess_output + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + output = sequence_model_parallel_all_to_all_4D(output, + scatter_dim=1, + gather_dim=2) + return output, replicated_output + + +class DistributedAttention_VSA(DistributedAttention): + """Distributed attention layer with VSA support. + """ + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: torch.Tensor | None = None, + replicated_k: torch.Tensor | None = None, + replicated_v: torch.Tensor | None = None, + gate_compress: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass for distributed attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens + replicated_k (Optional[torch.Tensor]): Replicated key tensor + replicated_v (Optional[torch.Tensor]): Replicated value tensor + + Returns: + Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: + - o (torch.Tensor): Output tensor after attention for the main sequence + - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided + """ + # Check text tokens are not supported for VSA now + assert replicated_q is None and replicated_k is None and replicated_v is None, "Replicated QKV is not supported for VSA now" + # Check input shapes + assert q.dim() == 4 and k.dim() == 4 and v.dim( + ) == 4, "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Stack QKV + qkvg = torch.cat([q, k, v, gate_compress], + dim=0) # [3, seq_len, num_heads, head_dim] + + # Redistribute heads across sequence dimension + qkvg = sequence_model_parallel_all_to_all_4D(qkvg, + scatter_dim=2, + gather_dim=1) + + qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata) + + q, k, v, gate_compress = qkvg.chunk(4, dim=0) + output = self.attn_impl.forward( + q, k, v, gate_compress, ctx_attn_metadata) # type: ignore[call-arg] + + # Redistribute back if using sequence parallelism + replicated_output = None + + # Apply backend-specific postprocess_output + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + output = sequence_model_parallel_all_to_all_4D(output, + scatter_dim=1, + gather_dim=2) + return output, replicated_output + + +class LocalAttention(nn.Module): + """Attention layer. + """ + + def __init__(self, + num_heads: int, + head_size: int, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + **extra_impl_args) -> None: + super().__init__() + if softmax_scale is None: + self.softmax_scale = head_size**-0.5 + else: + self.softmax_scale = softmax_scale + if num_kv_heads is None: + num_kv_heads = num_heads + + dtype = get_compute_dtype() + attn_backend = get_attn_backend( + head_size, + dtype, + supported_attention_backends=supported_attention_backends) + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls(num_heads=num_heads, + head_size=head_size, + softmax_scale=self.softmax_scale, + num_kv_heads=num_kv_heads, + causal=causal, + **extra_impl_args) + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + self.backend = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + """ + Apply local attention between query, key and value tensors. + + Args: + q (torch.Tensor): Query tensor of shape [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor of shape [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor of shape [batch_size, seq_len, num_heads, head_dim] + + Returns: + torch.Tensor: Output tensor after local attention + """ + # Check input shapes + assert q.dim() == 4 and k.dim() == 4 and v.dim( + ) == 4, "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + return output diff --git a/trainer/attention/selector.py b/trainer/attention/selector.py new file mode 100644 index 0000000000000000000000000000000000000000..916df05ac50db48de8b82fb5ee1541bc90198cc5 --- /dev/null +++ b/trainer/attention/selector.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py + +import os +from collections.abc import Generator +from contextlib import contextmanager +from functools import cache +from typing import cast + +import torch + +import trainer.envs as envs +from trainer.attention.backends.abstract import AttentionBackend +from trainer.logger import init_logger +from trainer.platforms import AttentionBackendEnum, current_platform +from trainer.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname + +logger = init_logger(__name__) + + +def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None: + """ + Convert a string backend name to a _Backend enum value. + + Returns: + * _Backend: enum value if backend_name is a valid in-tree type + * None: otherwise it's an invalid in-tree type or an out-of-tree platform is + loaded. + """ + assert backend_name is not None + return AttentionBackendEnum[backend_name] if backend_name in AttentionBackendEnum.__members__ else \ + None + + +def get_env_variable_attn_backend() -> AttentionBackendEnum | None: + ''' + Get the backend override specified by the Trainer attention + backend environment variable, if one is specified. + + Returns: + + * _Backend enum value if an override is specified + * None otherwise + ''' + backend_name = os.environ.get(STR_BACKEND_ENV_VAR) + return (None + if backend_name is None else backend_name_to_enum(backend_name)) + + +# Global state allows a particular choice of backend +# to be forced, overriding the logic which auto-selects +# a backend based on system & workload configuration +# (default behavior if this variable is None) +# +# THIS SELECTION TAKES PRECEDENCE OVER THE +# TRAINER ATTENTION BACKEND ENVIRONMENT VARIABLE +forced_attn_backend: AttentionBackendEnum | None = None + + +def global_force_attn_backend( + attn_backend: AttentionBackendEnum | None) -> None: + ''' + Force all attention operations to use a specified backend. + + Passing `None` for the argument re-enables automatic + backend selection., + + Arguments: + + * attn_backend: backend selection (None to revert to auto) + ''' + global forced_attn_backend + forced_attn_backend = attn_backend + + +def get_global_forced_attn_backend() -> AttentionBackendEnum | None: + ''' + Get the currently-forced choice of attention backend, + or None if auto-selection is currently enabled. + ''' + return forced_attn_backend + + +def get_attn_backend( + head_size: int, + dtype: torch.dtype, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, +) -> type[AttentionBackend]: + return _cached_get_attn_backend(head_size, dtype, + supported_attention_backends) + + +@cache +def _cached_get_attn_backend( + head_size: int, + dtype: torch.dtype, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, +) -> type[AttentionBackend]: + # Check whether a particular choice of backend was + # previously forced. + # + # THIS SELECTION OVERRIDES THE TRAINER_ATTENTION_BACKEND + # ENVIRONMENT VARIABLE. + if not supported_attention_backends: + raise ValueError("supported_attention_backends is empty") + selected_backend = None + backend_by_global_setting: AttentionBackendEnum | None = ( + get_global_forced_attn_backend()) + if backend_by_global_setting is not None: + selected_backend = backend_by_global_setting + else: + # Check the environment variable and override if specified + backend_by_env_var: str | None = envs.TRAINER_ATTENTION_BACKEND + if backend_by_env_var is not None: + selected_backend = backend_name_to_enum(backend_by_env_var) + + # get device-specific attn_backend + if selected_backend not in supported_attention_backends: + selected_backend = None + attention_cls = current_platform.get_attn_backend_cls( + selected_backend, head_size, dtype) + if not attention_cls: + raise ValueError( + f"Invalid attention backend for {current_platform.device_name}") + return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls)) + + +@contextmanager +def global_force_attn_backend_context_manager( + attn_backend: AttentionBackendEnum) -> Generator[None, None, None]: + ''' + Globally force a Trainer attention backend override within a + context manager, reverting the global attention backend + override to its prior state upon exiting the context + manager. + + Arguments: + + * attn_backend: attention backend to force + + Returns: + + * Generator + ''' + + # Save the current state of the global backend override (if any) + original_value = get_global_forced_attn_backend() + + # Globally force the new backend override + global_force_attn_backend(attn_backend) + + # Yield control back to the enclosed code block + try: + yield + finally: + # Revert the original global backend override, if any + global_force_attn_backend(original_value) diff --git a/trainer/configs/__init__.py b/trainer/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/configs/configs.py b/trainer/configs/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..e901430238528098a203e123f270ba515fb4a7de --- /dev/null +++ b/trainer/configs/configs.py @@ -0,0 +1,151 @@ +import dataclasses +from typing import Any, Optional + +from trainer.configs.utils import update_config_from_args +from trainer.utils import FlexibleArgumentParser, StoreBoolean + + +@dataclasses.dataclass +class PreprocessConfig: + """Configuration for preprocessing operations.""" + + # Model and dataset configuration + model_path: str = "" + dataset_path: str = "" + dataset_output_dir: str = "./output" + + # Dataloader configuration + dataloader_num_workers: int = 1 + preprocess_video_batch_size: int = 2 + + # Saver configuration + samples_per_file: int = 64 + flush_frequency: int = 256 + + # Video processing parameters + max_height: int = 480 + max_width: int = 848 + num_frames: int = 163 + video_length_tolerance_range: float = 2.0 + train_fps: int = 30 + speed_factor: float = 1.0 + drop_short_ratio: float = 1.0 + do_temporal_sample: bool = False + + # Model configuration + training_cfg_rate: float = 0.0 + + seed: int = 42 + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser, + prefix: str = "preprocess") -> FlexibleArgumentParser: + """Add preprocessing configuration arguments to the parser.""" + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + + preprocess_args = parser.add_argument_group("Preprocessing Arguments") + # Model & Dataset + preprocess_args.add_argument(f"--{prefix_with_dot}model-path", + type=str, + default=PreprocessConfig.model_path, + help="Path to the model for preprocessing") + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-path", + type=str, + default=PreprocessConfig.dataset_path, + help="Path to the dataset directory for preprocessing") + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-output-dir", + type=str, + default=PreprocessConfig.dataset_output_dir, + help="The output directory where the dataset will be written.") + + # Dataloader + preprocess_args.add_argument( + f"--{prefix_with_dot}dataloader-num-workers", + type=int, + default=PreprocessConfig.dataloader_num_workers, + help= + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}preprocess-video-batch-size", + type=int, + default=PreprocessConfig.preprocess_video_batch_size, + help="Batch size (per device) for the training dataloader.") + + # Saver + preprocess_args.add_argument(f"--{prefix_with_dot}samples-per-file", + type=int, + default=PreprocessConfig.samples_per_file, + help="Number of samples per output file") + preprocess_args.add_argument(f"--{prefix_with_dot}flush-frequency", + type=int, + default=PreprocessConfig.flush_frequency, + help="How often to save to parquet files") + + # Video processing parameters + preprocess_args.add_argument(f"--{prefix_with_dot}max-height", + type=int, + default=PreprocessConfig.max_height, + help="Maximum height for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}max-width", + type=int, + default=PreprocessConfig.max_width, + help="Maximum width for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}num-frames", + type=int, + default=PreprocessConfig.num_frames, + help="Number of frames to process") + preprocess_args.add_argument( + f"--{prefix_with_dot}video-length-tolerance-range", + type=float, + default=PreprocessConfig.video_length_tolerance_range, + help="Video length tolerance range") + preprocess_args.add_argument(f"--{prefix_with_dot}train-fps", + type=int, + default=PreprocessConfig.train_fps, + help="Training FPS") + preprocess_args.add_argument(f"--{prefix_with_dot}speed-factor", + type=float, + default=PreprocessConfig.speed_factor, + help="Speed factor for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}drop-short-ratio", + type=float, + default=PreprocessConfig.drop_short_ratio, + help="Ratio for dropping short videos") + preprocess_args.add_argument( + f"--{prefix_with_dot}do-temporal-sample", + action=StoreBoolean, + default=PreprocessConfig.do_temporal_sample, + help="Whether to do temporal sampling") + + # Model Training configuration + preprocess_args.add_argument(f"--{prefix_with_dot}training-cfg-rate", + type=float, + default=PreprocessConfig.training_cfg_rate, + help="Training CFG rate") + preprocess_args.add_argument(f"--{prefix_with_dot}seed", + type=int, + default=PreprocessConfig.seed, + help="Seed for random number generator") + + return parser + + @classmethod + def from_kwargs(cls, kwargs: dict[str, + Any]) -> Optional["PreprocessConfig"]: + """Create PreprocessConfig from keyword arguments.""" + preprocess_config = cls() + if not update_config_from_args( + preprocess_config, kwargs, prefix="preprocess", pop_args=True): + return None + return preprocess_config + + def check_preprocess_config(self) -> None: + if self.dataset_path == "": + raise ValueError("dataset_path must be set for preprocess mode") + if self.samples_per_file <= 0: + raise ValueError("samples_per_file must be greater than 0") + if self.flush_frequency <= 0: + raise ValueError("flush_frequency must be greater than 0") diff --git a/trainer/configs/fasthunyuan_t2v.json b/trainer/configs/fasthunyuan_t2v.json new file mode 100644 index 0000000000000000000000000000000000000000..6eb29ae44701746b8660666c4ef91378d89f91c6 --- /dev/null +++ b/trainer/configs/fasthunyuan_t2v.json @@ -0,0 +1,48 @@ +{ + "embedded_cfg_scale": 6, + "flow_shift": 17, + "dit_cpu_offload": false, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": true, + "vae_sp": true, + "vae_config": { + "load_encoder": false, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 4, + "use_tiling": true, + "use_temporal_tiling": true, + "use_parallel_tiling": true + }, + "dit_config": { + "prefix": "Hunyuan", + "quant_config": null + }, + "text_encoder_precisions": [ + "fp16", + "fp16" + ], + "text_encoder_configs": [ + { + "prefix": "llama", + "quant_config": null, + "lora_config": null + }, + { + "prefix": "clip", + "quant_config": null, + "lora_config": null, + "num_hidden_layers_override": null, + "require_post_norm": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false +} \ No newline at end of file diff --git a/trainer/configs/models/__init__.py b/trainer/configs/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e09b9c34bf5691fd41b43d8d729c66ab1d64dfaa --- /dev/null +++ b/trainer/configs/models/__init__.py @@ -0,0 +1,6 @@ +from trainer.configs.models.base import ModelConfig +from trainer.configs.models.dits.base import DiTConfig +from trainer.configs.models.encoders.base import EncoderConfig +from trainer.configs.models.vaes.base import VAEConfig + +__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"] diff --git a/trainer/configs/models/base.py b/trainer/configs/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..951dbafd524148c5411d902d996118f7d26dd6bc --- /dev/null +++ b/trainer/configs/models/base.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field, fields +from typing import Any + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model) +# 2. ArchConfig should be inherited & overridden by each model arch_config +# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users +@dataclass +class ArchConfig: + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=list + ) # mapping from huggingface weight names to custom names + + +@dataclass +class ModelConfig: + # Every model config parameter can be categorized into either ArchConfig or everything else + # Diffuser/Transformer parameters + arch_config: ArchConfig = field(default_factory=ArchConfig) + + # Trainer-specific parameters here + # i.e. STA, quantization, teacache + + def __getattr__(self, name): + # Only called if 'name' is not found in ModelConfig directly + if hasattr(self.arch_config, name): + return getattr(self.arch_config, name) + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'") + + def __getstate__(self): + # Return a dictionary of attributes to pickle + # Convert to dict and exclude any problematic attributes + state = self.__dict__.copy() + return state + + def __setstate__(self, state): + # Restore instance attributes from the unpickled state + self.__dict__.update(state) + + # This should be used only when loading from transformers/diffusers + def update_model_arch(self, source_model_dict: dict[str, Any]) -> None: + arch_config = self.arch_config + valid_fields = {f.name for f in fields(arch_config)} + + for key, value in source_model_dict.items(): + if key in valid_fields: + setattr(arch_config, key, value) + else: + raise AttributeError( + f"{type(arch_config).__name__} has no field '{key}'") + + if hasattr(arch_config, "__post_init__"): + arch_config.__post_init__() + + def update_model_config(self, source_model_dict: dict[str, Any]) -> None: + assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config." + + valid_fields = {f.name for f in fields(self)} + + for key, value in source_model_dict.items(): + if key in valid_fields: + setattr(self, key, value) + else: + logger.warning("%s does not contain field '%s'!", + type(self).__name__, key) + raise AttributeError(f"Invalid field: {key}") + + if hasattr(self, "__post_init__"): + self.__post_init__() diff --git a/trainer/configs/models/dits/__init__.py b/trainer/configs/models/dits/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a951e79486a0778123c7c27fe59f7bbcc9f1474 --- /dev/null +++ b/trainer/configs/models/dits/__init__.py @@ -0,0 +1,5 @@ +from trainer.configs.models.dits.hunyuanvideo import HunyuanVideoConfig +from trainer.configs.models.dits.stepvideo import StepVideoConfig +from trainer.configs.models.dits.wanvideo import WanVideoConfig + +__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"] diff --git a/trainer/configs/models/dits/base.py b/trainer/configs/models/dits/base.py new file mode 100644 index 0000000000000000000000000000000000000000..858acef31be1245fafee3cffc5767062bc667bfc --- /dev/null +++ b/trainer/configs/models/dits/base.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Any + +from trainer.configs.models.base import ArchConfig, ModelConfig +from trainer.layers.quantization import QuantizationConfig +from trainer.platforms import AttentionBackendEnum + + +@dataclass +class DiTArchConfig(ArchConfig): + _fsdp_shard_conditions: list = field(default_factory=list) + _compile_conditions: list = field(default_factory=list) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, + AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.VIDEO_SPARSE_ATTN) + + hidden_size: int = 0 + num_attention_heads: int = 0 + num_channels_latents: int = 0 + exclude_lora_layers: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if not self._compile_conditions: + self._compile_conditions = self._fsdp_shard_conditions.copy() + + +@dataclass +class DiTConfig(ModelConfig): + arch_config: DiTArchConfig = field(default_factory=DiTArchConfig) + + # TrainerDiT-specific parameters + prefix: str = "" + quant_config: QuantizationConfig | None = None + + @staticmethod + def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any: + """Add CLI arguments for DiTConfig fields""" + parser.add_argument( + f"--{prefix}.prefix", + type=str, + dest=f"{prefix.replace('-', '_')}.prefix", + default=DiTConfig.prefix, + help="Prefix for the DiT model", + ) + + parser.add_argument( + f"--{prefix}.quant-config", + type=str, + dest=f"{prefix.replace('-', '_')}.quant_config", + default=None, + help="Quantization configuration for the DiT model", + ) + + return parser diff --git a/trainer/configs/models/dits/hunyuanvideo.py b/trainer/configs/models/dits/hunyuanvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..4dbe6ad5dcf8667dc0405498ed99347e490ad968 --- /dev/null +++ b/trainer/configs/models/dits/hunyuanvideo.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +import torch + +from trainer.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_double_block(n: str, m) -> bool: + return "double" in n and str.isdigit(n.split(".")[-1]) + + +def is_single_block(n: str, m) -> bool: + return "single" in n and str.isdigit(n.split(".")[-1]) + + +# def is_refiner_block(n: str, m) -> bool: +# return "refiner" in n and str.isdigit(n.split(".")[-1]) + + +def is_txt_in(n: str, m) -> bool: + return n.split(".")[-1] == "txt_in" + + +@dataclass +class HunyuanVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [is_double_block, is_single_block]) + + _compile_conditions: list = field( + default_factory=lambda: [is_double_block, is_single_block, is_txt_in]) + + param_names_mapping: dict = field( + default_factory=lambda: { + # 1. context_embedder.time_text_embed submodules (specific rules, applied first): + r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": + r"txt_in.t_embedder.mlp.fc_in.\1", + r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$": + r"txt_in.t_embedder.mlp.fc_out.\1", + r"^context_embedder\.proj_in\.(.*)$": + r"txt_in.input_embedder.\1", + r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$": + r"txt_in.c_embedder.fc_in.\1", + r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$": + r"txt_in.c_embedder.fc_out.\1", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$": + r"txt_in.refiner_blocks.\1.norm1.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$": + r"txt_in.refiner_blocks.\1.norm2.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": + (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": + (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": + (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": + r"txt_in.refiner_blocks.\1.self_attn_proj.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": + r"txt_in.refiner_blocks.\1.mlp.fc_in.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": + r"txt_in.refiner_blocks.\1.mlp.fc_out.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": + r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", + + # 3. x_embedder mapping: + r"^x_embedder\.proj\.(.*)$": + r"img_in.proj.\1", + + # 4. Top-level time_text_embed mappings: + r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$": + r"time_in.mlp.fc_in.\1", + r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$": + r"time_in.mlp.fc_out.\1", + r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$": + r"guidance_in.mlp.fc_in.\1", + r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$": + r"guidance_in.mlp.fc_out.\1", + r"^time_text_embed\.text_embedder\.linear_1\.(.*)$": + r"vector_in.fc_in.\1", + r"^time_text_embed\.text_embedder\.linear_2\.(.*)$": + r"vector_in.fc_out.\1", + + # 5. transformer_blocks mapping: + r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$": + r"double_blocks.\1.img_mod.linear.\2", + r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$": + r"double_blocks.\1.txt_mod.linear.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": + r"double_blocks.\1.img_attn_q_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": + r"double_blocks.\1.img_attn_k_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": + (r"double_blocks.\1.img_attn_qkv.\2", 0, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": + (r"double_blocks.\1.img_attn_qkv.\2", 1, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": + (r"double_blocks.\1.img_attn_qkv.\2", 2, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": + (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": + (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": + (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": + r"double_blocks.\1.img_attn_proj.\2", + # Corrected: merge attn.to_add_out into the main projection. + r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$": + r"double_blocks.\1.txt_attn_proj.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$": + r"double_blocks.\1.txt_attn_q_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$": + r"double_blocks.\1.txt_attn_k_norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": + r"double_blocks.\1.img_mlp.fc_in.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": + r"double_blocks.\1.img_mlp.fc_out.\2", + r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$": + r"double_blocks.\1.txt_mlp.fc_in.\2", + r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$": + r"double_blocks.\1.txt_mlp.fc_out.\2", + + # 6. single_transformer_blocks mapping: + r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": + r"single_blocks.\1.q_norm.\2", + r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": + r"single_blocks.\1.k_norm.\2", + r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": + (r"single_blocks.\1.linear1.\2", 0, 4), + r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": + (r"single_blocks.\1.linear1.\2", 1, 4), + r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": + (r"single_blocks.\1.linear1.\2", 2, 4), + r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": + (r"single_blocks.\1.linear1.\2", 3, 4), + # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch. + r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$": + r"single_blocks.\1.linear2.\2", + r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$": + r"single_blocks.\1.modulation.linear.\2", + + # 7. Final layers mapping: + r"^norm_out\.linear\.(.*)$": + r"final_layer.adaLN_modulation.linear.\1", + r"^proj_out\.(.*)$": + r"final_layer.linear.\1", + }) + + # Reverse mapping for saving checkpoints: custom -> hf + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + patch_size: int = 2 + patch_size_t: int = 1 + in_channels: int = 16 + out_channels: int = 16 + num_attention_heads: int = 24 + attention_head_dim: int = 128 + mlp_ratio: float = 4.0 + num_layers: int = 20 + num_single_layers: int = 40 + num_refiner_layers: int = 2 + rope_axes_dim: tuple[int, int, int] = (16, 56, 56) + guidance_embeds: bool = False + dtype: torch.dtype | None = None + text_embed_dim: int = 4096 + pooled_projection_dim: int = 768 + rope_theta: int = 256 + qk_norm: str = "rms_norm" + exclude_lora_layers: list[str] = field( + default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"]) + + def __post_init__(self): + super().__post_init__() + self.hidden_size: int = self.attention_head_dim * self.num_attention_heads + self.num_channels_latents: int = self.in_channels + + +@dataclass +class HunyuanVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig) + + prefix: str = "Hunyuan" diff --git a/trainer/configs/models/dits/stepvideo.py b/trainer/configs/models/dits/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..834cde67f0574aa3b0e67cd99f3776bf4779d95b --- /dev/null +++ b/trainer/configs/models/dits/stepvideo.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class StepVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()]) + + param_names_mapping: dict = field( + default_factory=lambda: { + # transformer block + r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$": + r"transformer_blocks.\1.norm1.norm.\2", + r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$": + r"transformer_blocks.\1.norm2.norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$": + r"transformer_blocks.\1.ff.fc_in.weight", + r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$": + r"transformer_blocks.\1.ff.fc_out.weight", + + # adanorm block + r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$": + r"adaln_single.emb.mlp.fc_in.\1", + r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$": + r"adaln_single.emb.mlp.fc_out.\1", + + # caption projection + r"^caption_projection\.linear_1\.(weight|bias)$": + r"caption_projection.fc_in.\1", + r"^caption_projection\.linear_2\.(weight|bias)$": + r"caption_projection.fc_out.\1", + }) + + num_attention_heads: int = 48 + attention_head_dim: int = 128 + in_channels: int = 64 + out_channels: int | None = 64 + num_layers: int = 48 + dropout: float = 0.0 + patch_size: int = 1 + norm_type: str = "ada_norm_single" + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + caption_channels: int | list[int] | tuple[int, ...] | None = field( + default_factory=lambda: [6144, 1024]) + attention_type: str | None = "torch" + use_additional_conditions: bool | None = False + exclude_lora_layers: list[str] = field(default_factory=lambda: []) + + def __post_init__(self): + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.out_channels = self.in_channels if self.out_channels is None else self.out_channels + self.num_channels_latents = self.out_channels + + +@dataclass +class StepVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig) + + prefix: str = "StepVideo" diff --git a/trainer/configs/models/dits/wanvideo.py b/trainer/configs/models/dits/wanvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..bbd00f05b7f07147bba12252e71d04bfb0592496 --- /dev/null +++ b/trainer/configs/models/dits/wanvideo.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class WanVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^patch_embedding\.(.*)$": + r"patch_embedding.proj.\1", + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": + r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": + r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": + r"condition_embedder.time_modulation.linear.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_in.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_out.\1", + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": + r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": + r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": + r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": + r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": + r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": + r"blocks.\1.norm_k.\2", + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": + r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": + r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": + r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": + r"blocks.\1.self_attn_residual_norm.norm.\2", + }) + + # Reverse mapping for saving checkpoints: custom -> hf + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Some LoRA adapters use the original official layer names instead of hf layer names, + # so apply this before the param_names_mapping + lora_param_names_mapping: dict = field( + default_factory=lambda: { + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": + r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": + r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + }) + + patch_size: tuple[int, int, int] = (1, 2, 2) + text_len = 512 + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + image_dim: int | None = None + added_kv_proj_dim: int | None = None + rope_max_seq_len: int = 1024 + pos_embed_seq_len: int | None = None + exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class WanVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig) + + prefix: str = "Wan" diff --git a/trainer/configs/models/encoders/__init__.py b/trainer/configs/models/encoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0214004301ce1789c418ea935ba112637f865042 --- /dev/null +++ b/trainer/configs/models/encoders/__init__.py @@ -0,0 +1,14 @@ +from trainer.configs.models.encoders.base import (BaseEncoderOutput, + EncoderConfig, + ImageEncoderConfig, + TextEncoderConfig) +from trainer.configs.models.encoders.clip import (CLIPTextConfig, + CLIPVisionConfig) +from trainer.configs.models.encoders.llama import LlamaConfig +from trainer.configs.models.encoders.t5 import T5Config + +__all__ = [ + "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", + "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", + "T5Config" +] diff --git a/trainer/configs/models/encoders/base.py b/trainer/configs/models/encoders/base.py new file mode 100644 index 0000000000000000000000000000000000000000..5631af20114ecf23026fb06f5b9c8a9094982904 --- /dev/null +++ b/trainer/configs/models/encoders/base.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Any + +import torch + +from trainer.configs.models.base import ArchConfig, ModelConfig +from trainer.layers.quantization import QuantizationConfig +from trainer.platforms import AttentionBackendEnum + + +@dataclass +class EncoderArchConfig(ArchConfig): + architectures: list[str] = field(default_factory=lambda: []) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA) + output_hidden_states: bool = False + use_return_dict: bool = True + + +@dataclass +class TextEncoderArchConfig(EncoderArchConfig): + vocab_size: int = 0 + hidden_size: int = 0 + num_hidden_layers: int = 0 + num_attention_heads: int = 0 + pad_token_id: int = 0 + eos_token_id: int = 0 + text_len: int = 0 + hidden_state_skip_layer: int = 0 + decoder_start_token_id: int = 0 + output_past: bool = True + scalable_attention: bool = True + tie_word_embeddings: bool = False + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=list + ) # mapping from huggingface weight names to custom names + tokenizer_kwargs: dict[str, Any] = field(default_factory=dict) + _fsdp_shard_conditions: list = field(default_factory=lambda: []) + + def __post_init__(self) -> None: + self.tokenizer_kwargs = { + "truncation": True, + "max_length": self.text_len, + "return_tensors": "pt", + } + + +@dataclass +class ImageEncoderArchConfig(EncoderArchConfig): + pass + + +@dataclass +class BaseEncoderOutput: + last_hidden_state: torch.FloatTensor | None = None + pooler_output: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + attention_mask: torch.Tensor | None = None + + +@dataclass +class EncoderConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=EncoderArchConfig) + + prefix: str = "" + quant_config: QuantizationConfig | None = None + lora_config: Any | None = None + + +@dataclass +class TextEncoderConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig) + + +@dataclass +class ImageEncoderConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig) diff --git a/trainer/configs/models/encoders/clip.py b/trainer/configs/models/encoders/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..34cb411f3a8b65f13068da461a6e4c78c1aca740 --- /dev/null +++ b/trainer/configs/models/encoders/clip.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.encoders.base import (ImageEncoderArchConfig, + ImageEncoderConfig, + TextEncoderArchConfig, + TextEncoderConfig) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embeddings") + + +@dataclass +class CLIPTextArchConfig(TextEncoderArchConfig): + vocab_size: int = 49408 + hidden_size: int = 512 + intermediate_size: int = 2048 + projection_dim: int = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 8 + max_position_embeddings: int = 77 + hidden_act: str = "quick_gelu" + layer_norm_eps: float = 1e-5 + dropout: float = 0.0 + attention_dropout: float = 0.0 + initializer_range: float = 0.02 + initializer_factor: float = 1.0 + pad_token_id: int = 1 + bos_token_id: int = 49406 + eos_token_id: int = 49407 + text_len: int = 77 + stacked_params_mapping: list[tuple[str, str, + str]] = field(default_factory=lambda: [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ]) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_transformer_layer, _is_embeddings]) + + +@dataclass +class CLIPVisionArchConfig(ImageEncoderArchConfig): + hidden_size: int = 768 + intermediate_size: int = 3072 + projection_dim: int = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int = 224 + patch_size: int = 32 + hidden_act: str = "quick_gelu" + layer_norm_eps: float = 1e-5 + dropout: float = 0.0 + attention_dropout: float = 0.0 + initializer_range: float = 0.02 + initializer_factor: float = 1.0 + stacked_params_mapping: list[tuple[str, str, + str]] = field(default_factory=lambda: [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ]) + + +@dataclass +class CLIPTextConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field( + default_factory=CLIPTextArchConfig) + + num_hidden_layers_override: int | None = None + require_post_norm: bool | None = None + prefix: str = "clip" + + +@dataclass +class CLIPVisionConfig(ImageEncoderConfig): + arch_config: ImageEncoderArchConfig = field( + default_factory=CLIPVisionArchConfig) + + num_hidden_layers_override: int | None = None + require_post_norm: bool | None = None + prefix: str = "clip" diff --git a/trainer/configs/models/encoders/llama.py b/trainer/configs/models/encoders/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..916877e732122e8eb2c7e84deb036d57ebb1faa4 --- /dev/null +++ b/trainer/configs/models/encoders/llama.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.encoders.base import (TextEncoderArchConfig, + TextEncoderConfig) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m) -> bool: + return n.endswith("norm") + + +@dataclass +class LlamaArchConfig(TextEncoderArchConfig): + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int = 0 + bos_token_id: int = 1 + eos_token_id: int = 2 + pretraining_tp: int = 1 + tie_word_embeddings: bool = False + rope_theta: float = 10000.0 + rope_scaling: float | None = None + attention_bias: bool = False + attention_dropout: float = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + hidden_state_skip_layer: int = 2 + text_len: int = 256 + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), # type: ignore + (".gate_up_proj", ".up_proj", 1), # type: ignore + ]) + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [_is_transformer_layer, _is_embeddings, _is_final_norm]) + + +@dataclass +class LlamaConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig) + + prefix: str = "llama" diff --git a/trainer/configs/models/encoders/t5.py b/trainer/configs/models/encoders/t5.py new file mode 100644 index 0000000000000000000000000000000000000000..bc9982c5f231520bd247a2919b19c200b3c4081f --- /dev/null +++ b/trainer/configs/models/encoders/t5.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.encoders.base import (TextEncoderArchConfig, + TextEncoderConfig) + + +def _is_transformer_layer(n: str, m) -> bool: + return "block" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("shared") + + +def _is_final_layernorm(n: str, m) -> bool: + return n.endswith("final_layer_norm") + + +@dataclass +class T5ArchConfig(TextEncoderArchConfig): + vocab_size: int = 32128 + d_model: int = 512 + d_kv: int = 64 + d_ff: int = 2048 + num_layers: int = 6 + num_decoder_layers: int | None = None + num_heads: int = 8 + relative_attention_num_buckets: int = 32 + relative_attention_max_distance: int = 128 + dropout_rate: float = 0.1 + layer_norm_epsilon: float = 1e-6 + initializer_factor: float = 1.0 + feed_forward_proj: str = "relu" + dense_act_fn: str = "" + is_gated_act: bool = False + is_encoder_decoder: bool = True + use_cache: bool = True + pad_token_id: int = 0 + eos_token_id: int = 1 + classifier_dropout: float = 0.0 + text_len: int = 512 + stacked_params_mapping: list[tuple[str, str, + str]] = field(default_factory=lambda: [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q", "q"), + (".qkv_proj", ".k", "k"), + (".qkv_proj", ".v", "v"), + ]) + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [_is_transformer_layer, _is_embeddings, _is_final_layernorm]) + + # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py + def __post_init__(self): + super().__post_init__() + act_info = self.feed_forward_proj.split("-") + self.dense_act_fn: str = act_info[-1] + self.is_gated_act: bool = act_info[0] == "gated" + if self.feed_forward_proj == "gated-gelu": + self.dense_act_fn = "gelu_new" + + self.tokenizer_kwargs = { + "padding": "max_length", + "truncation": True, + "max_length": self.text_len, + "add_special_tokens": True, + "return_attention_mask": True, + "return_tensors": "pt", + } + + +@dataclass +class T5Config(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig) + + prefix: str = "t5" diff --git a/trainer/configs/models/vaes/__init__.py b/trainer/configs/models/vaes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..65bae61be28a71fa7bb976b500269bb09a61816a --- /dev/null +++ b/trainer/configs/models/vaes/__init__.py @@ -0,0 +1,9 @@ +from trainer.configs.models.vaes.hunyuanvae import HunyuanVAEConfig +from trainer.configs.models.vaes.stepvideovae import StepVideoVAEConfig +from trainer.configs.models.vaes.wanvae import WanVAEConfig + +__all__ = [ + "HunyuanVAEConfig", + "WanVAEConfig", + "StepVideoVAEConfig", +] diff --git a/trainer/configs/models/vaes/base.py b/trainer/configs/models/vaes/base.py new file mode 100644 index 0000000000000000000000000000000000000000..4dbc5dc1e22f53f07ac96c4d2d203bb8c48daa86 --- /dev/null +++ b/trainer/configs/models/vaes/base.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +import argparse +import dataclasses +from dataclasses import dataclass, field +from typing import Any + +import torch + +from trainer.configs.models.base import ArchConfig, ModelConfig +from trainer.utils import StoreBoolean + + +@dataclass +class VAEArchConfig(ArchConfig): + scaling_factor: float | torch.Tensor = 0 + + temporal_compression_ratio: int = 4 + spatial_compression_ratio: int = 8 + + +@dataclass +class VAEConfig(ModelConfig): + arch_config: VAEArchConfig = field(default_factory=VAEArchConfig) + + # TrainerVAE-specific parameters + load_encoder: bool = True + load_decoder: bool = True + + tile_sample_min_height: int = 256 + tile_sample_min_width: int = 256 + tile_sample_min_num_frames: int = 16 + tile_sample_stride_height: int = 192 + tile_sample_stride_width: int = 192 + tile_sample_stride_num_frames: int = 12 + blend_num_frames: int = 0 + + use_tiling: bool = True + use_temporal_tiling: bool = True + use_parallel_tiling: bool = True + use_temporal_scaling_frames: bool = True + + def __post_init__(self): + self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + @staticmethod + def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any: + """Add CLI arguments for VAEConfig fields""" + parser.add_argument( + f"--{prefix}.load-encoder", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.load_encoder", + default=VAEConfig.load_encoder, + help="Whether to load the VAE encoder", + ) + parser.add_argument( + f"--{prefix}.load-decoder", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.load_decoder", + default=VAEConfig.load_decoder, + help="Whether to load the VAE decoder", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-height", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_height", + default=VAEConfig.tile_sample_min_height, + help="Minimum height for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-width", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_width", + default=VAEConfig.tile_sample_min_width, + help="Minimum width for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames", + default=VAEConfig.tile_sample_min_num_frames, + help="Minimum number of frames for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-height", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height", + default=VAEConfig.tile_sample_stride_height, + help="Stride height for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-width", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width", + default=VAEConfig.tile_sample_stride_width, + help="Stride width for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames", + default=VAEConfig.tile_sample_stride_num_frames, + help="Stride number of frames for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.blend-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.blend_num_frames", + default=VAEConfig.blend_num_frames, + help="Number of frames to blend for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.use-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_tiling", + default=VAEConfig.use_tiling, + help="Whether to use tiling for VAE", + ) + parser.add_argument( + f"--{prefix}.use-temporal-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_temporal_tiling", + default=VAEConfig.use_temporal_tiling, + help="Whether to use temporal tiling for VAE", + ) + parser.add_argument( + f"--{prefix}.use-parallel-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_parallel_tiling", + default=VAEConfig.use_parallel_tiling, + help="Whether to use parallel tiling for VAE", + ) + + return parser + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig": + kwargs = {} + for attr in dataclasses.fields(cls): + value = getattr(args, attr.name, None) + if value is not None: + kwargs[attr.name] = value + return cls(**kwargs) diff --git a/trainer/configs/models/vaes/hunyuanvae.py b/trainer/configs/models/vaes/hunyuanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..65f31adb419ec339649adec20e0b227ebb444e32 --- /dev/null +++ b/trainer/configs/models/vaes/hunyuanvae.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class HunyuanVAEArchConfig(VAEArchConfig): + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + down_block_types: tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ) + up_block_types: tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ) + block_out_channels: tuple[int, ...] = (128, 256, 512, 512) + layers_per_block: int = 2 + act_fn: str = "silu" + norm_num_groups: int = 32 + scaling_factor: float = 0.476986 + spatial_compression_ratio: int = 8 + temporal_compression_ratio: int = 4 + mid_block_add_attention: bool = True + + def __post_init__(self): + self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) - + 1) + + +@dataclass +class HunyuanVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig) diff --git a/trainer/configs/models/vaes/stepvideovae.py b/trainer/configs/models/vaes/stepvideovae.py new file mode 100644 index 0000000000000000000000000000000000000000..7f7e2bac565c603aecc9908fce3331e3dfd932b0 --- /dev/null +++ b/trainer/configs/models/vaes/stepvideovae.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class StepVideoVAEArchConfig(VAEArchConfig): + in_channels: int = 3 + out_channels: int = 3 + z_channels: int = 64 + num_res_blocks: int = 2 + version: int = 2 + frame_len: int = 17 + world_size: int = 1 + + spatial_compression_ratio: int = 16 + temporal_compression_ratio: int = 8 + + scaling_factor: float = 1.0 + + +@dataclass +class StepVideoVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig) + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + use_temporal_scaling_frames: bool = False \ No newline at end of file diff --git a/trainer/configs/models/vaes/wanvae.py b/trainer/configs/models/vaes/wanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..0827f4d369bdf27b4500d3acec7e2dee1a0dab79 --- /dev/null +++ b/trainer/configs/models/vaes/wanvae.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +import torch + +from trainer.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class WanVAEArchConfig(VAEArchConfig): + base_dim: int = 96 + decoder_base_dim: int | None = None + z_dim: int = 16 + dim_mult: tuple[int, ...] = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple[float, ...] = () + temperal_downsample: tuple[bool, ...] = (False, True, True) + dropout: float = 0.0 + latents_mean: tuple[float, ...] = ( + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ) + latents_std: tuple[float, ...] = ( + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ) + is_residual: bool = False + in_channels: int = 3 + out_channels: int = 3 + patch_size: int | None = None + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 8 + clip_output: bool = True + + def __post_init__(self): + self.scaling_factor: torch.tensor = 1.0 / torch.tensor( + self.latents_std).view(1, self.z_dim, 1, 1, 1) + self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view( + 1, self.z_dim, 1, 1, 1) + self.temporal_compression_ratio = self.scale_factor_temporal + self.spatial_compression_ratio = self.scale_factor_spatial + + +@dataclass +class WanVAEConfig(VAEConfig): + arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig) + use_feature_cache: bool = True + + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def __post_init__(self): + self.blend_num_frames = (self.tile_sample_min_num_frames - + self.tile_sample_stride_num_frames) * 2 diff --git a/trainer/configs/pipelines/__init__.py b/trainer/configs/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4458995ba514ad9371a88f655e43daefb760985c --- /dev/null +++ b/trainer/configs/pipelines/__init__.py @@ -0,0 +1,15 @@ +from trainer.configs.pipelines.base import (PipelineConfig, + SlidingTileAttnConfig) +from trainer.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig +from trainer.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name) +from trainer.configs.pipelines.stepvideo import StepVideoT2VConfig +from trainer.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig, + WanT2V480PConfig, WanT2V720PConfig) + +__all__ = [ + "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig", + "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig", + "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig", + "get_pipeline_config_cls_from_name" +] diff --git a/trainer/configs/pipelines/base.py b/trainer/configs/pipelines/base.py new file mode 100644 index 0000000000000000000000000000000000000000..274b625c215ca513e3f72e42e234359f2940dcab --- /dev/null +++ b/trainer/configs/pipelines/base.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +from collections.abc import Callable +from dataclasses import asdict, dataclass, field, fields +from enum import Enum +from typing import Any, cast + +import torch + +from trainer.configs.models import (DiTConfig, EncoderConfig, ModelConfig, + VAEConfig) +from trainer.configs.models.encoders import BaseEncoderOutput +from trainer.configs.utils import update_config_from_args +from trainer.logger import init_logger +from trainer.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict + +logger = init_logger(__name__) + + +class STA_Mode(str, Enum): + """STA (Sliding Tile Attention) modes.""" + STA_INFERENCE = "STA_inference" + STA_SEARCHING = "STA_searching" + STA_TUNING = "STA_tuning" + STA_TUNING_CFG = "STA_tuning_cfg" + NONE = None + + +def preprocess_text(prompt: str) -> str: + return prompt + + +def postprocess_text(output: BaseEncoderOutput) -> torch.tensor: + raise NotImplementedError + + +# config for a single pipeline +@dataclass +class PipelineConfig: + """Base configuration for all pipeline architectures.""" + model_path: str = "" + pipeline_config_path: str | None = None + + # Video generation parameters + embedded_cfg_scale: float = 6.0 + flow_shift: float | None = None + disable_autocast: bool = False + + # Model configuration + dit_config: DiTConfig = field(default_factory=DiTConfig) + dit_precision: str = "bf16" + + # VAE configuration + vae_config: VAEConfig = field(default_factory=VAEConfig) + vae_precision: str = "fp32" + vae_tiling: bool = True + vae_sp: bool = True + + # Image encoder configuration + image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig) + image_encoder_precision: str = "fp32" + + # Text encoder configuration + DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32", ) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (EncoderConfig(), )) + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp32", )) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], + ...] = field(default_factory=lambda: + (postprocess_text, )) + + # StepVideo specific parameters + pos_magic: str | None = None + neg_magic: str | None = None + timesteps_scale: bool | None = None + + # STA (Sliding Tile Attention) parameters + mask_strategy_file_path: str | None = None + STA_mode: STA_Mode = STA_Mode.STA_INFERENCE + skip_time_steps: int = 15 + + # DMD parameters + dmd_denoising_steps: list[int] | None = field(default=None) + + # Compilation + # enable_torch_compile: bool = False + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser, + prefix: str = "") -> FlexibleArgumentParser: + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + + # model_path will be conflicting with the model_path in TrainerArgs, + # so we add it separately if prefix is not empty + if prefix_with_dot != "": + parser.add_argument( + f"--{prefix_with_dot}model-path", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}model_path", + default=PipelineConfig.model_path, + help="Path to the pretrained model", + ) + + parser.add_argument( + f"--{prefix_with_dot}pipeline-config-path", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}pipeline_config_path", + default=PipelineConfig.pipeline_config_path, + help="Path to the pipeline config", + ) + parser.add_argument( + f"--{prefix_with_dot}embedded-cfg-scale", + type=float, + dest=f"{prefix_with_dot.replace('-', '_')}embedded_cfg_scale", + default=PipelineConfig.embedded_cfg_scale, + help="Embedded CFG scale", + ) + parser.add_argument( + f"--{prefix_with_dot}flow-shift", + type=float, + dest=f"{prefix_with_dot.replace('-', '_')}flow_shift", + default=PipelineConfig.flow_shift, + help="Flow shift parameter", + ) + + # DiT configuration + parser.add_argument( + f"--{prefix_with_dot}dit-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}dit_precision", + default=PipelineConfig.dit_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for the DiT model", + ) + + # VAE configuration + parser.add_argument( + f"--{prefix_with_dot}vae-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}vae_precision", + default=PipelineConfig.vae_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for VAE", + ) + parser.add_argument( + f"--{prefix_with_dot}vae-tiling", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}vae_tiling", + default=PipelineConfig.vae_tiling, + help="Enable VAE tiling", + ) + parser.add_argument( + f"--{prefix_with_dot}vae-sp", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}vae_sp", + help="Enable VAE spatial parallelism", + ) + + # Text encoder configuration + parser.add_argument( + f"--{prefix_with_dot}text-encoder-precisions", + nargs="+", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}text_encoder_precisions", + default=PipelineConfig.DEFAULT_TEXT_ENCODER_PRECISIONS, + choices=["fp32", "fp16", "bf16"], + help="Precision for each text encoder", + ) + + # Image encoder configuration + parser.add_argument( + f"--{prefix_with_dot}image-encoder-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}image_encoder_precision", + default=PipelineConfig.image_encoder_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for image encoder", + ) + parser.add_argument( + f"--{prefix_with_dot}pos_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}pos_magic", + default=PipelineConfig.pos_magic, + help="Positive magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}neg_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}neg_magic", + default=PipelineConfig.neg_magic, + help="Negative magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}timesteps_scale", + type=bool, + dest=f"{prefix_with_dot.replace('-', '_')}timesteps_scale", + default=PipelineConfig.timesteps_scale, + help= + "Bool for applying scheduler scale in set_timesteps, used in stepvideo", + ) + + # DMD parameters + parser.add_argument( + f"--{prefix_with_dot}dmd-denoising-steps", + type=parse_int_list, + default=PipelineConfig.dmd_denoising_steps, + help= + "Comma-separated list of denoising steps (e.g., '1000,757,522')", + ) + + # Add VAE configuration arguments + from trainer.configs.models.vaes.base import VAEConfig + VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") + + # Add DiT configuration arguments + from trainer.configs.models.dits.base import DiTConfig + DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config") + + return parser + + def update_config_from_dict(self, + args: dict[str, Any], + prefix: str = "") -> None: + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + update_config_from_args(self, args, prefix, pop_args=True) + update_config_from_args(self.vae_config, + args, + f"{prefix_with_dot}vae_config", + pop_args=True) + update_config_from_args(self.dit_config, + args, + f"{prefix_with_dot}dit_config", + pop_args=True) + + @classmethod + def from_pretrained(cls, model_path: str) -> "PipelineConfig": + """ + use the pipeline class setting from model_path to match the pipeline config + """ + from trainer.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name) + pipeline_config_cls = get_pipeline_config_cls_from_name(model_path) + + return cast(PipelineConfig, pipeline_config_cls(model_path=model_path)) + + @classmethod + def from_kwargs(cls, + kwargs: dict[str, Any], + config_cli_prefix: str = "") -> "PipelineConfig": + """ + Load PipelineConfig from kwargs Dictionary. + kwargs: dictionary of kwargs + config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance + """ + from trainer.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name) + + prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip() + != "") else "" + model_path: str | None = kwargs.get(prefix_with_dot + 'model_path', + None) or kwargs.get('model_path') + pipeline_config_or_path: str | PipelineConfig | dict[ + str, Any] | None = kwargs.get(prefix_with_dot + 'pipeline_config', + None) or kwargs.get('pipeline_config') + if model_path is None: + raise ValueError("model_path is required in kwargs") + + # 1. Get the pipeline config class from the registry + pipeline_config_cls = get_pipeline_config_cls_from_name(model_path) + + # 2. Instantiate PipelineConfig + if pipeline_config_cls is None: + logger.warning( + "Couldn't find pipeline config for %s. Using the default pipeline config.", + model_path) + pipeline_config = cls() + else: + pipeline_config = pipeline_config_cls() + + # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided + if isinstance(pipeline_config_or_path, str): + pipeline_config.load_from_json(pipeline_config_or_path) + kwargs[prefix_with_dot + + 'pipeline_config_path'] = pipeline_config_or_path + elif isinstance(pipeline_config_or_path, PipelineConfig): + pipeline_config = pipeline_config_or_path + elif isinstance(pipeline_config_or_path, dict): + pipeline_config.update_pipeline_config(pipeline_config_or_path) + + # 4. Update PipelineConfig from CLI arguments if provided + kwargs[prefix_with_dot + 'model_path'] = model_path + pipeline_config.update_config_from_dict(kwargs, config_cli_prefix) + return pipeline_config + + def check_pipeline_config(self) -> None: + if self.vae_sp and not self.vae_tiling: + raise ValueError( + "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True." + ) + + if len(self.text_encoder_configs) != len(self.text_encoder_precisions): + raise ValueError( + f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})" + ) + + if len(self.text_encoder_configs) != len(self.preprocess_text_funcs): + raise ValueError( + f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" + ) + + if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs): + raise ValueError( + f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" + ) + + def dump_to_json(self, file_path: str): + output_dict = shallow_asdict(self) + del_keys = [] + for key, value in output_dict.items(): + if isinstance(value, ModelConfig): + model_dict = asdict(value) + # Model Arch Config should be hidden away from the users + model_dict.pop("arch_config") + output_dict[key] = model_dict + elif isinstance(value, tuple) and all( + isinstance(v, ModelConfig) for v in value): + model_dicts = [] + for v in value: + model_dict = asdict(v) + # Model Arch Config should be hidden away from the users + model_dict.pop("arch_config") + model_dicts.append(model_dict) + output_dict[key] = model_dicts + elif isinstance(value, tuple) and all(callable(f) for f in value): + # Skip dumping functions + del_keys.append(key) + + for key in del_keys: + output_dict.pop(key, None) + + with open(file_path, "w") as f: + json.dump(output_dict, f, indent=2) + + def load_from_json(self, file_path: str): + with open(file_path) as f: + input_pipeline_dict = json.load(f) + self.update_pipeline_config(input_pipeline_dict) + + def update_pipeline_config(self, source_pipeline_dict: dict[str, + Any]) -> None: + for f in fields(self): + key = f.name + if key in source_pipeline_dict: + current_value = getattr(self, key) + new_value = source_pipeline_dict[key] + + # If it's a nested ModelConfig, update it recursively + if isinstance(current_value, ModelConfig): + current_value.update_model_config(new_value) + elif isinstance(current_value, tuple) and all( + isinstance(v, ModelConfig) for v in current_value): + assert len(current_value) == len( + new_value + ), "Users shouldn't delete or add text encoder config objects in your json" + for target_config, source_config in zip(current_value, + new_value, + strict=True): + target_config.update_model_config(source_config) + else: + setattr(self, key, new_value) + + if hasattr(self, "__post_init__"): + self.__post_init__() + + +@dataclass +class SlidingTileAttnConfig(PipelineConfig): + """Configuration for sliding tile attention.""" + + # Override any BaseConfig defaults as needed + # Add sliding tile specific parameters + window_size: int = 16 + stride: int = 8 + + # You can provide custom defaults for inherited fields + height: int = 576 + width: int = 1024 + + # Additional configuration specific to sliding tile attention + pad_to_square: bool = False + use_overlap_optimization: bool = True + + +def parse_int_list(value: str) -> list[int]: + """Parse a comma-separated string of integers into a list.""" + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] diff --git a/trainer/configs/pipelines/hunyuan.py b/trainer/configs/pipelines/hunyuan.py new file mode 100644 index 0000000000000000000000000000000000000000..0346f0d45e4feb03bd3a06dac2747812324c811c --- /dev/null +++ b/trainer/configs/pipelines/hunyuan.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TypedDict + +import torch + +from trainer.configs.models import DiTConfig, EncoderConfig, VAEConfig +from trainer.configs.models.dits import HunyuanVideoConfig +from trainer.configs.models.encoders import (BaseEncoderOutput, + CLIPTextConfig, LlamaConfig) +from trainer.configs.models.vaes import HunyuanVAEConfig +from trainer.configs.pipelines.base import PipelineConfig + +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") + + +class PromptTemplate(TypedDict): + template: str + crop_start: int + + +prompt_template_video: PromptTemplate = { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, +} + + +def llama_preprocess_text(prompt: str) -> str: + return prompt_template_video["template"].format(prompt) + + +def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + hidden_state_skip_layer = 2 + assert outputs.hidden_states is not None + hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states + last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + + 1)] + crop_start = prompt_template_video.get("crop_start", -1) + last_hidden_state = last_hidden_state[:, crop_start:] + return last_hidden_state + + +def clip_preprocess_text(prompt: str) -> str: + return prompt + + +def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + pooler_output: torch.tensor = outputs.pooler_output + return pooler_output + + +@dataclass +class HunyuanConfig(PipelineConfig): + """Base configuration for HunYuan pipeline architecture.""" + + # HunyuanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig) + # Denoising stage + embedded_cfg_scale: int = 6 + flow_shift: int = 7 + + # Text encoding stage + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (LlamaConfig(), CLIPTextConfig())) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (llama_preprocess_text, clip_preprocess_text)) + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput], torch.tensor], + ...] = field(default_factory=lambda: + (llama_postprocess_text, clip_postprocess_text)) + + # Precision for each component + dit_precision: str = "bf16" + vae_precision: str = "fp16" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp16", "fp16")) + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class FastHunyuanConfig(HunyuanConfig): + """Configuration specifically optimized for FastHunyuan weights.""" + + # Override HunyuanConfig defaults + flow_shift: int = 17 + + # No need to re-specify guidance_scale or embedded_cfg_scale as they + # already have the desired values from HunyuanConfig diff --git a/trainer/configs/pipelines/registry.py b/trainer/configs/pipelines/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..8684a9668651b6bc0aba4cd8a2696447850ce7c0 --- /dev/null +++ b/trainer/configs/pipelines/registry.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry for pipeline weight-specific configurations.""" + +import os +from collections.abc import Callable + +from trainer.configs.pipelines.base import PipelineConfig +from trainer.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig +from trainer.configs.pipelines.stepvideo import StepVideoT2VConfig +from trainer.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config, + FastWan2_2_TI2V_5B_Config, + WanI2V480PConfig, WanI2V720PConfig, + WanT2V480PConfig, WanT2V720PConfig) +from trainer.logger import init_logger +from trainer.utils import (maybe_download_model_index, + verify_model_config_and_directory) + +logger = init_logger(__name__) + +# Registry maps specific model weights to their config classes +PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { + "Trainer/FastHunyuan-diffusers": FastHunyuanConfig, + "hunyuanvideo-community/HunyuanVideo": HunyuanConfig, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig, + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig, + "Trainer/FastWan2.1-T2V-1.3B-Diffusers": FastWan2_1_T2V_480P_Config, + "Trainer/FastWan2.1-T2V-14B-480P-Diffusers": FastWan2_1_T2V_480P_Config, + "Trainer/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config, + "Trainer/stepvideo-t2v-diffusers": StepVideoT2VConfig, + "Trainer/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig, + # Add other specific weight variants +} + +# For determining pipeline type from model ID +PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { + "hunyuan": lambda id: "hunyuan" in id.lower(), + "wanpipeline": lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), + "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), + "stepvideo": lambda id: "stepvideo" in id.lower(), + # Add other pipeline architecture detectors +} + +# Fallback configs when exact match isn't found but architecture is detected +PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { + "hunyuan": + HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant + "wanpipeline": + WanT2V480PConfig, # Base Wan config as fallback for any Wan variant + "wanimagetovideo": WanI2V480PConfig, + "wandmdpipeline": FastWan2_1_T2V_480P_Config, + "stepvideo": StepVideoT2VConfig + # Other fallbacks by architecture +} + + +def get_pipeline_config_cls_from_name( + pipeline_name_or_path: str) -> type[PipelineConfig]: + """Get the appropriate configuration class for a given pipeline name or path. + + This function implements a multi-step lookup process to find the most suitable + configuration class for a given pipeline. It follows this order: + 1. Exact match in the PIPE_NAME_TO_CONFIG + 2. Partial match in the PIPE_NAME_TO_CONFIG + 3. Fallback to class name in the model_index.json + 4. else raise an error + + Args: + pipeline_name_or_path (str): The name or path of the pipeline. This can be: + - A registered model ID (e.g., "Trainer/FastHunyuan-diffusers") + - A local path to a model directory + - A model ID that will be downloaded + + Returns: + Type[PipelineConfig]: The configuration class that best matches the pipeline. + This will be one of: + - A specific weight configuration class if an exact match is found + - A fallback configuration class based on the pipeline architecture + - The base PipelineConfig class if no matches are found + + Note: + - For local paths, the function will verify the model configuration + - For remote models, it will attempt to download the model index + - Warning messages are logged when falling back to less specific configurations + """ + + pipeline_config_cls: type[PipelineConfig] | None = None + + # First try exact match for specific weights + if pipeline_name_or_path in PIPE_NAME_TO_CONFIG: + pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in PIPE_NAME_TO_CONFIG.items(): + if registered_id in pipeline_name_or_path: + pipeline_config_cls = config_class + break + + # If no match, try to use the fallback config + if pipeline_config_cls is None: + if os.path.exists(pipeline_name_or_path): + config = verify_model_config_and_directory(pipeline_name_or_path) + else: + config = maybe_download_model_index(pipeline_name_or_path) + logger.warning( + "Trying to use the config from the model_index.json. Trainer may not correctly identify the optimal config for this model in this situation." + ) + + pipeline_name = config["_class_name"] + # Try to determine pipeline architecture for fallback + for pipeline_type, detector in PIPELINE_DETECTOR.items(): + if detector(pipeline_name.lower()): + pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get( + pipeline_type) + break + + if pipeline_config_cls is not None: + logger.warning( + "No match found for pipeline %s, using fallback config %s.", + pipeline_name_or_path, pipeline_config_cls) + + if pipeline_config_cls is None: + raise ValueError( + f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path." + ) + + return pipeline_config_cls diff --git a/trainer/configs/pipelines/stepvideo.py b/trainer/configs/pipelines/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..1ab2f7d1332253aea11a4986c682967fe147870e --- /dev/null +++ b/trainer/configs/pipelines/stepvideo.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.models import DiTConfig, VAEConfig +from trainer.configs.models.dits import StepVideoConfig +from trainer.configs.models.vaes import StepVideoVAEConfig +from trainer.configs.pipelines.base import PipelineConfig + + +@dataclass +class StepVideoT2VConfig(PipelineConfig): + """Base configuration for StepVideo pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=StepVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: int = 13 + timesteps_scale: bool = False + pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "bf16" diff --git a/trainer/configs/pipelines/wan.py b/trainer/configs/pipelines/wan.py new file mode 100644 index 0000000000000000000000000000000000000000..11dd2fa3aca83f0709fa1696492d2b7afe5551de --- /dev/null +++ b/trainer/configs/pipelines/wan.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from trainer.configs.models import DiTConfig, EncoderConfig, VAEConfig +from trainer.configs.models.dits import WanVideoConfig +from trainer.configs.models.encoders import (BaseEncoderOutput, + CLIPVisionConfig, T5Config) +from trainer.configs.models.vaes import WanVAEConfig +from trainer.configs.pipelines.base import PipelineConfig + + +def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + mask: torch.tensor = outputs.attention_mask + hidden_state: torch.tensor = outputs.last_hidden_state + seq_lens = mask.gt(0).sum(dim=1).long() + assert torch.isnan(hidden_state).sum() == 0 + prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] + prompt_embeds_tensor: torch.tensor = torch.stack([ + torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0) + return prompt_embeds_tensor + + +@dataclass +class WanT2V480PConfig(PipelineConfig): + """Base configuration for Wan T2V 1.3B pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=WanVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: int = 3 + + # Text encoding stage + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (T5Config(), )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], + ...] = field(default_factory=lambda: + (t5_postprocess_text, )) + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp32", )) + + # WanConfig-specific added parameters + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class WanT2V720PConfig(WanT2V480PConfig): + """Base configuration for Wan T2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 5 + + +@dataclass +class WanI2V480PConfig(WanT2V480PConfig): + """Base configuration for Wan I2V 14B 480P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Precision for each component + image_encoder_config: EncoderConfig = field( + default_factory=CLIPVisionConfig) + image_encoder_precision: str = "fp32" + + def __post_init__(self): + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class WanI2V720PConfig(WanI2V480PConfig): + """Base configuration for Wan I2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 5 + + +@dataclass +class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): + """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 8 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 757, 522]) + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): + flow_shift: int = 5 + ti2v_task: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): + flow_shift: int = 5 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 757, 522]) + + +@dataclass +class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): + pass + + +@dataclass +class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): + pass diff --git a/trainer/configs/sample/__init__.py b/trainer/configs/sample/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..825f83d3ce43a9c2874f1a99b14c3674659189dd --- /dev/null +++ b/trainer/configs/sample/__init__.py @@ -0,0 +1,3 @@ +from trainer.configs.sample.base import SamplingParam + +__all__ = ["SamplingParam"] diff --git a/trainer/configs/sample/base.py b/trainer/configs/sample/base.py new file mode 100644 index 0000000000000000000000000000000000000000..3d3b892fae75b362e96fd43cc1ec30de95c8f09b --- /dev/null +++ b/trainer/configs/sample/base.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass +from typing import Any + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class SamplingParam: + """ + Sampling parameters for video generation. + """ + # All fields below are copied from ForwardBatch + data_type: str = "video" + + # Image inputs + image_path: str | None = None + + # Text inputs + prompt: str | list[str] | None = None + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + prompt_path: str | None = None + output_path: str = "outputs/" + output_video_name: str | None = None + + # Batch info + num_videos_per_prompt: int = 1 + seed: int = 1024 + + # Original dimensions (before VAE scaling) + num_frames: int = 125 + num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus + height: int = 720 + width: int = 1280 + fps: int = 24 + + # Denoising parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_rescale: float = 0.0 + + # TeaCache parameters + enable_teacache: bool = False + + # Misc + save_video: bool = True + return_frames: bool = False + + def __post_init__(self) -> None: + self.data_type = "video" if self.num_frames > 1 else "image" + + def check_sampling_param(self): + if self.prompt_path and not self.prompt_path.endswith(".txt"): + raise ValueError("prompt_path must be a txt file") + + def update(self, source_dict: dict[str, Any]) -> None: + for key, value in source_dict.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + logger.exception("%s has no attribute %s", + type(self).__name__, key) + + self.__post_init__() + + @classmethod + def from_pretrained(cls, model_path: str) -> "SamplingParam": + from trainer.configs.sample.registry import ( + get_sampling_param_cls_for_name) + sampling_cls = get_sampling_param_cls_for_name(model_path) + if sampling_cls is not None: + sampling_param: SamplingParam = sampling_cls() + else: + logger.warning( + "Couldn't find an optimal sampling param for %s. Using the default sampling param.", + model_path) + sampling_param = cls() + + return sampling_param + + @staticmethod + def add_cli_args(parser: Any) -> Any: + """Add CLI arguments for SamplingParam fields""" + parser.add_argument( + "--prompt", + type=str, + default=SamplingParam.prompt, + help="Text prompt for video generation", + ) + parser.add_argument( + "--negative-prompt", + type=str, + default=SamplingParam.negative_prompt, + help="Negative text prompt for video generation", + ) + parser.add_argument( + "--prompt-path", + type=str, + default=SamplingParam.prompt_path, + help="Path to a text file containing the prompt", + ) + parser.add_argument( + "--output-path", + type=str, + default=SamplingParam.output_path, + help="Path to save the generated video", + ) + parser.add_argument( + "--output-video-name", + type=str, + default=SamplingParam.output_video_name, + help="Name of the output video", + ) + parser.add_argument( + "--num-videos-per-prompt", + type=int, + default=SamplingParam.num_videos_per_prompt, + help="Number of videos to generate per prompt", + ) + parser.add_argument( + "--seed", + type=int, + default=SamplingParam.seed, + help="Random seed for generation", + ) + parser.add_argument( + "--num-frames", + type=int, + default=SamplingParam.num_frames, + help="Number of frames to generate", + ) + parser.add_argument( + "--height", + type=int, + default=SamplingParam.height, + help="Height of generated video", + ) + parser.add_argument( + "--width", + type=int, + default=SamplingParam.width, + help="Width of generated video", + ) + parser.add_argument( + "--fps", + type=int, + default=SamplingParam.fps, + help="Frames per second for saved video", + ) + parser.add_argument( + "--num-inference-steps", + type=int, + default=SamplingParam.num_inference_steps, + help="Number of denoising steps", + ) + parser.add_argument( + "--guidance-scale", + type=float, + default=SamplingParam.guidance_scale, + help="Classifier-free guidance scale", + ) + parser.add_argument( + "--guidance-rescale", + type=float, + default=SamplingParam.guidance_rescale, + help="Guidance rescale factor", + ) + parser.add_argument( + "--save-video", + action="store_true", + default=SamplingParam.save_video, + help="Whether to save the video to disk", + ) + parser.add_argument( + "--no-save-video", + action="store_false", + dest="save_video", + help="Don't save the video to disk", + ) + parser.add_argument( + "--return-frames", + action="store_true", + default=SamplingParam.return_frames, + help="Whether to return the raw frames", + ) + parser.add_argument( + "--image-path", + type=str, + default=SamplingParam.image_path, + help="Path to input image for image-to-video generation", + ) + return parser + + +@dataclass +class CacheParams: + cache_type: str = "none" diff --git a/trainer/configs/sample/hunyuan.py b/trainer/configs/sample/hunyuan.py new file mode 100644 index 0000000000000000000000000000000000000000..49eb385af685ac388bfadd8764e10a877f1517e6 --- /dev/null +++ b/trainer/configs/sample/hunyuan.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.sample.base import SamplingParam +from trainer.configs.sample.teacache import TeaCacheParams + + +@dataclass +class HunyuanSamplingParam(SamplingParam): + num_inference_steps: int = 50 + + num_frames: int = 125 + height: int = 720 + width: int = 1280 + fps: int = 24 + + guidance_scale: float = 1.0 + + teacache_params: TeaCacheParams = field( + default_factory=lambda: TeaCacheParams( + teacache_thresh=0.15, + coefficients=[ + 7.33226126e+02, -4.01131952e+02, 6.75869174e+01, + -3.14987800e+00, 9.61237896e-02 + ])) + + +@dataclass +class FastHunyuanSamplingParam(HunyuanSamplingParam): + num_inference_steps: int = 6 diff --git a/trainer/configs/sample/registry.py b/trainer/configs/sample/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..31437c631dd6e51ed9ba355b178a6993617917f0 --- /dev/null +++ b/trainer/configs/sample/registry.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from collections.abc import Callable +from typing import Any + +from trainer.configs.sample.hunyuan import (FastHunyuanSamplingParam, + HunyuanSamplingParam) +from trainer.configs.sample.stepvideo import StepVideoT2VSamplingParam + +# isort: off +from trainer.configs.sample.wan import ( + FastWanT2V480PConfig, + Wan2_2_I2V_A14B_SamplingParam, + Wan2_2_T2V_A14B_SamplingParam, + Wan2_2_TI2V_5B_SamplingParam, + WanI2V_14B_480P_SamplingParam, + WanI2V_14B_720P_SamplingParam, + WanT2V_1_3B_SamplingParam, + WanT2V_14B_SamplingParam, +) +# isort: on +from trainer.logger import init_logger +from trainer.utils import (maybe_download_model_index, + verify_model_config_and_directory) + +logger = init_logger(__name__) +# Registry maps specific model weights to their config classes +SAMPLING_PARAM_REGISTRY: dict[str, Any] = { + "Trainer/FastHunyuan-diffusers": FastHunyuanSamplingParam, + "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, + "Trainer/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, + "Trainer/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + "Trainer/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, + # Add other specific weight variants +} + +# For determining pipeline type from model ID +SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { + "hunyuan": lambda id: "hunyuan" in id.lower(), + "wanpipeline": lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), + "stepvideo": lambda id: "stepvideo" in id.lower(), + # Add other pipeline architecture detectors +} + +# Fallback configs when exact match isn't found but architecture is detected +SAMPLING_FALLBACK_PARAM: dict[str, Any] = { + "hunyuan": + HunyuanSamplingParam, # Base Hunyuan config as fallback for any Hunyuan variant + "wanpipeline": + WanT2V_1_3B_SamplingParam, # Base Wan config as fallback for any Wan variant + "wanimagetovideo": WanI2V_14B_480P_SamplingParam, + "stepvideo": StepVideoT2VSamplingParam + # Other fallbacks by architecture +} + + +def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: + """Get the appropriate sampling param for specific pretrained weights.""" + + if os.path.exists(pipeline_name_or_path): + config = verify_model_config_and_directory(pipeline_name_or_path) + logger.warning( + "Trainer may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed." + ) + else: + config = maybe_download_model_index(pipeline_name_or_path) + + pipeline_name = config["_class_name"] + + # First try exact match for specific weights + if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: + return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): + if registered_id in pipeline_name_or_path: + return config_class + + # If no match, try to use the fallback config + fallback_config = None + # Try to determine pipeline architecture for fallback + for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items(): + if detector(pipeline_name.lower()): + fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type) + break + + logger.warning( + "No match found for pipeline %s, using fallback sampling param %s.", + pipeline_name_or_path, fallback_config) + return fallback_config diff --git a/trainer/configs/sample/stepvideo.py b/trainer/configs/sample/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..5404a570fc67c1d0648e1b873e55ee349c971285 --- /dev/null +++ b/trainer/configs/sample/stepvideo.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + +from trainer.configs.sample.base import SamplingParam + + +@dataclass +class StepVideoT2VSamplingParam(SamplingParam): + # Video parameters + height: int = 720 + width: int = 1280 + num_frames: int = 81 + + # Denoising stage + guidance_scale: float = 9.0 + num_inference_steps: int = 50 + + # neg magic and pos magic + # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" diff --git a/trainer/configs/sample/teacache.py b/trainer/configs/sample/teacache.py new file mode 100644 index 0000000000000000000000000000000000000000..5f92b429195c54df818da705c88f498af249c643 --- /dev/null +++ b/trainer/configs/sample/teacache.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.sample.base import CacheParams + + +@dataclass +class TeaCacheParams(CacheParams): + cache_type: str = "teacache" + teacache_thresh: float = 0.0 + coefficients: list[float] = field(default_factory=list) + + +@dataclass +class WanTeaCacheParams(CacheParams): + # Unfortunately, TeaCache is very different for Wan than other models + cache_type: str = "teacache" + teacache_thresh: float = 0.0 + use_ret_steps: bool = True + ret_steps_coeffs: list[float] = field(default_factory=list) + non_ret_steps_coeffs: list[float] = field(default_factory=list) + + @property + def coefficients(self) -> list[float]: + if self.use_ret_steps: + return self.ret_steps_coeffs + else: + return self.non_ret_steps_coeffs + + @property + def ret_steps(self) -> int: + if self.use_ret_steps: + return 5 * 2 + else: + return 1 * 2 + + def get_cutoff_steps(self, num_inference_steps: int) -> int: + if self.use_ret_steps: + return num_inference_steps * 2 + else: + return num_inference_steps * 2 - 2 \ No newline at end of file diff --git a/trainer/configs/sample/wan.py b/trainer/configs/sample/wan.py new file mode 100644 index 0000000000000000000000000000000000000000..d188767eb339110bd09992db73021f49455e133e --- /dev/null +++ b/trainer/configs/sample/wan.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from trainer.configs.sample.base import SamplingParam +from trainer.configs.sample.teacache import WanTeaCacheParams + + +@dataclass +class WanT2V_1_3B_SamplingParam(SamplingParam): + # Video parameters + height: int = 480 + width: int = 832 + num_frames: int = 81 + fps: int = 16 + + # Denoising stage + guidance_scale: float = 3.0 + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + num_inference_steps: int = 50 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.08, + ret_steps_coeffs=[ + -5.21862437e+04, 9.23041404e+03, -5.28275948e+02, + 1.36987616e+01, -4.99875664e-02 + ], + non_ret_steps_coeffs=[ + 2.39676752e+03, -1.31110545e+03, 2.01331979e+02, + -8.29855975e+00, 1.37887774e-01 + ])) + + +@dataclass +class WanT2V_14B_SamplingParam(SamplingParam): + # Video parameters + height: int = 720 + width: int = 1280 + num_frames: int = 81 + fps: int = 16 + + # Denoising stage + guidance_scale: float = 5.0 + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + num_inference_steps: int = 50 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.20, + use_ret_steps=False, + ret_steps_coeffs=[ + -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, + 5.87365115e+01, -3.15583525e-01 + ], + non_ret_steps_coeffs=[ + -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, + -13.02252404 + ])) + + +@dataclass +class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam): + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 40 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.26, + ret_steps_coeffs=[ + -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, + 5.87365115e+01, -3.15583525e-01 + ], + non_ret_steps_coeffs=[ + -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, + -13.02252404 + ])) + + +@dataclass +class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam): + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 40 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.3, + ret_steps_coeffs=[ + -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, + 5.87365115e+01, -3.15583525e-01 + ], + non_ret_steps_coeffs=[ + -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, + -13.02252404 + ])) + + +@dataclass +class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam): + # DMD parameters + # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) + num_inference_steps: int = 3 + num_frames: int = 61 + height: int = 448 + width: int = 832 + fps: int = 16 + + +# ============================================= +# ============= Wan2.2 TI2V Models ============= +# ============================================= +@dataclass +class Wan2_2_Base_SamplingParam(SamplingParam): + """Sampling parameters for Wan2.2 TI2V 5B model.""" + negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + + +@dataclass +class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam): + """Sampling parameters for Wan2.2 TI2V 5B model.""" + height: int = 704 + width: int = 1280 + num_frames: int = 121 + fps: int = 24 + guidance_scale: float = 5.0 + num_inference_steps: int = 50 + + +@dataclass +class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): + guidance_scale: float = 4.0 + guidance_scale_2: float = 3.0 + num_inference_steps: int = 40 + fps: int = 16 + + +@dataclass +class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): + guidance_scale: float = 3.5 + guidance_scale_2: float = 3.5 + num_inference_steps: int = 40 + fps: int = 16 diff --git a/trainer/configs/utils.py b/trainer/configs/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a2d47b306776c1150c644324b1434113162cbf69 --- /dev/null +++ b/trainer/configs/utils.py @@ -0,0 +1,61 @@ +import argparse +from typing import Any + + +def update_config_from_args(config: Any, + args_dict: dict[str, Any], + prefix: str = "", + pop_args: bool = False) -> bool: + """ + Update configuration object from arguments dictionary. + + Args: + config: The configuration object to update + args_dict: Dictionary containing arguments + prefix: Prefix for the configuration parameters in the args_dict. + If None, assumes direct attribute mapping without prefix. + """ + # Handle top-level attributes (no prefix) + args_not_to_remove = [ + 'model_path', + ] + args_to_remove = [] + if prefix.strip() == "": + for key, value in args_dict.items(): + if hasattr(config, key) and value is not None: + if key == "text_encoder_precisions" and isinstance(value, list): + setattr(config, key, tuple(value)) + else: + setattr(config, key, value) + if pop_args: + args_to_remove.append(key) + else: + # Handle nested attributes with prefix + prefix_with_dot = f"{prefix}." + for key, value in args_dict.items(): + if key.startswith(prefix_with_dot) and value is not None: + attr_name = key[len(prefix_with_dot):] + if hasattr(config, attr_name): + setattr(config, attr_name, value) + if pop_args: + args_to_remove.append(key) + + if pop_args: + for key in args_to_remove: + if key not in args_not_to_remove: + args_dict.pop(key) + + return len(args_to_remove) > 0 + + +def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]: + """ + Clean the arguments by removing the ones that not explicitly provided by the user. + """ + provided_args = {} + for k, v in vars(args).items(): + if (v is not None and hasattr(args, '_provided') + and k in args._provided): + provided_args[k] = v + + return provided_args diff --git a/trainer/configs/wan_1.3B_t2v_pipeline.json b/trainer/configs/wan_1.3B_t2v_pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..2ed1b3bfeb6426a9b854b2ea0541bce77beb7d10 --- /dev/null +++ b/trainer/configs/wan_1.3B_t2v_pipeline.json @@ -0,0 +1,41 @@ +{ + "embedded_cfg_scale": 6.0, + "flow_shift": 3, + "dit_cpu_offload": true, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": false, + "vae_sp": false, + "vae_config": { + "load_encoder": false, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 8, + "use_tiling": false, + "use_temporal_tiling": false, + "use_parallel_tiling": false, + "use_feature_cache": true + }, + "dit_config": { + "prefix": "Wan", + "quant_config": null + }, + "text_encoder_precisions": [ + "fp32" + ], + "text_encoder_configs": [ + { + "prefix": "t5", + "quant_config": null, + "lora_config": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false +} \ No newline at end of file diff --git a/trainer/configs/wan_14B_i2v_480p_pipeline.json b/trainer/configs/wan_14B_i2v_480p_pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..21743d66b3621375e4597b0ba4a9662c71a15df0 --- /dev/null +++ b/trainer/configs/wan_14B_i2v_480p_pipeline.json @@ -0,0 +1,49 @@ +{ + "embedded_cfg_scale": 6.0, + "flow_shift": 3, + "dit_cpu_offload": true, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": false, + "vae_sp": false, + "vae_config": { + "load_encoder": true, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 8, + "use_tiling": false, + "use_temporal_tiling": false, + "use_parallel_tiling": false, + "use_feature_cache": true + }, + "dit_config": { + "prefix": "Wan", + "quant_config": null + }, + "text_encoder_precisions": [ + "fp32" + ], + "text_encoder_configs": [ + { + "prefix": "t5", + "quant_config": null, + "lora_config": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false, + "image_encoder_config": { + "prefix": "clip", + "quant_config": null, + "lora_config": null, + "num_hidden_layers_override": null, + "require_post_norm": null + }, + "image_encoder_precision": "fp32" +} \ No newline at end of file diff --git a/trainer/dataset/__init__.py b/trainer/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..607c6198e03a466b7a18fd8d45194c6cf036fdad --- /dev/null +++ b/trainer/dataset/__init__.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +from torchvision import transforms +from torchvision.transforms import Lambda + +from trainer.dataset.ar_camera_hunyuan_w_mem_dataset import build_ar_camera_hunyuan_w_mem_dataloader +from trainer.dataset.preprocessing_datasets import VideoCaptionMergedDataset +from trainer.dataset.transform import (CenterCropResizeVideo, Normalize255, + TemporalRandomCrop) +from trainer.dataset.validation_dataset import ValidationDataset + + +def getdataset(args) -> VideoCaptionMergedDataset: + if args.do_temporal_sample: + temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x + else: + temporal_sample = None + norm_fun = Lambda(lambda x: 2.0 * x - 1.0) + resize_topcrop = [ + CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), + ] + resize = [ + CenterCropResizeVideo((args.max_height, args.max_width)), + ] + transform = transforms.Compose([ + # Normalize255(), + *resize, + ]) + transform_topcrop = transforms.Compose([ + Normalize255(), + *resize_topcrop, + norm_fun, + ]) + return VideoCaptionMergedDataset(data_merge_path=args.data_merge_path, + args=args, + transform=transform, + temporal_sample=temporal_sample, + transform_topcrop=transform_topcrop, + seed=args.seed) + + +__all__ = [ + "build_ar_camera_hunyuan_w_mem_dataloader", +] diff --git a/trainer/dataset/ar_camera_hunyuan_w_mem_dataset.py b/trainer/dataset/ar_camera_hunyuan_w_mem_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2cdf4b93853403f8350f8065dc2d3e194eab4753 --- /dev/null +++ b/trainer/dataset/ar_camera_hunyuan_w_mem_dataset.py @@ -0,0 +1,738 @@ +import json +import os +import sys + +sys.path.append(os.path.abspath('.')) +import torch +import pandas as pd +import numpy as np +import random +from pathlib import Path +from typing import List, Tuple, Dict +from scipy.spatial.transform import Rotation as R_scipy +import math +import torch.multiprocessing as mp + +from torch.utils.data import Dataset, Sampler +from torchdata.stateful_dataloader import StatefulDataLoader +from scipy.spatial.transform import (Rotation + as R) +from trainer.distributed import ( + get_local_torch_device, + ) +from trainer.distributed import (get_sp_world_size, + get_world_rank, + get_world_size) + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +def generate_points_in_sphere(n_points: int, radius: float) -> torch.Tensor: + """ + Uniformly sample points within a sphere of a specified radius. + + :param n_points: The number of points to generate. + :param radius: The radius of the sphere. + :return: A tensor of shape (n_points, 3), representing the (x, y, z) coordinates of the points. + """ + samples_r = torch.rand(n_points) + samples_phi = torch.rand(n_points) + samples_u = torch.rand(n_points) + + r = radius * torch.pow(samples_r, 1 / 3) + phi = 2 * math.pi * samples_phi + theta = torch.acos(1 - 2 * samples_u) + + # transfer the coordinates from spherical to cartesian + x = r * torch.sin(theta) * torch.cos(phi) + y = r * torch.sin(theta) * torch.sin(phi) + z = r * torch.cos(theta) + + points = torch.stack((x, y, z), dim=1) + return points + + +def rotation_matrix_to_angles(R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Estimate the Pitch and Yaw angles from a 3x3 rotation matrix R in the camera coordinate system. + + Assumed Camera Coordinate System: X=Right, Y=Up, Z=Backward + (or NeRF style: X=Right, Y=Down, Z=Forward). + Here we adopt the common Computer Vision convention: Z-axis is Forward. + + Note: The angle calculations here are directly based on the conventions of your `is_inside_fov_3d_hv` function: + - Yaw/Azimuth angle is in the XZ plane (atan2(x, z)). + - Pitch/Elevation angle is relative to the horizontal plane (atan2(y, sqrt(x^2 + z^2))). + + For the third column R[:, 2] of the W2C matrix R (the direction of the World Z-axis in the Camera frame), + this typically corresponds to the direction the camera is looking + (the representation of the world Z-axis in the camera frame). + + To simplify and match your `is_inside_fov` logic, we directly use the camera's Z-axis vector: + Camera Z-axis direction in World Frame (Forward Vector): fwd = R_w2c_inv @ [0, 0, 1] + More simply, the Z-axis vector of the C2W matrix is the camera's forward vector in the world frame. + C2W = W2C_inv + """ + + R_c2w = R.T + fwd = R_c2w[:, 2] + + x = fwd[0] + y = fwd[1] + z = fwd[2] + + # compute yaw and pitch + yaw_rad = torch.atan2(x, z) + yaw_deg = yaw_rad * (180.0 / math.pi) + pitch_rad = torch.atan2(y, torch.sqrt(x ** 2 + z ** 2)) + pitch_deg = pitch_rad * (180.0 / math.pi) + + return pitch_deg, yaw_deg + +def is_inside_fov_3d_hv(points: torch.Tensor, center: torch.Tensor, + center_pitch: torch.Tensor, center_yaw: torch.Tensor, + fov_half_h: torch.Tensor, fov_half_v: torch.Tensor) -> torch.Tensor: + """ + Check whether points are inside a 3D view frustum defined by a center coordinate, pitch angle, and yaw angle. + + :param points: Tensor of shape (N, 3) or (N, B, 3) representing the coordinates of the sampled points. + :param center: Tensor of shape (3) or (B, 3) representing the camera center coordinates. + :param center_pitch: Tensor of shape (1) or (B) representing the pitch angle of center view direction. + :param center_yaw: Tensor of shape (1) or (B) representing the yaw angle of the center view direction. + :param fov_half_h: The horizontal half field-of-view angle (in degrees). + :param fov_half_v: The vertical half field-of-view angle (in degrees). + :return: Boolean tensor of shape (N) or (N, B), indicating whether each point is inside the FOV. + """ + if points.ndim == 2: # N, 3 + vectors = points - center[None, :] + C = 1 + elif points.ndim == 3: # N, B, 3 + vectors = points - center[None, ...] + center_pitch = center_pitch[None, :] if center_pitch.ndim == 1 else center_pitch + center_yaw = center_yaw[None, :] if center_yaw.ndim == 1 else center_yaw + else: + raise ValueError("points' shape should be (N, 3) or (N, B, 3)") + + x = vectors[..., 0] + y = vectors[..., 1] + z = vectors[..., 2] + + # Calculate the horizontal angle (yaw/azimuth), assuming the Z-axis is forward. + azimuth = torch.atan2(x, z) * (180 / math.pi) + + # Calculate the vertical angle (pitch/elevation). + elevation = torch.atan2(y, torch.sqrt(x ** 2 + z ** 2)) * (180 / math.pi) + + # Calculate the angular difference from the center view direction (handling angle wrapping). + diff_azimuth = (azimuth - center_yaw) + diff_azimuth = torch.remainder(diff_azimuth + 180, 360) - 180 + + diff_elevation = (elevation - center_pitch) + diff_elevation = torch.remainder(diff_elevation + 180, 360) - 180 + + # Check if within FOV + in_fov_h = diff_azimuth.abs() < fov_half_h + in_fov_v = diff_elevation.abs() < fov_half_v + + return in_fov_h & in_fov_v + + +def calculate_fov_overlap_similarity( + w2c_matrix_curr: torch.Tensor, + w2c_matrix_hist: torch.Tensor, + fov_h_deg: float = 105.0, + fov_v_deg: float = 75.0, + device=None, + points_local=None, +) -> float: + """ + Calculate the Field-of-View (FOV) overlap similarity between two W2C poses using Monte Carlo sampling. + + Similarity = (Number of points in Curr_FOV ∩ Hist_FOV) / (Number of points in Curr_FOV). + + :param w2c_matrix_curr: The (4, 4) W2C matrix for the current frame. + :param w2c_matrix_hist: The (4, 4) W2C matrix for the historical frame. + :param num_samples, radius, fov_h_deg, fov_v_deg: Sampling and FOV parameters. + :return: The overlap ratio (a float between 0.0 and 1.0). + """ + w2c_matrix_curr = torch.tensor(w2c_matrix_curr, device=device) + w2c_matrix_hist = torch.tensor(w2c_matrix_hist, device=device) + + c2w_matrix_curr = torch.linalg.inv(w2c_matrix_curr) + c2w_matrix_hist = torch.linalg.inv(w2c_matrix_hist) + C_inv = w2c_matrix_curr + + w2c_matrix_curr = torch.linalg.inv(C_inv @ c2w_matrix_curr) + w2c_matrix_hist = torch.linalg.inv(C_inv @ c2w_matrix_hist) + + R_curr, t_curr = w2c_matrix_curr[:3, :3], w2c_matrix_curr[:3, 3] + R_hist, t_hist = w2c_matrix_hist[:3, :3], w2c_matrix_hist[:3, 3] + P_w_curr = -R_curr.T @ t_curr + P_w_hist = -R_hist.T @ t_hist + + # pitch, yaw + pitch_curr, yaw_curr = rotation_matrix_to_angles(R_curr) + pitch_hist, yaw_hist = rotation_matrix_to_angles(R_hist) + + fov_half_h = torch.tensor(fov_h_deg / 2.0, device=device) + fov_half_v = torch.tensor(fov_v_deg / 2.0, device=device) + + # move to P_w_curr (N, 3) + points_world = points_local + P_w_curr[None, :] + + in_fov_curr = is_inside_fov_3d_hv( + points_world, P_w_curr[None, :], + pitch_curr[None], yaw_curr[None], + fov_half_h, fov_half_v + ) + + # compute based on angle + in_fov_hist = is_inside_fov_3d_hv( + points_world, P_w_hist[None, :], + pitch_hist[None], yaw_hist[None], + fov_half_h, fov_half_v + ) + + # compute based on distance + dist = torch.norm(points_world - P_w_hist.reshape(1, -1), dim=1) < 8.0 + in_fov_hist = in_fov_hist.bool() & dist.reshape(1, -1).bool() + + overlap_count = (in_fov_curr.bool() & in_fov_hist.bool()).sum().float() + fov_curr_count = in_fov_curr.sum().float() + + if fov_curr_count == 0: + return 0.0 + + overlap_ratio = overlap_count / fov_curr_count + + return overlap_ratio.item() + + +def select_aligned_memory_frames( + w2c_list: List[np.ndarray], current_frame_idx: int, + memory_frames: int, temporal_context_size: int, + pred_latent_size: int, pos_weight: float = 1.0, + ang_weight: float = 1.0, device=None, + points_local=None) -> List[int]: + """ + Selects memory and context frames for a given frame based on a four-frame segment distance calculation. + + :param w2c_list: List of all N 4x4 World-to-Camera (W2C) extrinsic matrices (np.ndarray). + :param current_frame_idx: The index of the current frame to be processed. + :param memory_frames: The total number of memory frames to select. + :param context_size: The total number of context frames to select. + :param pos_weight: The weight applied to the spatial (position) distance component. + :param ang_weight: The weight applied to the angular distance component. + + :return: List[int]: A list containing the indices of the selected memory frames and context frames. + """ + if current_frame_idx <= memory_frames: + return list(range(0, current_frame_idx)) + + num_total_frames = len(w2c_list) + if current_frame_idx >= num_total_frames or current_frame_idx < 3: + raise ValueError( + f"The current frame index must be within the valid range of w2c_list and must be at least 3." + f"{current_frame_idx}, {len(w2c_list)}" + ) + + start_context_idx = max(0, current_frame_idx - temporal_context_size) + context_frames_indices = list(range(start_context_idx, current_frame_idx)) + + candidate_distances = [] + query_clip_indices = list( + range( + current_frame_idx, + current_frame_idx + pred_latent_size + if current_frame_idx + pred_latent_size <= num_total_frames + else num_total_frames + ) + ) + + historical_clip_indices = list(range(4, current_frame_idx - temporal_context_size, 4)) + + memory_frames_indices = [0,1,2,3] # add the first chunk as context + memory_frames = memory_frames - temporal_context_size + + for hist_idx in historical_clip_indices: + total_dist = 0 + hist_w2c_1 = w2c_list[hist_idx] + hist_w2c_2 = w2c_list[hist_idx + 2] + for query_idx in query_clip_indices: + dist_1_for_query_idx = 1.0 - calculate_fov_overlap_similarity(w2c_list[query_idx], hist_w2c_1, + fov_h_deg=60.0, fov_v_deg=35.0, + device=device, points_local=points_local) + dist_2_for_query_idx = 1.0 - calculate_fov_overlap_similarity(w2c_list[query_idx], hist_w2c_2, + fov_h_deg=60.0, fov_v_deg=35.0, + device=device, points_local=points_local) + dist_for_query_idx = (dist_1_for_query_idx + dist_2_for_query_idx) / 2.0 + total_dist += dist_for_query_idx + + final_clip_distance = total_dist / len(query_clip_indices) + candidate_distances.append((hist_idx, final_clip_distance)) + + candidate_distances.sort(key=lambda x: x[1]) + + for start_idx, _ in candidate_distances: + # check the memory frame number + if len(memory_frames_indices) >= memory_frames: + break + + if start_idx not in memory_frames_indices: + memory_frames_indices.extend(range(start_idx, start_idx + 4)) + + # exclude the repeated frames + selected_frames_set = set(context_frames_indices) + selected_frames_set.update(memory_frames_indices) + + final_selected_frames = sorted(list(selected_frames_set)) + + return final_selected_frames + + +class DP_SP_BatchSampler(Sampler[list[int]]): + """ + A simple sequential batch sampler that yields batches of indices. + """ + + def __init__( + self, + batch_size: int, + dataset_size: int, + num_sp_groups: int, + sp_world_size: int, + global_rank: int, + drop_last: bool = True, + drop_first_row: bool = False, + seed: int = 0, + ): + self.batch_size = batch_size + self.dataset_size = dataset_size + self.drop_last = drop_last + self.seed = seed + self.num_sp_groups = num_sp_groups + self.global_rank = global_rank + self.sp_world_size = sp_world_size + + # ── epoch-level RNG ──────────────────────────────────────────────── + rng = torch.Generator().manual_seed(self.seed) + # Create a random permutation of all indices + global_indices = torch.randperm(self.dataset_size, generator=rng) + + if drop_first_row: + # drop 0 in global_indices + global_indices = global_indices[global_indices != 0] + self.dataset_size = self.dataset_size - 1 + + if self.drop_last: + # For drop_last=True, we: + # 1. Ensure total samples is divisible by (batch_size * num_sp_groups) + # 2. This guarantees each SP group gets same number of complete batches + # 3. Prevents uneven batch sizes across SP groups at end of epoch + num_batches = self.dataset_size // self.batch_size + num_global_batches = num_batches // self.num_sp_groups + global_indices = global_indices[:num_global_batches * + self.num_sp_groups * + self.batch_size] + else: + if self.dataset_size % (self.num_sp_groups * self.batch_size) != 0: + # add more indices to make it divisible by (batch_size * num_sp_groups) + padding_size = self.num_sp_groups * self.batch_size - ( + self.dataset_size % (self.num_sp_groups * self.batch_size)) + logger.info("Padding the dataset from %d to %d", + self.dataset_size, self.dataset_size + padding_size) + global_indices = torch.cat( + [global_indices, global_indices[:padding_size]]) + + # shard the indices to each sp group + ith_sp_group = self.global_rank // self.sp_world_size + sp_group_local_indices = global_indices[ith_sp_group::self. + num_sp_groups] + self.sp_group_local_indices = sp_group_local_indices + logger.info("Dataset size for each sp group: %d", + len(sp_group_local_indices)) + + def __iter__(self): + indices = self.sp_group_local_indices + for i in range(0, len(indices), self.batch_size): + batch_indices = indices[i:i + self.batch_size] + yield batch_indices.tolist() + + def __len__(self): + return len(self.sp_group_local_indices) // self.batch_size + +class CameraJsonWMemDataset(Dataset): + def __init__(self, json_path, causal, window_frames, batch_size, cfg_rate, i2v_rate, drop_last, drop_first_row, + seed, device, shared_state): + self.json_data = json.load(open(json_path, 'r')) + self.all_length = len(self.json_data) + self.causal = causal + self.window_frames = window_frames + self.memory_frames = 20 + self.cfg_rate = cfg_rate + self.rng = random.Random(seed) + self.i2v_rate = i2v_rate + self.device = device + self.shared_state = shared_state + + self.sampler = DP_SP_BatchSampler( + batch_size=batch_size, + dataset_size=self.all_length, + num_sp_groups=get_world_size() // get_sp_world_size(), + sp_world_size=get_sp_world_size(), + global_rank=get_world_rank(), + drop_last=drop_last, + drop_first_row=drop_first_row, + seed=seed, + ) + + self.mapping = { + (0, 0, 0, 0): 0, + (1, 0, 0, 0): 1, + (0, 1, 0, 0): 2, + (0, 0, 1, 0): 3, + (0, 0, 0, 1): 4, + (1, 0, 1, 0): 5, + (1, 0, 0, 1): 6, + (0, 1, 1, 0): 7, + (0, 1, 0, 1): 8, + } + + self.points_local = generate_points_in_sphere(50000, 8.0).to(device) + + self.neg_prompt_pt = torch.load( + "/your_path/to/hunyuan_neg_prompt.pt", + map_location="cpu", + weights_only=True, + ) + + self.neg_byt5_pt = torch.load( + "/your_path/to/hunyuan_neg_byt5_prompt.pt", + map_location="cpu", + weights_only=True, + ) + + def __len__(self): + return self.all_length + + def camera_center_normalization(self, w2c): + c2w = np.linalg.inv(w2c) + C0_inv = np.linalg.inv(c2w[0]) + c2w_aligned = np.array([C0_inv @ C for C in c2w]) + return np.linalg.inv(c2w_aligned) + + def one_hot_to_one_dimension(self, one_hot): + y = torch.tensor([self.mapping[tuple(row.tolist())] for row in one_hot]) + return y + + def perturb_pose(self, w2c_list, current_frame_idx, generate_length): + total_latent_num = w2c_list.shape[0] + context_w2c = w2c_list[:current_frame_idx] + + pertubed_left = self.rng.random() < 0.5 + if pertubed_left: + generate_w2c = w2c_list[current_frame_idx - 4: current_frame_idx - 4 + generate_length] + else: + if total_latent_num >= current_frame_idx + 4 + generate_length: + generate_w2c = w2c_list[current_frame_idx + 4: current_frame_idx + 4 + generate_length] + else: + generate_w2c = w2c_list[current_frame_idx: current_frame_idx + generate_length] + + return np.concatenate((context_w2c, generate_w2c), axis=0) + + def update_max_frames(self, training_step): + if training_step < 500: + self.shared_state["max_frames"] = 32 + elif training_step < 1000: + self.shared_state["max_frames"] = 64 + elif training_step < 2000: + self.shared_state["max_frames"] = 96 + elif training_step < 3000: + self.shared_state["max_frames"] = 128 + else: + self.shared_state["max_frames"] = 160 + + def __getitem__(self, idx): + while True: + try: + json_data = self.json_data[idx] + latent_pt_path = json_data['latent_path'] + pose_path = json_data['pose_path'] + + latent_pt = torch.load( + os.path.join(latent_pt_path), + map_location="cpu", + weights_only=True, + ) + latent = latent_pt['latent'][0] + latent_length = latent.shape[1] + + # if self.causal: + if latent_length < self.window_frames: + idx = self.rng.randint(0, self.all_length - 1) + continue + else: + max_frames = int(self.shared_state["max_frames"]) // 4 * 4 + max_length = min(max_frames, latent_length // 4 * 4) + + latent = latent[:, :max_length, ...] # mostly larger than window_frames, need chunk + + prompt_embed = latent_pt['prompt_embeds'][0] + prompt_mask = latent_pt['prompt_mask'][0] + + image_cond = latent_pt['image_cond'][0] + vision_states = latent_pt['vision_states'][0] + byt5_text_states = latent_pt['byt5_text_states'][0] + byt5_text_mask = latent_pt['byt5_text_mask'][0] + + if self.rng.random() < self.cfg_rate: + prompt_embed = self.neg_prompt_pt['negative_prompt_embeds'][0] + prompt_mask = self.neg_prompt_pt['negative_prompt_mask'][0] + byt5_text_states = self.neg_byt5_pt['byt5_text_states'][0] + byt5_text_mask = self.neg_byt5_pt['byt5_text_mask'][0] + + pose_json = json.load(open(pose_path, 'r')) + pose_keys = list(pose_json.keys()) + intrinsic_list = [] + w2c_list = [] + for i in range(latent.shape[1]): + t_key = pose_keys[0] if i == 0 else pose_keys[4 * (i - 1) + 4] + intrinsic = np.array(pose_json[t_key]['intrinsic']) + w2c = np.array(pose_json[t_key]['w2c']) + + intrinsic[0, 0] /= intrinsic[0, 2] * 2 + intrinsic[1, 1] /= intrinsic[1, 2] * 2 + intrinsic[0, 2] = 0.5 + intrinsic[1, 2] = 0.5 + w2c_list.append(w2c) + intrinsic_list.append(intrinsic) + + w2c_list = np.array(w2c_list) + w2c_list = self.camera_center_normalization(w2c_list) + intrinsic_list = torch.tensor(np.array(intrinsic_list)) + + if 'latent_dataset_w_action' in latent_pt_path: # prepare for dataset with action labels + trans_one_hot = np.zeros((intrinsic_list.shape[0], 4), dtype=np.int32) + rotate_one_hot = np.zeros((intrinsic_list.shape[0], 4), dtype=np.int32) + action_json = json.load(open(json_data["action_path"], 'r')) + action_keys = list(action_json.keys()) + for action_idx in range(1, trans_one_hot.shape[0]): + t_key = action_keys[4 * (action_idx - 1) + 4] + t_move_action = action_json[t_key]["move_action"] + t_view_action = action_json[t_key]["view_action"] + if "W" in t_move_action and "S" not in t_move_action: + trans_one_hot[action_idx, 0] = 1 + if "S" in t_move_action and "W" not in t_move_action: + trans_one_hot[action_idx, 1] = 1 + if "D" in t_move_action and "A" not in t_move_action: + trans_one_hot[action_idx, 2] = 1 + if "A" in t_move_action and "D" not in t_move_action: + trans_one_hot[action_idx, 3] = 1 + + if t_view_action == "LR": + rotate_one_hot[action_idx, 0] = 1 + elif t_view_action == "LL": + rotate_one_hot[action_idx, 1] = 1 + elif t_view_action == "LU": + rotate_one_hot[action_idx, 2] = 1 + elif t_view_action == "LD": + rotate_one_hot[action_idx, 3] = 1 + + trans_one_label = self.one_hot_to_one_dimension(trans_one_hot) + rotate_one_label = self.one_hot_to_one_dimension(rotate_one_hot) + action_for_pe = trans_one_label * 9 + rotate_one_label + + else: # prepare action labels on the fly + c2ws = np.linalg.inv(w2c_list) + C_inv = np.linalg.inv(c2ws[:-1]) + relative_c2w = np.zeros_like(c2ws) + relative_c2w[0, ...] = c2ws[0, ...] + relative_c2w[1:, ...] = C_inv @ c2ws[1:, ...] + trans_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) + rotate_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) + + move_norm_valid = 0.01 + for i in range(1, relative_c2w.shape[0]): + move_dirs = relative_c2w[i, :3, 3] + move_norms = np.linalg.norm(move_dirs) + # compute translation angles + if move_norms > move_norm_valid: + move_norm_dirs = move_dirs / move_norms + angles_rad = np.arccos(move_norm_dirs.clip(-1.0, 1.0)) + trans_angles_deg = angles_rad * (180.0 / torch.pi) + + if trans_angles_deg[2] < 60: + trans_one_hot[i, 0] = 1 + elif trans_angles_deg[2] > 120: + trans_one_hot[i, 1] = 1 + + if trans_angles_deg[0] < 60: + trans_one_hot[i, 2] = 1 + elif trans_angles_deg[0] > 120: + trans_one_hot[i, 3] = 1 + else: + trans_angles_deg = torch.zeros(3) + + R_rel = relative_c2w[i, :3, :3] + r = R.from_matrix(R_rel) + rot_angles_deg = r.as_euler('xyz', degrees=True) + + # compute rotation angles + if rot_angles_deg[1] > 5e-2: + rotate_one_hot[i, 0] = 1 + elif rot_angles_deg[1] < -5e-2: + rotate_one_hot[i, 1] = 1 + + if rot_angles_deg[0] > 5e-2: + rotate_one_hot[i, 2] = 1 + elif rot_angles_deg[0] < -5e-2: + rotate_one_hot[i, 3] = 1 + + trans_one_hot = torch.tensor(trans_one_hot) + rotate_one_hot = torch.tensor(rotate_one_hot) + + trans_one_label = self.one_hot_to_one_dimension(trans_one_hot) + rotate_one_label = self.one_hot_to_one_dimension(rotate_one_hot) + action_for_pe = trans_one_label * 9 + rotate_one_label + + select_window_out_flag = 0 # whether to select the latents with length > window_frames + select_prob = self.rng.random() + + if select_prob < 0.8: + select_window_out_flag = 1 # mean to select frames outside the window + max_index = latent.shape[1] - (self.window_frames - self.memory_frames) + + start_chunk_id = (self.window_frames) // 4 + end_chunk_id = max_index // 4 + current_frame_idx = self.rng.randint(start_chunk_id, end_chunk_id) * 4 # include the left and right + + # -------------------- for ar, only search the memory for the current chunk + selected_history_frame_id = select_aligned_memory_frames(w2c_list, + current_frame_idx, + memory_frames=self.memory_frames, + temporal_context_size=12, + pred_latent_size=4, + points_local=self.points_local, + device=self.device) # align the training objective: refine the fov selection + selected_history_frame_id.extend(range(current_frame_idx, current_frame_idx + 4)) + latent = latent[:, selected_history_frame_id] + reset_w2c_list = w2c_list[selected_history_frame_id] + w2c_list = reset_w2c_list + reset_intrinsic_list = intrinsic_list[selected_history_frame_id] + intrinsic_list = reset_intrinsic_list + reset_action_for_pe = action_for_pe[selected_history_frame_id] + action_for_pe = reset_action_for_pe + + else: + pred_latent_size = self.window_frames + latent = latent[:, :pred_latent_size, ...] + w2c_list = w2c_list[:pred_latent_size] + intrinsic_list = intrinsic_list[:pred_latent_size] + action_for_pe = action_for_pe[:pred_latent_size] + + i2v_mask = torch.ones_like(latent) + + batch = { + "i2v_mask": i2v_mask, + "latent": latent, + "prompt_embed": prompt_embed, + "w2c": torch.tensor(w2c_list), + "intrinsic": intrinsic_list, + "action": action_for_pe, + "action_for_pe": action_for_pe, + "context_frames_list": None, # selected context frames for each chunk + "select_window_out_flag": select_window_out_flag, # select frames outside the window or not + "video_path": json_data["pose_path"], + "max_length": max_frames, + + "image_cond": image_cond, + "vision_states": vision_states, + "prompt_mask": prompt_mask, + "byt5_text_states": byt5_text_states, + "byt5_text_mask": byt5_text_mask, + } + break + except Exception as e: + print('error:', e, latent_pt_path, flush=True) + idx = self.rng.randint(0, self.all_length - 1) + return batch + + +def cycle(dl): + while True: + for data in dl: + yield data + + +def latent_collate_function(batch): + latent = torch.stack([b["latent"] for b in batch], dim=0) + prompt_embed = torch.stack([b["prompt_embed"] for b in batch], dim=0) + w2c = torch.stack([b["w2c"] for b in batch], dim=0) + intrinsic = torch.stack([b["intrinsic"] for b in batch], dim=0) + action = torch.stack([b["action"] for b in batch], dim=0) + action_for_pe = torch.stack([b["action_for_pe"] for b in batch], dim=0) + i2v_mask = torch.stack([b["i2v_mask"] for b in batch], dim=0) + + image_cond = torch.stack([b["image_cond"] for b in batch], dim=0) + vision_states = torch.stack([b["vision_states"] for b in batch], dim=0) + prompt_mask = torch.stack([b["prompt_mask"] for b in batch], dim=0) + byt5_text_states = torch.stack([b["byt5_text_states"] for b in batch], dim=0) + byt5_text_mask = torch.stack([b["byt5_text_mask"] for b in batch], dim=0) + + + context_frames_list = [b["context_frames_list"] for b in batch] + select_window_out_flag = [b["select_window_out_flag"] for b in batch] + video_path = [b["video_path"] for b in batch] + max_length = [b["max_length"] for b in batch] + + return { + "i2v_mask": i2v_mask, + "latent": latent, + "prompt_embed": prompt_embed, + "w2c": w2c, + "intrinsic": intrinsic, + "action": action, + "video_path": video_path, + "context_frames_list": context_frames_list, + "select_window_out_flag": select_window_out_flag, + "action_for_pe": action_for_pe, + "max_length": max_length, + "image_cond": image_cond, + "vision_states": vision_states, + "prompt_mask": prompt_mask, + "byt5_text_states": byt5_text_states, + "byt5_text_mask": byt5_text_mask, + } + + +def build_ar_camera_hunyuan_w_mem_dataloader( + json_path, + causal, + window_frames, + # memory_frames, + batch_size, + num_data_workers, + drop_last, + drop_first_row, + seed, + cfg_rate, + i2v_rate, ) -> tuple[CameraJsonWMemDataset, StatefulDataLoader]: + manager = mp.Manager() + shared_state = manager.dict() + shared_state["max_frames"] = window_frames + + dataset = CameraJsonWMemDataset(json_path, causal, window_frames, batch_size, cfg_rate, i2v_rate, + drop_last=drop_last, drop_first_row=drop_first_row, seed=seed, + device=get_local_torch_device(), shared_state=shared_state) + + loader = StatefulDataLoader( + dataset, + batch_sampler=dataset.sampler, + collate_fn=latent_collate_function, + num_workers=num_data_workers, + pin_memory=True, + persistent_workers=num_data_workers > 0, + ) + return dataset, loader diff --git a/trainer/dataset/dataloader/schema.py b/trainer/dataset/dataloader/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..ed00fd74bf173ee5fab0cb3cc1747447c2258496 --- /dev/null +++ b/trainer/dataset/dataloader/schema.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# schema.py +""" +Unified data schema and format for saving and loading image/video data after +preprocessing. + +It uses apache arrow in-memory format that can be consumed by modern data +frameworks that can handle parquet or lance file. +""" + +import pyarrow as pa + +pyarrow_schema_i2v = pa.schema([ + pa.field("id", pa.string()), + # --- Image/Video VAE latents --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("vae_latent_bytes", pa.binary()), + # e.g., [C, T, H, W] or [C, H, W] + pa.field("vae_latent_shape", pa.list_(pa.int64())), + # e.g., 'float32' + pa.field("vae_latent_dtype", pa.string()), + # --- Text encoder output tensor --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("text_embedding_bytes", pa.binary()), + # e.g., [SeqLen, Dim] + pa.field("text_embedding_shape", pa.list_(pa.int64())), + # e.g., 'bfloat16' or 'float32' + pa.field("text_embedding_dtype", pa.string()), + #I2V + pa.field("clip_feature_bytes", pa.binary()), + pa.field("clip_feature_shape", pa.list_(pa.int64())), + pa.field("clip_feature_dtype", pa.string()), + pa.field("first_frame_latent_bytes", pa.binary()), + pa.field("first_frame_latent_shape", pa.list_(pa.int64())), + pa.field("first_frame_latent_dtype", pa.string()), + # I2V Validation + pa.field("pil_image_bytes", pa.binary()), + pa.field("pil_image_shape", pa.list_(pa.int64())), + pa.field("pil_image_dtype", pa.string()), + # --- Metadata --- + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), # 'image' or 'video' + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + # -- Video-specific (can be null/default for images) --- + # Number of frames processed (e.g., 1 for image, N for video) + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), +]) + +pyarrow_schema_t2v = pa.schema([ + pa.field("id", pa.string()), + # --- Image/Video VAE latents --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("vae_latent_bytes", pa.binary()), + # e.g., [C, T, H, W] or [C, H, W] + pa.field("vae_latent_shape", pa.list_(pa.int64())), + # e.g., 'float32' + pa.field("vae_latent_dtype", pa.string()), + # --- Text encoder output tensor --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("text_embedding_bytes", pa.binary()), + # e.g., [SeqLen, Dim] + pa.field("text_embedding_shape", pa.list_(pa.int64())), + # e.g., 'bfloat16' or 'float32' + pa.field("text_embedding_dtype", pa.string()), + # --- Metadata --- + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), # 'image' or 'video' + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + # -- Video-specific (can be null/default for images) --- + # Number of frames processed (e.g., 1 for image, N for video) + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), +]) diff --git a/trainer/dataset/parquet_dataset_iterable_style.py b/trainer/dataset/parquet_dataset_iterable_style.py new file mode 100644 index 0000000000000000000000000000000000000000..c85eac824fbb813f8a44d2fd888cf796061cabd2 --- /dev/null +++ b/trainer/dataset/parquet_dataset_iterable_style.py @@ -0,0 +1,305 @@ +import os +import pickle +import random + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import tqdm +from torch.utils.data import IterableDataset, get_worker_info +from torchdata.stateful_dataloader import StatefulDataLoader + +from trainer.dataset.utils import collate_latents_embs_masks +from trainer.distributed import (get_sp_world_size, get_world_group, + get_world_rank, get_world_size) +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class BatchIterator: + # TODO: Implement state_dict and load_state_dict to support resume. + def __init__(self, files, batch_size, text_padding_length, keys, + worker_num_samples, read_batch_size): + self.files = files + self.batch_size = batch_size + self.text_padding_length = text_padding_length + self.keys = keys + self.worker_num_samples = worker_num_samples + self.processed_samples = 0 + self.buffer = [] + self.read_batch_size = read_batch_size + + def __iter__(self): + for file in self.files: + if self.processed_samples >= self.worker_num_samples: + return + + reader = pq.ParquetFile(file) + for batch in reader.iter_batches(batch_size=self.read_batch_size): + if self.processed_samples >= self.worker_num_samples: + return + + self.buffer.extend(batch.to_pylist()) + + while len(self.buffer) >= self.batch_size: + if self.processed_samples >= self.worker_num_samples: + return + + batch_to_process = self.buffer[:self.batch_size] + self.buffer = self.buffer[self.batch_size:] + + all_latents, all_embs, all_masks, caption_text = collate_latents_embs_masks( + batch_to_process, self.text_padding_length, self.keys) + self.processed_samples += self.batch_size + yield all_latents, all_embs, all_masks, caption_text + + +class LatentsParquetIterStyleDataset(IterableDataset): + """Efficient loader for video-text data from a directory of Parquet files.""" + + # Modify this in the future if we want to add more keys, for example, in image to video. + keys = [("vae_latent", "latent"), ("text_embedding")] + + def __init__(self, + path: str, + batch_size: int = 1024, + cfg_rate: float = 0.1, + num_workers: int = 1, + drop_last: bool = True, + text_padding_length: int = 512, + seed: int = 42, + read_batch_size: int = 32, + parquet_schema: pa.Schema = None): + super().__init__() + self.path = str(path) + self.batch_size = batch_size + self.parquet_schema = parquet_schema + self.cfg_rate = cfg_rate + self.text_padding_length = text_padding_length + self.seed = seed + self.read_batch_size = read_batch_size + # Get distributed training info + self.global_rank = get_world_rank() + self.world_size = get_world_size() + self.sp_world_size = get_sp_world_size() + self.num_sp_groups = self.world_size // self.sp_world_size + num_workers = 1 if num_workers == 0 else num_workers + # Get sharding info + shard_parquet_files, shard_total_samples, shard_parquet_lengths = shard_parquet_files_across_sp_groups_and_workers( + self.path, self.num_sp_groups, num_workers, seed) + + if drop_last: + self.worker_num_samples = min( + shard_total_samples) // batch_size * batch_size + # Assign files to current rank's SP group + ith_sp_group = self.global_rank // self.sp_world_size + self.sp_group_parquet_files = shard_parquet_files[ith_sp_group::self + .num_sp_groups] + self.sp_group_parquet_lengths = shard_parquet_lengths[ + ith_sp_group::self.num_sp_groups] + self.sp_group_num_samples = shard_total_samples[ith_sp_group::self. + num_sp_groups] + logger.info( + "In total %d parquet files, %d samples, after sharding we retain %d samples due to drop_last", + sum([len(shard) for shard in shard_parquet_files]), + sum(shard_total_samples), + self.worker_num_samples * self.num_sp_groups * num_workers) + else: + raise ValueError("drop_last must be True") + logger.info("Each dataloader worker will load %d samples", + self.worker_num_samples) + + def __iter__(self): + worker_info = get_worker_info() + worker_id = worker_info.id if worker_info is not None else 1 + + worker_files = self.sp_group_parquet_files[worker_id] + + batch_iterator = BatchIterator( + files=worker_files, + batch_size=self.batch_size, + text_padding_length=self.text_padding_length, + keys=self.keys, + worker_num_samples=self.worker_num_samples, + read_batch_size=self.read_batch_size) # type: ignore + + yield from batch_iterator + + if batch_iterator.processed_samples != self.worker_num_samples: + raise ValueError( + "Rank %d, Worker %d: Not enough samples to process, this should not happen", + self.global_rank, worker_id) + + +def shard_parquet_files_across_sp_groups_and_workers( + path: str, + num_sp_groups: int, + num_workers: int, + seed: int = 42, +) -> tuple[list[list[str]], list[int], list[dict[str, int]]]: + """ + Shard parquet files across SP groups and workers in a balanced way. + + Args: + path: Directory containing parquet files + num_sp_groups: Number of SP groups to shard across + num_workers: Number of workers per SP group + seed: Random seed for shuffling + + Returns: + Tuple containing: + - List of lists of parquet files for each shard + - List of total samples per shard + - List of dictionaries mapping file paths to their lengths + """ + # Check if sharding plan already exists + sharding_info_dir = os.path.join( + path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers") + + # Only rank 0 handles cache checking and file scanning + if get_world_rank() == 0: + cache_loaded = False + shard_parquet_files = None + shard_total_samples = None + shard_parquet_lengths = None + + # First try to load existing sharding plan + if os.path.exists(sharding_info_dir): + logger.info("Loading sharding plan from %s", sharding_info_dir) + try: + with open( + os.path.join(sharding_info_dir, + "shard_parquet_files.pkl"), "rb") as f: + shard_parquet_files = pickle.load(f) + with open( + os.path.join(sharding_info_dir, + "shard_total_samples.pkl"), "rb") as f: + shard_total_samples = pickle.load(f) + with open( + os.path.join(sharding_info_dir, + "shard_parquet_lengths.pkl"), "rb") as f: + shard_parquet_lengths = pickle.load(f) + cache_loaded = True + logger.info("Successfully loaded sharding plan") + except Exception as e: + logger.error("Error loading sharding plan: %s", str(e)) + logger.info("Falling back to creating new sharding plan") + cache_loaded = False + + # If cache not loaded (either doesn't exist or failed to load), create sharding plan + if not cache_loaded: + logger.info("Creating new sharding plan") + logger.info("Scanning for parquet files in %s", path) + + # Find all parquet files + parquet_files = [] + + for root, _, files in os.walk(path): + for file in files: + if file.endswith('.parquet'): + parquet_files.append(os.path.join(root, file)) + + if not parquet_files: + raise ValueError("No parquet files found in %s", path) + + # Calculate file lengths efficiently using a single pass + logger.info("Calculating file lengths...") + lengths = [] + for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"): + lengths.append(pq.ParquetFile(file).metadata.num_rows) + + total_samples = sum(lengths) + logger.info("Found %d files with %d total samples", + len(parquet_files), total_samples) + + # Sort files by length for better balancing + sorted_indices = np.argsort(lengths) + sorted_files = [parquet_files[i] for i in sorted_indices] + sorted_lengths = [lengths[i] for i in sorted_indices] + + # Create shards + num_shards = num_sp_groups * num_workers + shard_parquet_files = [[] for _ in range(num_shards)] + shard_total_samples = [0] * num_shards + shard_parquet_lengths = [{} for _ in range(num_shards)] + + # Distribute files to shards using a greedy approach + logger.info("Distributing files to shards...") + for file, length in zip(reversed(sorted_files), + reversed(sorted_lengths), + strict=True): + # Find shard with minimum current length + target_shard = np.argmin(shard_total_samples) + shard_parquet_files[target_shard].append(file) + shard_total_samples[target_shard] += length + shard_parquet_lengths[target_shard][file] = length + #randomize each shard + for shard in shard_parquet_files: + rng = random.Random(seed) + rng.shuffle(shard) + + # Save the sharding plan + os.makedirs(sharding_info_dir, exist_ok=True) + with open( + os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), + "wb") as f: + pickle.dump(shard_parquet_files, f) + with open( + os.path.join(sharding_info_dir, "shard_total_samples.pkl"), + "wb") as f: + pickle.dump(shard_total_samples, f) + with open( + os.path.join(sharding_info_dir, + "shard_parquet_lengths.pkl"), "wb") as f: + pickle.dump(shard_parquet_lengths, f) + logger.info("Saved sharding info to %s", sharding_info_dir) + + # Wait for rank 0 to finish creating/loading sharding plan + world_group = get_world_group() + world_group.barrier() + + # Now all ranks load the sharding plan (it should exist and be valid now) + logger.info("Loading sharding plan from %s after barrier", + sharding_info_dir) + with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), + "rb") as f: + shard_parquet_files = pickle.load(f) + with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"), + "rb") as f: + shard_total_samples = pickle.load(f) + with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), + "rb") as f: + shard_parquet_lengths = pickle.load(f) + + return shard_parquet_files, shard_total_samples, shard_parquet_lengths + + +def build_parquet_iterable_style_dataloader( + path: str, + batch_size: int, + num_data_workers: int, + cfg_rate: float = 0.0, + drop_last: bool = True, + text_padding_length: int = 512, + seed: int = 42, + read_batch_size: int = 32 +) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]: + """Build a dataloader for the LatentsParquetIterStyleDataset.""" + dataset = LatentsParquetIterStyleDataset( + path=path, + batch_size=batch_size, + cfg_rate=cfg_rate, + num_workers=num_data_workers, + drop_last=drop_last, + text_padding_length=text_padding_length, + seed=seed, + read_batch_size=read_batch_size) + + loader = StatefulDataLoader( + dataset, + batch_size=1, + num_workers=num_data_workers, + pin_memory=True, + ) + return dataset, loader diff --git a/trainer/dataset/preprocessing_datasets.py b/trainer/dataset/preprocessing_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..82e7b0ccfbaed19e942b89d50a448681cc5344b1 --- /dev/null +++ b/trainer/dataset/preprocessing_datasets.py @@ -0,0 +1,630 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import math +import os +import random +from abc import ABC, abstractmethod +from collections import Counter +from dataclasses import dataclass +from os.path import join as opj +from typing import Any + +import numpy as np +import torch +import torchvision +from einops import rearrange +from PIL import Image +from transformers import AutoTokenizer + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class PreprocessBatch: + """ + Batch information for dataset processing stages. + + This class holds all the information about a video-caption or image-caption pair + as it moves through the processing pipeline. Fields are populated by different stages. + """ + # Raw metadata + path: str + cap: str | list[str] + resolution: dict | None = None + fps: float | None = None + duration: float | None = None + + # Processed metadata + num_frames: int | None = None + sample_frame_index: list[int] | None = None + sample_num_frames: int | None = None + + # Processed data + pixel_values: torch.Tensor | None = None + text: str | None = None + input_ids: torch.Tensor | None = None + cond_mask: torch.Tensor | None = None + + @property + def is_video(self) -> bool: + """Check if this is a video item.""" + return self.path.endswith(".mp4") + + @property + def is_image(self) -> bool: + """Check if this is an image item.""" + return self.path.endswith(".jpg") + + +class DatasetStage(ABC): + """ + Abstract base class for dataset processing stages. + + Similar to PipelineStage but designed for dataset preprocessing operations. + """ + + @abstractmethod + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """ + Process the dataset batch. + + Args: + batch: Dataset batch to process + **kwargs: Additional processing parameters + + Returns: + Processed batch + """ + raise NotImplementedError + + +class DatasetFilterStage(ABC): + """ + Abstract base class for dataset filtering stages. + + These stages can filter out items during metadata processing. + """ + + @abstractmethod + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """ + Check if batch should be kept. + + Args: + batch: Dataset batch to check + **kwargs: Additional parameters + + Returns: + True if batch should be kept, False otherwise + """ + raise NotImplementedError + + @abstractmethod + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """ + Process the dataset batch (for non-filtering operations). + + Args: + batch: Dataset batch to process + **kwargs: Additional processing parameters + + Returns: + Processed batch + """ + raise NotImplementedError + + +class DataValidationStage(DatasetFilterStage): + """Stage for validating data items.""" + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """ + Validate data item. + + Args: + batch: Dataset batch to validate + + Returns: + True if valid, False if invalid + """ + # Check for caption + if batch.cap is None: + return False + + if batch.is_video: + # Validate video-specific fields + if batch.duration is None or batch.fps is None: + return False + elif not batch.is_image: + return False + + return True + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process does nothing for validation - filtering is handled by should_keep.""" + return batch + + +class ResolutionFilterStage(DatasetFilterStage): + """Stage for filtering data items based on resolution constraints.""" + + def __init__(self, + max_h_div_w_ratio: float = 17 / 16, + min_h_div_w_ratio: float = 8 / 16, + max_height: int = 1024, + max_width: int = 1024): + self.max_h_div_w_ratio = max_h_div_w_ratio + self.min_h_div_w_ratio = min_h_div_w_ratio + self.max_height = max_height + self.max_width = max_width + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """ + Check if data item passes resolution filtering. + + Args: + batch: Dataset batch with resolution information + + Returns: + True if passes filter, False otherwise + """ + # Only apply to videos + if not batch.is_video: + return True + + if batch.resolution is None: + return False + + height = batch.resolution.get("height", None) + width = batch.resolution.get("width", None) + if height is None or width is None: + return False + + # Check aspect ratio + aspect = self.max_height / self.max_width + hw_aspect_thr = 1.5 + + return self.filter_resolution( + height, + width, + max_h_div_w_ratio=hw_aspect_thr * aspect, + min_h_div_w_ratio=1 / hw_aspect_thr * aspect, + ) + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process does nothing for resolution filtering - filtering is handled by should_keep.""" + return batch + + def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float, + min_h_div_w_ratio: float) -> bool: + """Filter based on height/width ratio.""" + return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio + + +class FrameSamplingStage(DatasetFilterStage): + """Stage for temporal frame sampling and indexing.""" + + def __init__(self, + num_frames: int, + train_fps: int, + speed_factor: int = 1, + video_length_tolerance_range: float = 5.0, + drop_short_ratio: float = 0.0, + seed: int = 42): + self.num_frames = num_frames + self.train_fps = train_fps + self.speed_factor = speed_factor + self.video_length_tolerance_range = video_length_tolerance_range + self.drop_short_ratio = drop_short_ratio + # Create a seeded random generator for deterministic sampling + self.rng = random.Random(seed) + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """ + Check if video should be kept based on length constraints. + + Args: + batch: Dataset batch + + Returns: + True if should be kept, False otherwise + """ + if batch.is_image: + return True + + if batch.duration is None or batch.fps is None: + return False + + num_frames = math.ceil(batch.fps * batch.duration) + + # Check if video is too long + if (num_frames / batch.fps > self.video_length_tolerance_range * + (self.num_frames / self.train_fps * self.speed_factor)): + return False + + # Resample frame indices to check length + frame_interval = batch.fps / self.train_fps + start_frame_idx = 0 + frame_indices = np.arange(start_frame_idx, num_frames, + frame_interval).astype(int) + + # Filter short videos + return not (len(frame_indices) < self.num_frames + and self.rng.random() < self.drop_short_ratio) + + def process(self, + batch: PreprocessBatch, + temporal_sample_fn=None, + **kwargs) -> PreprocessBatch: + """ + Process frame sampling for video data items. + + Args: + batch: Dataset batch + temporal_sample_fn: Function for temporal sampling + + Returns: + Updated batch with frame sampling info + """ + if batch.is_image: + # For images, just add sample info + batch.sample_frame_index = [0] + batch.sample_num_frames = 1 + return batch + + assert batch.duration is not None and batch.fps is not None + batch.num_frames = math.ceil(batch.fps * batch.duration) + + # Resample frame indices + frame_interval = batch.fps / self.train_fps + start_frame_idx = 0 + frame_indices = np.arange(start_frame_idx, batch.num_frames, + frame_interval).astype(int) + + # Temporal crop if too long + if len(frame_indices) > self.num_frames: + if temporal_sample_fn is not None: + begin_index, end_index = temporal_sample_fn(len(frame_indices)) + frame_indices = frame_indices[begin_index:end_index] + else: + frame_indices = frame_indices[:self.num_frames] + + batch.sample_frame_index = frame_indices.tolist() + batch.sample_num_frames = len(frame_indices) + + return batch + + +class VideoTransformStage(DatasetStage): + """Stage for video data transformation.""" + + def __init__(self, transform) -> None: + self.transform = transform + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """ + Transform video data. + + Args: + batch: Dataset batch with video information + + Returns: + Batch with transformed video tensor + """ + if not batch.is_video: + return batch + + assert os.path.exists(batch.path), f"file {batch.path} do not exist!" + assert batch.sample_frame_index is not None, "Frame indices must be set before transformation" + + torchvision_video, _, metadata = torchvision.io.read_video( + batch.path, output_format="TCHW") + video = torchvision_video[batch.sample_frame_index] + if self.transform is not None: + video = self.transform(video) + video = rearrange(video, "t c h w -> c t h w") + video = video.to(torch.uint8) + + h, w = video.shape[-2:] + assert ( + h / w <= 17 / 16 and h / w >= 8 / 16 + ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}" + + video = video.float() / 127.5 - 1.0 + batch.pixel_values = video + return batch + + +class ImageTransformStage(DatasetStage): + """Stage for image data transformation.""" + + def __init__(self, transform, transform_topcrop) -> None: + self.transform = transform + self.transform_topcrop = transform_topcrop + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """ + Transform image data. + + Args: + batch: Dataset batch with image information + + Returns: + Batch with transformed image tensor + """ + if not batch.is_image: + return batch + + image = Image.open(batch.path).convert("RGB") + image = torch.from_numpy(np.array(image)) + image = rearrange(image, "h w c -> c h w").unsqueeze(0) + + if self.transform_topcrop is not None: + image = self.transform_topcrop(image) + elif self.transform is not None: + image = self.transform(image) + + image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W] + image = image.float() / 127.5 - 1.0 + batch.pixel_values = image + return batch + + +class TextEncodingStage(DatasetStage): + """Stage for text tokenization and encoding.""" + + def __init__(self, + tokenizer, + text_max_length: int, + cfg_rate: float = 0.0, + seed: int = 42): + self.tokenizer = tokenizer + self.text_max_length = text_max_length + self.cfg_rate = cfg_rate + # Create a seeded random generator for deterministic CFG + self.rng = random.Random(seed) + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """ + Process text data. + + Args: + batch: Dataset batch with caption information + + Returns: + Batch with encoded text information + """ + text = batch.cap + if not isinstance(text, list): + text = [text] + text = [self.rng.choice(text)] + + text = text[0] if self.rng.random() > self.cfg_rate else "" + text_tokens_and_mask = self.tokenizer( + text, + max_length=self.text_max_length, + padding="max_length", + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors="pt", + ) + + batch.text = text + batch.input_ids = text_tokens_and_mask["input_ids"] + batch.cond_mask = text_tokens_and_mask["attention_mask"] + return batch + + +class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, + torch.distributed.checkpoint.stateful.Stateful): + """ + Merged dataset for video and caption data with stage-based processing. + Assumes that data_merge_path is a txt file with the following format: + , + + The folder should contain videos. + + The json file should be a list of dictionaries with the following format: + [ + { + "path": "1gGQy4nxyUo-Scene-016.mp4", + "resolution": { + "width": 1920, + "height": 1080 + }, + "size": 2439112, + "fps": 25.0, + "duration": 6.88, + "num_frames": 172, + "cap": [ + "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open." + ] + }, + ... + ] + + This dataset processes video and image data through a series of stages: + - Data validation + - Resolution filtering + - Frame sampling + - Transformation + - Text encoding + """ + + def __init__(self, + data_merge_path: str, + args, + transform, + temporal_sample, + transform_topcrop, + start_idx: int = 0, + seed: int = 42): + self.data_merge_path = data_merge_path + self.start_idx = start_idx + self.args = args + self.temporal_sample = temporal_sample + self.seed = seed + + # Initialize tokenizer + tokenizer_path = os.path.join(args.model_path, "tokenizer") + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, + cache_dir=args.cache_dir) + + # Initialize processing stages + self._init_stages(args, transform, transform_topcrop, tokenizer) + + # Process metadata + self.processed_batches = self._process_metadata() + + def _init_stages(self, args, transform, transform_topcrop, + tokenizer) -> None: + """Initialize all processing stages.""" + self.validation_stage = DataValidationStage() + self.resolution_filter_stage = ResolutionFilterStage( + max_height=args.max_height, max_width=args.max_width) + self.frame_sampling_stage = FrameSamplingStage( + num_frames=args.num_frames, + train_fps=args.train_fps, + speed_factor=args.speed_factor, + video_length_tolerance_range=args.video_length_tolerance_range, + drop_short_ratio=args.drop_short_ratio, + seed=self.seed) + self.video_transform_stage = VideoTransformStage(transform) + self.image_transform_stage = ImageTransformStage( + transform, transform_topcrop) + self.text_encoding_stage = TextEncodingStage( + tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=args.training_cfg_rate, + seed=self.seed) + + def _load_raw_data(self) -> list[dict]: + """Load raw data from JSON files.""" + # Read folder-annotation pairs + with open(self.data_merge_path) as f: + folder_anno_pairs = [ + line.strip().split(",") for line in f if line.strip() + ] + assert len( + folder_anno_pairs) == 1, "Only support one folder-annotation pair" + assert len(folder_anno_pairs[0] + ) == 2, "Folder-annotation pair should have two elements" + folder, annotation_file = folder_anno_pairs[0] + + data_items: list[dict] = [] + with open(annotation_file) as f: + data_items = json.load(f) + + # Update paths with folder prefix + for item in data_items: + item["path"] = opj(folder, item["path"]) + + return data_items + + def _process_metadata(self) -> list[PreprocessBatch]: + """Process the raw metadata through all filtering stages.""" + raw_data = self._load_raw_data() + processed_batches = [] + + # Initialize counters + filter_counts = { + "validation_failed": 0, + "resolution_failed": 0, + "frame_sampling_failed": 0 + } + sample_num_frames: list[int] = [] + + for item in raw_data: + batch = PreprocessBatch(path=item["path"], + cap=item["cap"], + resolution=item.get("resolution"), + fps=item.get("fps"), + duration=item.get("duration")) + + # Apply filtering stages + if not self._apply_filter_stages(batch, filter_counts): + continue + + # Apply frame sampling processing + batch = self.frame_sampling_stage.process( + batch, temporal_sample_fn=self.temporal_sample) + + processed_batches.append(batch) + assert batch.sample_num_frames is not None + sample_num_frames.append(batch.sample_num_frames) + + self._log_filtering_stats(filter_counts, sample_num_frames, + len(raw_data), len(processed_batches)) + return processed_batches + + def _apply_filter_stages(self, batch: PreprocessBatch, + filter_counts: dict[str, int]) -> bool: + """Apply all filter stages and update counters. Returns True if batch should be kept.""" + if not self.validation_stage.should_keep(batch): + filter_counts["validation_failed"] += 1 + return False + + if not self.resolution_filter_stage.should_keep(batch): + filter_counts["resolution_failed"] += 1 + return False + + if not self.frame_sampling_stage.should_keep(batch): + filter_counts["frame_sampling_failed"] += 1 + return False + + return True + + def _log_filtering_stats(self, filter_counts: dict[str, int], + sample_num_frames: list[int], before_count: int, + after_count: int): + """Log filtering statistics.""" + logger.info( + "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, " + "Counter(sample_num_frames): %s, before filter: %d, after filter: %d", + filter_counts['validation_failed'], + filter_counts['resolution_failed'], + filter_counts['frame_sampling_failed'], Counter(sample_num_frames), + before_count, after_count) + + def __iter__(self): + """Iterate through processed data items.""" + for idx in range(len(self.processed_batches)): + yield self._get_item(idx) + + def __len__(self): + return len(self.processed_batches) + + def _get_item(self, idx: int) -> dict: + """Get a single processed data item.""" + batch = self.processed_batches[idx] + + # Apply transformation stages + batch = self.video_transform_stage.process(batch) + batch = self.image_transform_stage.process(batch) + batch = self.text_encoding_stage.process(batch) + + # Build result dictionary + result = { + "pixel_values": batch.pixel_values, + "text": batch.text, + "input_ids": batch.input_ids, + "cond_mask": batch.cond_mask, + "path": batch.path, + } + + # Add video-specific fields + if batch.is_video: + result.update({"fps": batch.fps, "duration": batch.duration}) + + return result + + def state_dict(self) -> dict[str, Any]: + """Return state dict for checkpointing.""" + return {"processed_batches": self.processed_batches} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load state dict from checkpoint.""" + self.processed_batches = state_dict["processed_batches"] diff --git a/trainer/dataset/transform.py b/trainer/dataset/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..53c9c49f3e04ccf4707a2e6617c9e0d1678b90e2 --- /dev/null +++ b/trainer/dataset/transform.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +import random + +import torch + + +def _is_tensor_video_clip(clip) -> bool: + if not torch.is_tensor(clip): + raise TypeError(f"clip should be Tensor. Got {type(clip)}") + + if not clip.ndimension() == 4: + raise ValueError(f"clip should be 4D. Got {clip.dim()}D") + + return True + + +def crop(clip, i, j, h, w) -> torch.Tensor: + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) + """ + if len(clip.size()) != 4: + raise ValueError("clip should be a 4D tensor") + return clip[..., i:i + h, j:j + w] + + +def resize(clip, target_size, interpolation_mode) -> torch.Tensor: + if len(target_size) != 2: + raise ValueError( + f"target size should be tuple (height, width), instead got {target_size}" + ) + return torch.nn.functional.interpolate( + clip, + size=target_size, + mode=interpolation_mode, + align_corners=True, + antialias=True, + ) + + +def center_crop_th_tw(clip, th, tw, top_crop) -> torch.Tensor: + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + + # import ipdb;ipdb.set_trace() + h, w = clip.size(-2), clip.size(-1) + tr = th / tw + if h / w > tr: + new_h = int(w * tr) + new_w = w + else: + new_h = h + new_w = int(h / tr) + + i = 0 if top_crop else int(round((h - new_h) / 2.0)) + j = int(round((w - new_w) / 2.0)) + return crop(clip, i, j, new_h, new_w) + + +def normalize_video(clip) -> torch.Tensor: + """ + Convert tensor data type from uint8 to float, divide value by 255.0 and + permute the dimensions of clip tensor + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) + Return: + clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) + """ + _is_tensor_video_clip(clip) + if not clip.dtype == torch.uint8: + raise TypeError( + f"clip tensor should have data type uint8. Got {clip.dtype}") + # return clip.float().permute(3, 0, 1, 2) / 255.0 + return clip.float() / 255.0 + + +class CenterCropResizeVideo: + """ + First use the short side for cropping length, + center crop video, then resize to the specified size + """ + + def __init__( + self, + size, + top_crop=False, + interpolation_mode="bilinear", + ) -> None: + if len(size) != 2: + raise ValueError( + f"size should be tuple (height, width), instead got {size}") + self.size = size + self.top_crop = top_crop + self.interpolation_mode = interpolation_mode + + def __call__(self, clip) -> torch.Tensor: + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) + Returns: + torch.tensor: scale resized / center cropped video clip. + size is (T, C, crop_size, crop_size) + """ + clip_center_crop = center_crop_th_tw(clip, + self.size[0], + self.size[1], + top_crop=self.top_crop) + clip_center_crop_resize = resize( + clip_center_crop, + target_size=self.size, + interpolation_mode=self.interpolation_mode, + ) + return clip_center_crop_resize + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" + + +class Normalize255: + """ + Convert tensor data type from uint8 to float, divide value by 255.0 and + """ + + def __init__(self) -> None: + pass + + def __call__(self, clip) -> torch.Tensor: + """ + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) + Return: + clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) + """ + return normalize_video(clip) + + def __repr__(self) -> str: + return self.__class__.__name__ + + +class TemporalRandomCrop: + """Temporally crop the given frame indices at a random location. + + Args: + size (int): Desired length of frames will be seen in the model. + """ + + def __init__(self, size) -> None: + self.size = size + + def __call__(self, total_frames) -> tuple[int, int]: + rand_end = max(0, total_frames - self.size - 1) + begin_index = random.randint(0, rand_end) + end_index = min(begin_index + self.size, total_frames) + return begin_index, end_index diff --git a/trainer/dataset/upload_to_hf.py b/trainer/dataset/upload_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..8e95e83ab50f1bc8a3763bf05d5799b9cb49d6be --- /dev/null +++ b/trainer/dataset/upload_to_hf.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +from huggingface_hub import HfApi, upload_folder + +api = HfApi() +repo_id = "weizhou03/HD-Mixkit-Finetune-Wan" # customize this +api.create_repo(repo_id=repo_id, repo_type="dataset") + +upload_folder(repo_id=repo_id, + folder_path="/workspace/data/HD-Mixkit-Finetune-Wan", + repo_type="dataset", + path_in_repo="") diff --git a/trainer/dataset/utils.py b/trainer/dataset/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ab1bf128a677bbaa898a558893abcdd7c47f752f --- /dev/null +++ b/trainer/dataset/utils.py @@ -0,0 +1,225 @@ +import random +from typing import Any, cast + +import numpy as np +import torch + + +def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor: + """ + Pad or crop an embedding [L, D] to exactly padding_length tokens. + Return: + - [L, D] tensor in pinned CPU memory + - [L] attention mask in pinned CPU memory + """ + L, D = t.shape + if padding_length > L: # pad + pad = torch.zeros(padding_length - L, D, dtype=t.dtype, device=t.device) + return torch.cat([t, pad], 0), torch.cat( + [torch.ones(L), torch.zeros(padding_length - L)], 0) + else: # crop + return t[:padding_length], torch.ones(padding_length) + + +def get_torch_tensors_from_row_dict(row_dict, + keys, + cfg_rate, + rng=None) -> dict[str, Any]: + """ + Get the latents and prompts from a row dictionary. + """ + return_dict = {} + for key in keys: + shape, bytes = None, None + if isinstance(key, tuple): + for k in key: + try: + shape = row_dict[f"{k}_shape"] + bytes = row_dict[f"{k}_bytes"] + except KeyError: + continue + key = key[0] + if shape is None or bytes is None: + raise ValueError(f"Key {key} not found in row_dict") + else: + shape = row_dict[f"{key}_shape"] + bytes = row_dict[f"{key}_bytes"] + + # TODO (peiyuan): read precision + if key == 'text_embedding' and (rng.random() + if rng else random.random()) < cfg_rate: + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() + data = torch.from_numpy(data) + if len(data.shape) == 3: + B, L, D = data.shape + assert B == 1, "Batch size must be 1" + data = data.squeeze(0) + return_dict[key] = data + return return_dict + + +def collate_latents_embs_masks( + batch_to_process, + text_padding_length, + keys, + cfg_rate=0.0, + rng=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]: + # Initialize tensors to hold padded embeddings and masks + all_latents = [] + all_embs = [] + all_masks = [] + caption_text = [] + # Process each row individually + for i, row in enumerate(batch_to_process): + # Get tensors from row + data = get_torch_tensors_from_row_dict(row, keys, cfg_rate, rng) + latents, emb = data["vae_latent"], data["text_embedding"] + + padded_emb, mask = pad(emb, text_padding_length) + # Store in batch tensors + all_latents.append(latents) + all_embs.append(padded_emb) + all_masks.append(mask) + # TODO(py): remove this once we fix preprocess + try: + caption_text.append(row["prompt"]) + except KeyError: + caption_text.append(row["caption"]) + + # Pin memory for faster transfer to GPU + all_latents = torch.stack(all_latents) + all_embs = torch.stack(all_embs) + all_masks = torch.stack(all_masks) + + return all_latents, all_embs, all_masks, caption_text + + +def collate_rows_from_parquet_schema(rows, + parquet_schema, + text_padding_length, + cfg_rate=0.0, + rng=None) -> dict[str, Any]: + """ + Collate rows from parquet files based on the provided schema. + Dynamically processes tensor fields based on schema and returns batched data. + + Args: + rows: List of row dictionaries from parquet files + parquet_schema: PyArrow schema defining the structure of the data + + Returns: + Dict containing batched tensors and metadata + """ + if not rows: + return cast(dict[str, Any], {}) + + # Initialize containers for different data types + batch_data: dict[str, Any] = {} + + # Get tensor and metadata field names from schema (fields ending with '_bytes') + tensor_fields = [] + metadata_fields = [] + for field in parquet_schema.names: + if field.endswith('_bytes'): + shape_field = field.replace('_bytes', '_shape') + dtype_field = field.replace('_bytes', '_dtype') + tensor_name = field.replace('_bytes', '') + tensor_fields.append(tensor_name) + assert shape_field in parquet_schema.names, f"Shape field {shape_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors." + assert dtype_field in parquet_schema.names, f"Dtype field {dtype_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors." + elif not field.endswith('_shape') and not field.endswith('_dtype'): + # Only add actual metadata fields, not the shape/dtype helper fields + metadata_fields.append(field) + + # Process each tensor field + for tensor_name in tensor_fields: + tensor_list = [] + + for row in rows: + # Get tensor data from row using the existing helper function pattern + shape_key = f"{tensor_name}_shape" + bytes_key = f"{tensor_name}_bytes" + + if shape_key in row and bytes_key in row: + shape = row[shape_key] + bytes_data = row[bytes_key] + + if len(bytes_data) == 0: + tensor = torch.zeros(0, dtype=torch.bfloat16) + else: + # Convert bytes to tensor using float32 as default + if tensor_name == 'text_embedding' and (rng.random( + ) if rng else random.random()) < cfg_rate: + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = np.frombuffer( + bytes_data, dtype=np.float32).reshape(shape).copy() + tensor = torch.from_numpy(data) + # if len(data.shape) == 3: + # B, L, D = tensor.shape + # assert B == 1, "Batch size must be 1" + # tensor = tensor.squeeze(0) + + tensor_list.append(tensor) + else: + # Handle missing tensor data + tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) + + # Stack tensors with special handling for text embeddings + if tensor_name == 'text_embedding': + # Handle text embeddings with padding + padded_tensors = [] + attention_masks = [] + + for tensor in tensor_list: + if tensor.numel() > 0: + padded_tensor, mask = pad(tensor, text_padding_length) + padded_tensors.append(padded_tensor) + attention_masks.append(mask) + else: + # Handle empty embeddings - assume default embedding dimension + padded_tensors.append( + torch.zeros(text_padding_length, + 768, + dtype=torch.bfloat16)) + attention_masks.append(torch.zeros(text_padding_length)) + + batch_data[tensor_name] = torch.stack(padded_tensors) + batch_data['text_attention_mask'] = torch.stack(attention_masks) + else: + # Stack all tensors to preserve batch consistency + # Don't filter out None or empty tensors as this breaks batch sizing + try: + batch_data[tensor_name] = torch.stack(tensor_list) + except ValueError as e: + shapes = [ + t.shape + if t is not None and hasattr(t, 'shape') else 'None/Invalid' + for t in tensor_list + ] + raise ValueError( + f"Failed to stack tensors for field '{tensor_name}'. " + f"Tensor shapes: {shapes}. " + f"All tensors in a batch must have compatible shapes. " + f"Original error: {e}") from e + + # Process metadata fields into info_list + info_list = [] + for row in rows: + info = {} + for field in metadata_fields: + info[field] = row.get(field, "") + + # Add prompt field for backward compatibility + info["prompt"] = info.get("caption", "") + info_list.append(info) + + batch_data['info_list'] = info_list + + # Add caption_text for backward compatibility + if info_list and 'caption' in info_list[0]: + batch_data['caption_text'] = [info['caption'] for info in info_list] + + return batch_data diff --git a/trainer/dataset/validation_dataset.py b/trainer/dataset/validation_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..a562e6ee273b2eede7ea83cd89d504cefcbaebf9 --- /dev/null +++ b/trainer/dataset/validation_dataset.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py +import os +import pathlib + +import datasets +from torch.utils.data import IterableDataset + +from trainer.distributed import (get_sp_world_size, get_world_rank, + get_world_size) +from trainer.logger import init_logger +from trainer.models.vision_utils import load_image, load_video + +logger = init_logger(__name__) + + +class ValidationDataset(IterableDataset): + + def __init__(self, filename: str): + super().__init__() + + self.filename = pathlib.Path(filename) + # get directory of filename + self.dir = os.path.abspath(self.filename.parent) + + if not self.filename.exists(): + raise FileNotFoundError( + f"File {self.filename.as_posix()} does not exist") + + if self.filename.suffix == ".csv": + data = datasets.load_dataset("csv", + data_files=self.filename.as_posix(), + split="train") + elif self.filename.suffix == ".json": + data = datasets.load_dataset("json", + data_files=self.filename.as_posix(), + split="train", + field="data") + elif self.filename.suffix == ".parquet": + data = datasets.load_dataset("parquet", + data_files=self.filename.as_posix(), + split="train") + elif self.filename.suffix == ".arrow": + data = datasets.load_dataset("arrow", + data_files=self.filename.as_posix(), + split="train") + else: + _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"] + raise ValueError( + f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}" + ) + + # Get distributed training info + self.global_rank = get_world_rank() + self.world_size = get_world_size() + self.sp_world_size = get_sp_world_size() + self.num_sp_groups = self.world_size // self.sp_world_size + + # Convert to list to get total samples + self.all_samples = list(data) + self.original_total_samples = len(self.all_samples) + + # Extend samples to be a multiple of DP degree (num_sp_groups) + remainder = self.original_total_samples % self.num_sp_groups + if remainder != 0: + samples_to_add = self.num_sp_groups - remainder + + # Duplicate samples cyclically to reach the target + additional_samples = [] + for i in range(samples_to_add): + additional_samples.append( + self.all_samples[i % self.original_total_samples]) + + self.all_samples.extend(additional_samples) + + self.total_samples = len(self.all_samples) + + # Calculate which SP group this rank belongs to + self.sp_group_id = self.global_rank // self.sp_world_size + + # Now all SP groups will have equal number of samples + self.samples_per_sp_group = self.total_samples // self.num_sp_groups + + # Calculate start and end indices for this SP group + self.start_idx = self.sp_group_id * self.samples_per_sp_group + self.end_idx = self.start_idx + self.samples_per_sp_group + + # Get samples for this SP group + self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx] + + logger.info( + "Rank %s (SP group %s): " + "Original samples: %s, " + "Extended samples: %s, " + "SP group samples: %s, " + "Range: [%s:%s]", + self.global_rank, + self.sp_group_id, + self.original_total_samples, + self.total_samples, + len(self.sp_group_samples), + self.start_idx, + self.end_idx, + local_main_process_only=False) + + def __len__(self): + """Return the number of samples for this SP group.""" + return len(self.sp_group_samples) + + def __iter__(self): + for sample in self.sp_group_samples: + # For consistency reasons, we mandate that "caption" is always present in the validation dataset. + # However, since the model specifications use "prompt", we create an alias here. + sample["prompt"] = sample["caption"] + + # Load image or video if the path is provided + # TODO(aryan): need to handle custom columns here for control conditions + sample["image"] = None + sample["video"] = None + + if sample.get("image_path", None) is not None: + image_path = sample["image_path"] + image_path = os.path.join(self.dir, image_path) + sample["image_path"] = image_path + if not pathlib.Path(image_path).is_file( + ) and not image_path.startswith("http"): + logger.warning("Image file %s does not exist.", image_path) + else: + sample["image"] = load_image(image_path) + + if sample.get("video_path", None) is not None: + video_path = sample["video_path"] + video_path = os.path.join(self.dir, video_path) + sample["video_path"] = video_path + if not pathlib.Path(video_path).is_file( + ) and not video_path.startswith("http"): + logger.warning("Video file %s does not exist.", video_path) + else: + sample["video"] = load_video(video_path) + + if sample.get("control_image_path", None) is not None: + control_image_path = sample["control_image_path"] + control_image_path = os.path.join(self.dir, control_image_path) + sample["control_image_path"] = control_image_path + if not pathlib.Path(control_image_path).is_file( + ) and not control_image_path.startswith("http"): + logger.warning("Control Image file %s does not exist.", + control_image_path) + else: + sample["control_image"] = load_image(control_image_path) + + if sample.get("control_video_path", None) is not None: + control_video_path = sample["control_video_path"] + control_video_path = os.path.join(self.dir, control_video_path) + sample["control_video_path"] = control_video_path + if not pathlib.Path(control_video_path).is_file( + ) and not control_video_path.startswith("http"): + logger.warning("Control Video file %s does not exist.", + control_video_path) + else: + sample["control_video"] = load_video(control_video_path) + + sample = {k: v for k, v in sample.items() if v is not None} + yield sample diff --git a/trainer/distributed/__init__.py b/trainer/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..970f3afed7a1e03fd2974327a40c1f4b98a4b52f --- /dev/null +++ b/trainer/distributed/__init__.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 + +from trainer.distributed.communication_op import * +from trainer.distributed.parallel_state import ( + cleanup_dist_env_and_memory, get_dp_group, get_dp_rank, get_dp_world_size, + get_local_torch_device, get_sp_group, get_sp_parallel_rank, + get_sp_world_size, get_tp_group, get_tp_rank, get_tp_world_size, + get_world_group, get_world_rank, get_world_size, + init_distributed_environment, initialize_model_parallel, + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized) +from trainer.distributed.utils import * + +__all__ = [ + # Initialization + "init_distributed_environment", + "initialize_model_parallel", + "cleanup_dist_env_and_memory", + "model_parallel_is_initialized", + "maybe_init_distributed_environment_and_model_parallel", + + # World group + "get_world_group", + "get_world_rank", + "get_world_size", + + # Data parallel group + "get_dp_group", + "get_dp_rank", + "get_dp_world_size", + + # Sequence parallel group + "get_sp_group", + "get_sp_parallel_rank", + "get_sp_world_size", + + # Tensor parallel group + "get_tp_group", + "get_tp_rank", + "get_tp_world_size", + + # Get torch device + "get_local_torch_device", +] diff --git a/trainer/distributed/communication_op.py b/trainer/distributed/communication_op.py new file mode 100644 index 0000000000000000000000000000000000000000..27951fa94562d421f8dfb3d14482b0c8970baea5 --- /dev/null +++ b/trainer/distributed/communication_op.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py + +import torch +import torch.distributed + +from trainer.distributed.parallel_state import get_sp_group, get_tp_group + + +def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + return get_tp_group().all_reduce(input_) + + +def tensor_model_parallel_all_gather(input_: torch.Tensor, + dim: int = -1) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_tp_group().all_gather(input_, dim) + + +# TODO: remove model, make it sequence_parallel +def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1) -> torch.Tensor: + """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group.""" + return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim) + + +def sequence_model_parallel_all_gather(input_: torch.Tensor, + dim: int = -1) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_sp_group().all_gather(input_, dim) diff --git a/trainer/distributed/device_communicators/__init__.py b/trainer/distributed/device_communicators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/distributed/device_communicators/base_device_communicator.py b/trainer/distributed/device_communicators/base_device_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..858e78e18794491a66bfc7343c73eba95ad88505 --- /dev/null +++ b/trainer/distributed/device_communicators/base_device_communicator.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py + +from typing import Any + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.distributed import ProcessGroup, ReduceOp + + +class DistributedAutograd: + """Collection of autograd functions for distributed operations. + + This class provides custom autograd functions for distributed operations like all_reduce, + all_gather, and all_to_all. Each operation is implemented as a static inner class with + proper forward and backward implementations. + """ + + class AllReduce(torch.autograd.Function): + """Differentiable all_reduce operation. + + The gradient of all_reduce is another all_reduce operation since the operation + combines values from all ranks equally. + """ + + @staticmethod + def forward(ctx: Any, + group: ProcessGroup, + input_: Tensor, + op: dist.ReduceOp | None = None) -> Tensor: + ctx.group = group + ctx.op = op + output = input_.clone() + dist.all_reduce(output, group=group, op=op) + return output + + @staticmethod + def backward(ctx: Any, + grad_output: Tensor) -> tuple[None, Tensor, None]: + grad_output = grad_output.clone() + dist.all_reduce(grad_output, group=ctx.group, op=ctx.op) + return None, grad_output, None + + class AllGather(torch.autograd.Function): + """Differentiable all_gather operation. + + The operation gathers tensors from all ranks and concatenates them along a specified dimension. + The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks. + """ + + @staticmethod + def forward(ctx: Any, group: ProcessGroup, input_: Tensor, + world_size: int, dim: int) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.dim = dim + ctx.input_shape = input_.shape + + input_size = input_.size() + output_size = (input_size[0] * world_size, ) + input_size[1:] + output_tensor = torch.empty(output_size, + dtype=input_.dtype, + device=input_.device) + + dist.all_gather_into_tensor(output_tensor, input_, group=group) + + output_tensor = output_tensor.reshape((world_size, ) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape(input_size[:dim] + + (world_size * + input_size[dim], ) + + input_size[dim + 1:]) + return output_tensor + + @staticmethod + def backward(ctx: Any, + grad_output: Tensor) -> tuple[None, Tensor, None, None]: + # Split the gradient tensor along the gathered dimension + dim_size = grad_output.size(ctx.dim) // ctx.world_size + grad_chunks = grad_output.reshape(grad_output.shape[:ctx.dim] + + (ctx.world_size, dim_size) + + grad_output.shape[ctx.dim + 1:]) + grad_chunks = grad_chunks.movedim(ctx.dim, 0) + + # Each rank only needs its corresponding gradient + grad_input = torch.empty(ctx.input_shape, + dtype=grad_output.dtype, + device=grad_output.device) + dist.reduce_scatter_tensor(grad_input, + grad_chunks.contiguous(), + group=ctx.group) + + return None, grad_input, None, None + + class AllToAll4D(torch.autograd.Function): + """Differentiable all_to_all operation specialized for 4D tensors. + + This operation is particularly useful for attention operations where we need to + redistribute data across ranks for efficient parallel processing. + + The operation supports two modes: + 1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads + 2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions + """ + + @staticmethod + def forward(ctx: Any, group: ProcessGroup, input_: Tensor, + world_size: int, scatter_dim: int, + gather_dim: int) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + + if world_size == 1: + return input_ + + assert input_.dim( + ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + + if scatter_dim == 2 and gather_dim == 1: + bs, shard_seqlen, hn, hd = input_.shape + seqlen = shard_seqlen * world_size + shard_hn = hn // world_size + + input_ = input_.transpose( + 0, 2).contiguous() # hn, shard_seqlen, bs, hd + output = torch.empty_like(input_) + + dist.all_to_all_single(output, input_, + group=group) # hn, shard_seqlen, bs, hd + + output = torch.cat(output.split(shard_hn), + dim=1) # sharded hn, seqlen, bs, hd + + output = output.transpose( + 0, 2).contiguous() # bs, seqlen, sharded_hn, hd + + return output + elif scatter_dim == 1 and gather_dim == 2: + bs, seqlen, shard_hn, hd = input_.shape + hn = shard_hn * world_size + shard_seqlen = seqlen // world_size + + input_ = input_.transpose( + 0, 2).contiguous() # shard_hn, seqlen, bs, hd + + input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs, + hd).transpose(0, 1).reshape( + shard_hn * world_size, shard_seqlen, + bs, hd).contiguous() + + output = torch.empty_like(input_) + + dist.all_to_all_single(output, input_, group=group) + + output = output.transpose( + 0, 2).contiguous() # bs, seqlen, sharded_hn, hd + + return output + else: + raise RuntimeError( + f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " + f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported." + ) + + @staticmethod + def backward( + ctx: Any, + grad_output: Tensor) -> tuple[None, Tensor, None, None, None]: + if ctx.world_size == 1: + return None, grad_output, None, None, None + + # For backward pass, we swap scatter_dim and gather_dim + output = DistributedAutograd.AllToAll4D.apply( + ctx.group, grad_output, ctx.world_size, ctx.gather_dim, + ctx.scatter_dim) + return None, output, None, None, None + + +class DeviceCommunicatorBase: + """ + Base class for device-specific communicator with autograd support. + It can use the `cpu_group` to initialize the communicator. + If the device has PyTorch integration (PyTorch can recognize its + communication backend), the `device_group` will also be given. + """ + + def __init__(self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = ""): + self.device = device or torch.device("cpu") + self.cpu_group = cpu_group + self.device_group = device_group + self.unique_name = unique_name + self.rank = dist.get_rank(cpu_group) + self.world_size = dist.get_world_size(cpu_group) + self.ranks = dist.get_process_group_ranks(cpu_group) + self.global_rank = dist.get_rank() + self.global_world_size = dist.get_world_size() + self.rank_in_group = dist.get_group_rank(self.cpu_group, + self.global_rank) + + def all_reduce(self, + input_: torch.Tensor, + op: dist.ReduceOp | None = ReduceOp.SUM) -> torch.Tensor: + """Performs an all_reduce operation with gradient support.""" + return DistributedAutograd.AllReduce.apply(self.device_group, input_, + op) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Performs an all_gather operation with gradient support.""" + if dim < 0: + dim += input_.dim() + return DistributedAutograd.AllGather.apply(self.device_group, input_, + self.world_size, dim) + + def all_to_all_4D(self, + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1) -> torch.Tensor: + """Performs a 4D all-to-all operation with gradient support.""" + return DistributedAutograd.AllToAll4D.apply(self.device_group, input_, + self.world_size, + scatter_dim, gather_dim) + + def gather(self, + input_: torch.Tensor, + dst: int = 0, + dim: int = -1) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}") + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + # Gather. + torch.distributed.gather(input_, + gather_list, + dst=self.ranks[dst], + group=self.device_group) + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv(self, + size: torch.Size, + dtype: torch.dtype, + src: int | None = None) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + pass diff --git a/trainer/distributed/device_communicators/cpu_communicator.py b/trainer/distributed/device_communicators/cpu_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..4759650f5cdf8c13d0339158388eb809b6b2ba7c --- /dev/null +++ b/trainer/distributed/device_communicators/cpu_communicator.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py + +import os + +import torch +from torch.distributed import ProcessGroup + +from trainer.platforms import current_platform +from trainer.platforms.interface import CpuArchEnum + +from .base_device_communicator import DeviceCommunicatorBase + + +class CpuCommunicator(DeviceCommunicatorBase): + + def __init__(self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = ""): + super().__init__(cpu_group, device, device_group, unique_name) + self.dist_module = torch.distributed + + if (current_platform.get_cpu_architecture() + == CpuArchEnum.X86) and hasattr( + torch.ops._C, + "init_shm_manager") and unique_name.startswith("tp"): + self.dist_module = _CPUSHMDistributed(self) + + def all_reduce( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM + ) -> torch.Tensor: + self.dist_module.all_reduce(input_, group=self.device_group, op=op) + return input_ + + def gather(self, + input_: torch.Tensor, + dst: int = 0, + dim: int = -1) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}") + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + + # Gather. + self.dist_module.gather(input_, + gather_list, + dst=self.ranks[dst], + group=self.device_group) + + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + input_size = input_.size() + # NOTE: we have to use concat-style all-gather here, + # stack-style all-gather has compatibility issues with + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 + output_size = (input_size[0] * self.world_size, ) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty(output_size, + dtype=input_.dtype, + device=input_.device) + # All-gather. + self.dist_module.all_gather_into_tensor(output_tensor, + input_, + group=self.device_group) + + # Reshape + output_tensor = output_tensor.reshape((self.world_size, ) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape(input_size[:dim] + + (self.world_size * + input_size[dim], ) + + input_size[dim + 1:]) + return output_tensor + + +class _CPUSHMDistributed: + + def __init__(self, communicator: CpuCommunicator): + instance_identifier = os.environ["VLLM_DIST_IDENT"] + unique_name = communicator.unique_name + instance_identifier = f"{instance_identifier}-{unique_name}" + self.communicator = communicator + + group_ranks = [str(rank) for rank in self.communicator.ranks] + shm_group_identifier = f"[{'-'.join(group_ranks)}]" + self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm" + + self.handle = self._init_cpu_shm() + + def _init_cpu_shm(self) -> int: + handle = torch.ops._C.init_shm_manager( + self.group_name, + self.communicator.world_size, + self.communicator.rank, + ) + torch.distributed.barrier(self.communicator.device_group) + torch.ops._C.join_shm_manager( + handle, + self.group_name, + ) + torch.distributed.barrier(self.communicator.device_group) + + return int(handle) + + def all_reduce(self, + input: torch.Tensor, + group: ProcessGroup | None = None) -> None: + torch.ops._C.shm_allreduce(self.handle, input) + + def gather(self, + input: torch.Tensor, + gather_list: list[torch.Tensor] | None, + dst: int = -1, + group: ProcessGroup | None = None) -> None: + # Note: different from the torch gather, here we use local dst rank. + torch.ops._C.shm_gather(self.handle, input, gather_list, + torch.distributed.get_group_rank(group, dst)) + + def all_gather_into_tensor(self, + output: torch.Tensor, + input: torch.Tensor, + group: ProcessGroup | None = None) -> None: + torch.ops._C.shm_all_gather(self.handle, input, output) diff --git a/trainer/distributed/device_communicators/cuda_communicator.py b/trainer/distributed/device_communicators/cuda_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..55c2f463b5983c9e62c0bb778c0a2fdf7d1d3015 --- /dev/null +++ b/trainer/distributed/device_communicators/cuda_communicator.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py + +import torch +from torch.distributed import ProcessGroup + +from trainer.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase) + + +class CudaCommunicator(DeviceCommunicatorBase): + + def __init__(self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = ""): + super().__init__(cpu_group, device, device_group, unique_name) + + from trainer.distributed.device_communicators.pynccl import ( + PyNcclCommunicator) + + self.pynccl_comm: PyNcclCommunicator | None = None + if self.world_size > 1: + self.pynccl_comm = PyNcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None): + pynccl_comm = self.pynccl_comm + assert pynccl_comm is not None + out = pynccl_comm.all_reduce(input_, op=op) + if out is None: + # fall back to the default all-reduce using PyTorch. + # this usually happens during testing. + # when we run the model, allreduce only happens for the TP + # group, where we always have either custom allreduce or pynccl. + out = input_.clone() + torch.distributed.all_reduce(out, group=self.device_group, op=op) + return out + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv(self, + size: torch.Size, + dtype: torch.dtype, + src: int | None = None) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + if self.pynccl_comm is not None: + self.pynccl_comm = None diff --git a/trainer/distributed/device_communicators/pynccl.py b/trainer/distributed/device_communicators/pynccl.py new file mode 100644 index 0000000000000000000000000000000000000000..14d0a1f668a1f12a3e823ae2033f6b9e0ccab35d --- /dev/null +++ b/trainer/distributed/device_communicators/pynccl.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from trainer.distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum, + ncclRedOpTypeEnum, ncclUniqueId) +from trainer.distributed.utils import StatelessProcessGroup +from trainer.logger import init_logger +from trainer.utils import current_stream + +logger = init_logger(__name__) + + +class PyNcclCommunicator: + + def __init__( + self, + group: ProcessGroup | StatelessProcessGroup, + device: int | str | torch.device, + library_path: str | None = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyNcclCommunicator to. If None, + it will be bind to f"cuda:{local_rank}". + library_path: the path to the NCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "PyNcclCommunicator should be attached to a non-NCCL group.") + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + return + try: + self.nccl = NCCLLibrary(library_path) + except Exception: + # disable because of missing NCCL library + # e.g. in a non-GPU environment + self.available = False + self.disabled = True + return + + self.available = True + self.disabled = False + + logger.info("Trainer is using nccl==%s", self.nccl.ncclGetVersion()) + + if self.rank == 0: + # get the unique id from NCCL + self.unique_id = self.nccl.ncclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = ncclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + # nccl communicator and stream will use this device + # `torch.cuda.device` is a context manager that changes the + # current cuda device to the specified one + with torch.cuda.device(device): + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank) + + stream = current_stream() + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + if stream is not None: + stream.synchronize() + del data + + def all_reduce(self, + in_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None) -> torch.Tensor: + if self.disabled: + return None + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert in_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {in_tensor.device}") + + out_tensor = torch.empty_like(in_tensor) + + if stream is None: + stream = current_stream() + self.nccl.ncclAllReduce(buffer_type(in_tensor.data_ptr()), + buffer_type(out_tensor.data_ptr()), + in_tensor.numel(), + ncclDataTypeEnum.from_torch(in_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), self.comm, + cudaStream_t(stream.cuda_stream)) + return out_tensor + + def all_gather(self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + stream=None): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}") + if stream is None: + stream = current_stream() + self.nccl.ncclAllGather(buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + input_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + self.comm, cudaStream_t(stream.cuda_stream)) + + def reduce_scatter(self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}") + if stream is None: + stream = current_stream() + self.nccl.ncclReduceScatter( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), output_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), self.comm, + cudaStream_t(stream.cuda_stream)) + + def send(self, tensor: torch.Tensor, dst: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}") + if stream is None: + stream = current_stream() + self.nccl.ncclSend(buffer_type(tensor.data_ptr()), tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), dst, + self.comm, cudaStream_t(stream.cuda_stream)) + + def recv(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}") + if stream is None: + stream = current_stream() + self.nccl.ncclRecv(buffer_type(tensor.data_ptr()), tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), src, + self.comm, cudaStream_t(stream.cuda_stream)) + + def broadcast(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}") + if stream is None: + stream = current_stream() + if src == self.rank: + sendbuff = buffer_type(tensor.data_ptr()) + # NCCL requires the sender also to have a receive buffer + recvbuff = buffer_type(tensor.data_ptr()) + else: + sendbuff = buffer_type() + recvbuff = buffer_type(tensor.data_ptr()) + self.nccl.ncclBroadcast(sendbuff, recvbuff, tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), src, + self.comm, cudaStream_t(stream.cuda_stream)) diff --git a/trainer/distributed/device_communicators/pynccl_wrapper.py b/trainer/distributed/device_communicators/pynccl_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..0e43e190e5277540fea193fbdcbeeaf0e383fbb7 --- /dev/null +++ b/trainer/distributed/device_communicators/pynccl_wrapper.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py + +# This file is a pure Python wrapper for the NCCL library. +# The main purpose is to use NCCL combined with CUDA graph. +# Before writing this script, we tried the following approach: +# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself +# often gets stuck when initializing the NCCL communicator. +# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` +# contains many other potential cuda APIs, that are not allowed during +# capturing the CUDA graph. For further details, please check +# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . +# +# Another rejected idea is to write a C/C++ binding for NCCL. It is usually +# doable, but we often encounter issues related with nccl versions, and need +# to switch between different versions of NCCL. See +# https://github.com/NVIDIA/nccl/issues/1234 for more details. +# A C/C++ binding is not flexible enough to handle this. It requires +# recompilation of the code every time we want to switch between different +# versions. This current implementation, with a **pure** Python wrapper, is +# more flexible. We can easily switch between different versions of NCCL by +# changing the environment variable `TRAINER_NCCL_SO_PATH`, or the `so_file` +# variable in the code. + +#TODO(will): support TRAINER_NCCL_SO_PATH + +import ctypes +import platform +from dataclasses import dataclass +from typing import Any + +import torch +from torch.distributed import ReduceOp + +from trainer.logger import init_logger +from trainer.utils import find_nccl_library + +logger = init_logger(__name__) + +# === export types and functions from nccl to Python === +# for the original nccl definition, please check +# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in + +ncclResult_t = ctypes.c_int +ncclComm_t = ctypes.c_void_p + + +class ncclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +cudaStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +ncclDataType_t = ctypes.c_int + + +class ncclDataTypeEnum: + ncclInt8 = 0 + ncclChar = 0 + ncclUint8 = 1 + ncclInt32 = 2 + ncclInt = 2 + ncclUint32 = 3 + ncclInt64 = 4 + ncclUint64 = 5 + ncclFloat16 = 6 + ncclHalf = 6 + ncclFloat32 = 7 + ncclFloat = 7 + ncclFloat64 = 8 + ncclDouble = 8 + ncclBfloat16 = 9 + ncclNumTypes = 10 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.ncclInt8 + if dtype == torch.uint8: + return cls.ncclUint8 + if dtype == torch.int32: + return cls.ncclInt32 + if dtype == torch.int64: + return cls.ncclInt64 + if dtype == torch.float16: + return cls.ncclFloat16 + if dtype == torch.float32: + return cls.ncclFloat32 + if dtype == torch.float64: + return cls.ncclFloat64 + if dtype == torch.bfloat16: + return cls.ncclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +ncclRedOp_t = ctypes.c_int + + +class ncclRedOpTypeEnum: + ncclSum = 0 + ncclProd = 1 + ncclMax = 2 + ncclMin = 3 + ncclAvg = 4 + ncclNumOps = 5 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.ncclSum + if op == ReduceOp.PRODUCT: + return cls.ncclProd + if op == ReduceOp.MAX: + return cls.ncclMax + if op == ReduceOp.MIN: + return cls.ncclMin + if op == ReduceOp.AVG: + return cls.ncclAvg + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: list[Any] + + +class NCCLLibrary: + exported_functions = [ + # const char* ncclGetErrorString(ncclResult_t result) + Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), + # ncclResult_t ncclGetVersion(int *version); + Function("ncclGetVersion", ncclResult_t, + [ctypes.POINTER(ctypes.c_int)]), + # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); + Function("ncclGetUniqueId", ncclResult_t, + [ctypes.POINTER(ncclUniqueId)]), + # ncclResult_t ncclCommInitRank( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); + # note that ncclComm_t is a pointer type, so the first argument + # is a pointer to a pointer + Function("ncclCommInitRank", ncclResult_t, [ + ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int + ]), + # ncclResult_t ncclAllReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function("ncclAllReduce", ncclResult_t, [ + buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, + ncclRedOp_t, ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclAllGather( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function("ncclAllGather", ncclResult_t, [ + buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, + ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclReduceScatter( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function("ncclReduceScatter", ncclResult_t, [ + buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, + ncclRedOp_t, ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclSend( + # const void* sendbuff, size_t count, ncclDataType_t datatype, + # int dest, ncclComm_t comm, cudaStream_t stream); + Function("ncclSend", ncclResult_t, [ + buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, + ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclRecv( + # void* recvbuff, size_t count, ncclDataType_t datatype, + # int src, ncclComm_t comm, cudaStream_t stream); + Function("ncclRecv", ncclResult_t, [ + buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, + ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclBroadcast( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, int root, ncclComm_t comm, + # cudaStream_t stream); + Function("ncclBroadcast", ncclResult_t, [ + buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, + ctypes.c_int, ncclComm_t, cudaStream_t + ]), + + # be cautious! this is a collective call, it will block until all + # processes in the communicator have called this function. + # because Python object destruction can happen in random order, + # it is better not to call it at all. + # ncclResult_t ncclCommDestroy(ncclComm_t comm); + Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: dict[str, dict[str, Any]] = {} + + def __init__(self, so_file: str | None = None): + + so_file = so_file or find_nccl_library() + + try: + if so_file not in NCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + NCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = NCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + logger.error( + "Failed to load NCCL library from %s ." + "It is expected if you are not running on NVIDIA/AMD GPUs." + "Otherwise, the nccl library might not exist, be corrupted " + "or it does not support the current platform %s." + "If you already have the library, please set the " + "environment variable TRAINER_NCCL_SO_PATH" + " to point to the correct nccl library path.", so_file, + platform.platform()) + raise e + + if so_file not in NCCLLibrary.path_to_dict_mapping: + _funcs: dict[str, Any] = {} + for func in NCCLLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + NCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] + + def ncclGetErrorString(self, result: ncclResult_t) -> str: + return str(self._funcs["ncclGetErrorString"](result).decode("utf-8")) + + def NCCL_CHECK(self, result: ncclResult_t) -> None: + if result != 0: + error_str = self.ncclGetErrorString(result) + raise RuntimeError(f"NCCL error: {error_str}") + + def ncclGetVersion(self) -> str: + version = ctypes.c_int() + self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) + version_str = str(version.value) + # something like 21903 --> "2.19.3" + major = version_str[0].lstrip("0") + minor = version_str[1:3].lstrip("0") + patch = version_str[3:].lstrip("0") + return f"{major}.{minor}.{patch}" + + def ncclGetUniqueId(self) -> ncclUniqueId: + unique_id = ncclUniqueId() + self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) + return unique_id + + def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId, + rank: int) -> ncclComm_t: + comm = ncclComm_t() + self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm), + world_size, unique_id, + rank)) + return comm + + def ncclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, op: int, comm: ncclComm_t, + stream: cudaStream_t) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK(self._funcs["ncclAllReduce"](sendbuff, recvbuff, count, + datatype, op, comm, + stream)) + + def ncclReduceScatter(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, op: int, comm: ncclComm_t, + stream: cudaStream_t) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK(self._funcs["ncclReduceScatter"](sendbuff, recvbuff, + count, datatype, op, + comm, stream)) + + def ncclAllGather(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, comm: ncclComm_t, + stream: cudaStream_t) -> None: + # `datatype` actually should be `ncclDataType_t` + # which is an aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK(self._funcs["ncclAllGather"](sendbuff, recvbuff, count, + datatype, comm, stream)) + + def ncclSend(self, sendbuff: buffer_type, count: int, datatype: int, + dest: int, comm: ncclComm_t, stream: cudaStream_t) -> None: + self.NCCL_CHECK(self._funcs["ncclSend"](sendbuff, count, datatype, dest, + comm, stream)) + + def ncclRecv(self, recvbuff: buffer_type, count: int, datatype: int, + src: int, comm: ncclComm_t, stream: cudaStream_t) -> None: + self.NCCL_CHECK(self._funcs["ncclRecv"](recvbuff, count, datatype, src, + comm, stream)) + + def ncclBroadcast(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, root: int, comm: ncclComm_t, + stream: cudaStream_t) -> None: + self.NCCL_CHECK(self._funcs["ncclBroadcast"](sendbuff, recvbuff, count, + datatype, root, comm, + stream)) + + def ncclCommDestroy(self, comm: ncclComm_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) + + +__all__ = [ + "NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId", + "ncclComm_t", "cudaStream_t", "buffer_type" +] diff --git a/trainer/distributed/parallel_state.py b/trainer/distributed/parallel_state.py new file mode 100644 index 0000000000000000000000000000000000000000..733ecac6e50c8e233eed2d9f35e1ab89835ad763 --- /dev/null +++ b/trainer/distributed/parallel_state.py @@ -0,0 +1,1279 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Adapted from +"""Trainer distributed state. +It takes over the control of the distributed environment from PyTorch. +The typical workflow is: + +- call `init_distributed_environment` to initialize the distributed environment. +- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to + initialize the model parallel groups. + +- any code dealing with the distributed stuff + +- call `destroy_model_parallel` to destroy the model parallel groups. +- call `destroy_distributed_environment` to destroy the distributed environment. + +If you only need to use the distributed environment without model parallelism, + you can skip the model parallel initialization and destruction steps. +""" +import contextlib +import os +import pickle +import weakref +from collections import namedtuple +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from multiprocessing import shared_memory +from typing import Any, Optional +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup, ReduceOp + +import trainer.envs as envs +from trainer.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase) +from trainer.distributed.device_communicators.cpu_communicator import ( + CpuCommunicator) +from trainer.distributed.utils import StatelessProcessGroup +from trainer.logger import init_logger +from trainer.platforms import current_platform + +logger = init_logger(__name__) + + +@dataclass +class GraphCaptureContext: + stream: torch.cuda.Stream | None + + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + + +def _split_tensor_dict( + tensor_dict: dict[str, torch.Tensor | Any] +) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]: + """Split the tensor dictionary into two parts: + 1. A list of (key, value) pairs. If the value is a tensor, it is replaced + by its metadata. + 2. A list of tensors. + """ + metadata_list: list[tuple[str, Any]] = [] + tensor_list: list[torch.Tensor] = [] + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor): + # Note: we cannot use `value.device` here, + # because it contains not only the device type but also the device + # index (e.g. "cuda:0"). We only need the device type. + # receiving side will set the device index. + device = value.device.type + metadata_list.append( + (key, TensorMetadata(device, value.dtype, value.size()))) + tensor_list.append(value) + else: + metadata_list.append((key, value)) + return metadata_list, tensor_list + + +_group_name_counter: dict[str, int] = {} + + +def _get_unique_name(name: str) -> str: + """Get a unique name for the group. + Example: + _get_unique_name("tp") -> "tp:0" + _get_unique_name("tp") -> "tp:1" + """ + if name not in _group_name_counter: + _group_name_counter[name] = 0 + newname = f"{name}:{_group_name_counter[name]}" + _group_name_counter[name] += 1 + return newname + + +_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._all_reduce_out_place(tensor) + + +def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + return torch.empty_like(tensor) + + +class GroupCoordinator: + """ + PyTorch ProcessGroup wrapper for a group of processes. + PyTorch ProcessGroup is bound to one specific communication backend, + e.g. NCCL, Gloo, MPI, etc. + GroupCoordinator takes charge of all the communication operations among + the processes in the group. It manages both CPU and device + communication. + """ + + # available attributes: + rank: int # global rank + ranks: list[int] # global ranks in the group + world_size: int # size of the group + # difference between `local_rank` and `rank_in_group`: + # if we have a group of size 4 across two nodes: + # Process | Node | Rank | Local Rank | Rank in Group + # 0 | 0 | 0 | 0 | 0 + # 1 | 0 | 1 | 1 | 1 + # 2 | 1 | 2 | 0 | 2 + # 3 | 1 | 3 | 1 | 3 + local_rank: int # local rank used to assign devices + rank_in_group: int # rank inside the group + cpu_group: ProcessGroup # group for CPU communication + device_group: ProcessGroup # group for device communication + use_device_communicator: bool # whether to use device communicator + device_communicator: DeviceCommunicatorBase # device communicator + mq_broadcaster: Any | None # shared memory broadcaster + + def __init__( + self, + group_ranks: list[list[int]], + local_rank: int, + torch_distributed_backend: str | Backend, + use_device_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, + ): + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + + for ranks in group_ranks: + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend) + # a group with `gloo` backend, to allow direct coordination between + # processes through the CPU. + cpu_group = torch.distributed.new_group(ranks, backend="gloo") + if self.rank in ranks: + self.ranks = ranks + self.world_size = len(ranks) + self.rank_in_group = ranks.index(self.rank) + self.device_group = device_group + self.cpu_group = cpu_group + try: + assert self.cpu_group is not None + assert self.device_group is not None + except Exception as e: + print(f"rank: {self.rank} group not found") + raise e + + from trainer.platforms import current_platform + + # TODO: fix it for other platforms + self.device = get_local_torch_device() + + self.use_device_communicator = use_device_communicator + + self.device_communicator: DeviceCommunicatorBase = None # type: ignore + if use_device_communicator and self.world_size > 1: + # Platform-aware device communicator selection + if current_platform.is_cuda_alike(): + from trainer.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator) + self.device_communicator = CudaCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + else: + # For MPS and CPU, use the CPU communicator + self.device_communicator = CpuCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + + self.mq_broadcaster = None + + from trainer.platforms import current_platform + + # TODO(will): check if this is needed + # self.use_custom_op_call = current_platform.is_cuda_alike() + self.use_custom_op_call = False + + @property + def first_rank(self): + """Return the global rank of the first process in the group""" + return self.ranks[0] + + @property + def last_rank(self): + """Return the global rank of the last process in the group""" + return self.ranks[-1] + + @property + def is_first_rank(self): + """Return whether the caller is the first process in the group""" + return self.rank == self.first_rank + + @property + def is_last_rank(self): + """Return whether the caller is the last process in the group""" + return self.rank == self.last_rank + + @property + def next_rank(self): + """Return the global rank of the process that follows the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group + 1) % world_size] + + @property + def prev_rank(self): + """Return the global rank of the process that precedes the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group - 1) % world_size] + + @contextmanager + def graph_capture(self, + graph_capture_context: GraphCaptureContext | None = None): + # Platform-aware graph capture + from trainer.platforms import current_platform + + if current_platform.is_cuda_alike(): + if graph_capture_context is None: + stream = torch.cuda.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = torch.cuda.current_stream() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with torch.cuda.stream(stream): + yield graph_capture_context + else: + # For non-CUDA platforms (MPS, CPU), just yield the context without stream management + if graph_capture_context is None: + # Create a dummy context for non-CUDA platforms + graph_capture_context = GraphCaptureContext(None) + yield graph_capture_context + + def all_reduce( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + """ + User-facing all-reduce function before we actually call the + all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we always make the all-reduce operation + out-of-place. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + if self.use_custom_op_call: + return torch.ops.vllm.all_reduce(input_, + group_name=self.unique_name) + else: + return self._all_reduce_out_place(input_, op=op) + + def _all_reduce_out_place( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + return self.device_communicator.all_reduce(input_, op=op) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}") + + return self.device_communicator.all_gather(input_, dim) + + def gather(self, + input_: torch.Tensor, + dst: int = 0, + dim: int = -1) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + return self.device_communicator.gather(input_, dst, dim) + + def all_to_all_4D(self, + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1) -> torch.Tensor: + if self.world_size == 1: + return input_ + return self.device_communicator.all_to_all_4D(input_, scatter_dim, + gather_dim) + + def broadcast(self, input_: torch.Tensor, src: int = 0): + """Broadcast the input tensor. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + # Broadcast. + torch.distributed.broadcast(input_, + src=self.ranks[src], + group=self.device_group) + return input_ + + def broadcast_object(self, obj: Any | None = None, src: int = 0): + """Broadcast the input object. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj + if self.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_broadcaster.broadcast_object(obj) + if self.rank_in_group == src: + torch.distributed.broadcast_object_list([obj], + src=self.ranks[src], + group=self.cpu_group) + return obj + else: + recv = [None] + torch.distributed.broadcast_object_list(recv, + src=self.ranks[src], + group=self.cpu_group) + return recv[0] + + def broadcast_object_list(self, + obj_list: list[Any], + src: int = 0, + group: ProcessGroup | None = None): + """Broadcast the input object list. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj_list + # Broadcast. + torch.distributed.broadcast_object_list(obj_list, + src=self.ranks[src], + group=self.device_group) + return obj_list + + def send_object(self, obj: Any, dst: int) -> None: + """Send the input object list to the destination rank.""" + """NOTE: `dst` is the local rank of the destination rank.""" + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank.") + + # Serialize object to tensor and get the size as well + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) + + size_tensor = torch.tensor([object_tensor.numel()], + dtype=torch.long, + device="cpu") + + # Send object size + + torch.distributed.send(size_tensor, + dst=self.ranks[dst], + group=self.cpu_group) + + # Send object + torch.distributed.send(object_tensor, + dst=self.ranks[dst], + group=self.cpu_group) + + return None + + def recv_object(self, src: int) -> Any: + """Receive the input object list from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + + assert src < self.world_size, f"Invalid src rank ({src})" + + assert src != self.rank_in_group, ( + "Invalid source rank. Source rank is the same as the current rank.") + + size_tensor = torch.empty(1, dtype=torch.long, device="cpu") + + # Receive object size + rank_size = torch.distributed.recv(size_tensor, + src=self.ranks[src], + group=self.cpu_group) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device="cpu") + + rank_object = torch.distributed.recv(object_tensor, + src=self.ranks[src], + group=self.cpu_group) + + assert rank_object == rank_size, ( + "Received object sender rank does not match the size sender rank.") + + obj = pickle.loads(object_tensor.numpy().tobytes()) + + return obj + + def broadcast_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any] | None = None, + src: int = 0, + group: ProcessGroup | None = None, + metadata_group: ProcessGroup | None = None + ) -> dict[str, torch.Tensor | Any] | None: + """Broadcast the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if (not torch.distributed.is_initialized() or self.world_size == 1): + return tensor_dict + + group = self.device_group + metadata_group = self.cpu_group + assert src < self.world_size, f"Invalid src rank ({src})" + + rank_in_group = self.rank_in_group + if rank_in_group == src: + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, + dict), (f"Expecting a dictionary, got {type(tensor_dict)}") + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `broadcast_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.broadcast_object(metadata_list, src=src) + async_handles = [] + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast(tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast(tensor, + src=self.ranks[src], + group=group, + async_op=True) + async_handles.append(handle) + for async_handle in async_handles: + async_handle.wait() + + else: + metadata_list = self.broadcast_object(None, src=src) + tensor_dict = {} + async_handles = [] + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, + dtype=value.dtype, + device=value.device) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=group, + async_op=True) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + for async_handle in async_handles: + async_handle.wait() + return tensor_dict + + def send_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any], + dst: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Send the input tensor dictionary. + NOTE: `dst` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + all_gather_size = (1 if all_gather_group is None else + all_gather_group.world_size) + all_gather_rank = (0 if all_gather_group is None else + all_gather_group.rank_in_group) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, + dict), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `send_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.send_object(metadata_list, dst=dst) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if (all_gather_group is not None + and tensor.numel() % all_gather_size == 0): + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.send(tensor, + dst=self.ranks[dst], + group=metadata_group) + else: + # use group for GPU tensors + torch.distributed.send(tensor, dst=self.ranks[dst], group=group) + return None + + def recv_tensor_dict( + self, + src: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Recv the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return None + + all_gather_size = (1 if all_gather_group is None else + all_gather_group.world_size) + all_gather_rank = (0 if all_gather_group is None else + all_gather_group.rank_in_group) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: dict[str, Any] = {} + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, + dtype=value.dtype, + device=value.device) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = (all_gather_group is not None + and tensor.numel() % all_gather_size == 0) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, + -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.recv(tensor, + src=self.ranks[src], + group=metadata_group) + else: + # use group for GPU tensors + torch.distributed.recv(tensor, + src=self.ranks[src], + group=group) + if use_all_gather: + # do the allgather + tensor = all_gather_group.all_gather( # type: ignore + tensor, dim=0) + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + return tensor_dict + + def barrier(self) -> None: + """Barrier synchronization among the group. + NOTE: don't use `device_group` here! `barrier` in NCCL is + terrible because it is internally a broadcast operation with + secretly created GPU tensors. It is easy to mess up the current + device. Use the CPU group instead. + """ + torch.distributed.barrier(group=self.cpu_group) + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + self.device_communicator.send(tensor, dst) + + def recv(self, + size: torch.Size, + dtype: torch.dtype, + src: int | None = None) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + return self.device_communicator.recv(size, dtype, src) + + def destroy(self) -> None: + if self.device_group is not None: + torch.distributed.destroy_process_group(self.device_group) + self.device_group = None + if self.cpu_group is not None: + torch.distributed.destroy_process_group(self.cpu_group) + self.cpu_group = None + if self.device_communicator is not None: + self.device_communicator.destroy() + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: GroupCoordinator | None = None +_NODE: GroupCoordinator | None = None + + +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, ("world group is not initialized") + return _WORLD + + +def init_world_group(ranks: list[int], local_rank: int, + backend: str) -> GroupCoordinator: + return GroupCoordinator( + group_ranks=[ranks], + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: list[list[int]], + local_rank: int, + backend: str, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, +) -> GroupCoordinator: + + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: GroupCoordinator | None = None + + +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, ("tensor model parallel group is not initialized") + return _TP + + +_ENABLE_CUSTOM_ALL_REDUCE = True + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def init_distributed_environment( + world_size: int = 1, + rank: int = 0, + distributed_init_method: str = "env://", + local_rank: int = 0, + backend: str = "nccl", + device_id: torch.device | None = None, +): + # Determine the appropriate backend based on the platform + from trainer.platforms import current_platform + if backend == "nccl" and not current_platform.is_cuda_alike(): + # Use gloo backend for non-CUDA platforms (MPS, CPU) + backend = "gloo" + logger.info("Using gloo backend for %s platform", + current_platform.device_name) + + logger.debug( + "world_size=%d rank=%d local_rank=%d " + "distributed_init_method=%s backend=%s", world_size, rank, local_rank, + distributed_init_method, backend) + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment") + + # For MPS, don't pass device_id as it doesn't support device indices + if current_platform.is_mps(): + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank) + else: + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + device_id=device_id) + # set the local rank + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = envs.LOCAL_RANK + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert _WORLD.world_size == torch.distributed.get_world_size(), ( + "world group already initialized with a different world size") + + +_SP: GroupCoordinator | None = None + + +def get_sp_group() -> GroupCoordinator: + assert _SP is not None, ("sequence model parallel group is not initialized") + return _SP + + +_DP: GroupCoordinator | None = None + + +def get_dp_group() -> GroupCoordinator: + assert _DP is not None, ("data parallel group is not initialized") + return _DP + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + sequence_model_parallel_size: int = 1, + data_parallel_size: int = 1, + backend: str | None = None, +) -> None: + """ + Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism (used for language encoder). + sequence_model_parallel_size: number of GPUs used for sequence model + parallelism (used for DiT). + """ + # Get world size and rank. Ensure some consistencies. + assert _WORLD is not None, "world group is not initialized, please call init_distributed_environment first" + world_size: int = get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + assert world_size >= tensor_model_parallel_size, f"world_size({world_size}) must be greater than or equal to tensor_model_parallel_size({tensor_model_parallel_size})" + num_tensor_model_parallel_groups: int = (world_size // + tensor_model_parallel_size) + global _TP + assert _TP is None, ("tensor model parallel group is already initialized") + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size)) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group(group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="tp") + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = (world_size // + sequence_model_parallel_size) + global _SP + assert _SP is None, ("sequence model parallel group is already initialized") + group_ranks = [] + + # Since SP is incompatible with TP and PP, we can use a simpler group creation logic + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range(i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size)) + group_ranks.append(ranks) + + _SP = init_model_parallel_group(group_ranks, + get_world_group().local_rank, + backend, + group_name="sp") + + # Build the data parallel groups. + num_data_parallel_groups: int = sequence_model_parallel_size + global _DP + assert _DP is None, ("data parallel group is already initialized") + group_ranks = [] + + for i in range(num_data_parallel_groups): + ranks = list(range(i, world_size, num_data_parallel_groups)) + group_ranks.append(ranks) + + _DP = init_model_parallel_group(group_ranks, + get_world_group().local_rank, + backend, + group_name="dp") + + +def get_sp_world_size() -> int: + """Return world size for the sequence model parallel group.""" + return get_sp_group().world_size + + +def get_sp_parallel_rank() -> int: + """Return my rank for the sequence model parallel group.""" + return get_sp_group().rank_in_group + + +def get_world_size() -> int: + """Return world size for the world group.""" + return get_world_group().world_size + + +def get_world_rank() -> int: + """Return my rank for the world group.""" + return get_world_group().rank + + +def get_dp_world_size() -> int: + """Return world size for the data parallel group.""" + return get_dp_group().world_size + + +def get_dp_rank() -> int: + """Return my rank for the data parallel group.""" + return get_dp_group().rank_in_group + + +def get_local_torch_device() -> torch.device: + """Return the torch device for the current rank.""" + return torch.device(f"cuda:{envs.LOCAL_RANK}" + ) if current_platform.is_cuda_alike() else torch.device( + "mps") + + +def maybe_init_distributed_environment_and_model_parallel( + tp_size: int, sp_size: int, distributed_init_method: str = "env://"): + if _WORLD is not None and model_parallel_is_initialized(): + # make sure the tp and sp sizes are correct + assert get_tp_world_size( + ) == tp_size, f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" + assert get_sp_world_size( + ) == sp_size, f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + return + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + rank = int(os.environ.get("RANK", 0)) + device = get_local_torch_device() + logger.info( + "Initializing distributed environment with world_size=%d, device=%s", + world_size, device) + + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method=distributed_init_method, + device_id=device) + initialize_model_parallel(tensor_model_parallel_size=tp_size, + sequence_model_parallel_size=sp_size) + + # Only set CUDA device if we're on a CUDA platform + if current_platform.is_cuda_alike(): + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + + +def model_parallel_is_initialized() -> bool: + """Check if tensor, sequence parallel groups are initialized.""" + return _TP is not None and _SP is not None and _DP is not None + + +_TP_STATE_PATCHED = False + + +@contextmanager +def patch_tensor_parallel_group(tp_group: GroupCoordinator): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _TP_STATE_PATCHED + assert not _TP_STATE_PATCHED, "Should not call when it's already patched" + + _TP_STATE_PATCHED = True + old_tp_group = get_tp_group() + global _TP + _TP = tp_group + try: + yield + finally: + # restore the original state + _TP_STATE_PATCHED = False + _TP = old_tp_group + + +def get_tp_world_size() -> int: + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tp_rank() -> int: + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def destroy_model_parallel() -> None: + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _SP + if _SP: + _SP.destroy() + _SP = None + + global _DP + if _DP: + _DP.destroy() + _DP = None + + +def destroy_distributed_environment() -> None: + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def cleanup_dist_env_and_memory(shutdown_ray: bool = False): + destroy_model_parallel() + destroy_distributed_environment() + with contextlib.suppress(AssertionError): + torch.distributed.destroy_process_group() + if shutdown_ray: + import ray # Lazy import Ray + ray.shutdown() + + +def is_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup, + source_rank: int = 0) -> list[int]: + """ + This is a collective operation that returns if each rank is in the same node + as the source rank. It tests if processes are attached to the same + memory system (shared access to shared memory). + """ + if isinstance(pg, ProcessGroup): + assert torch.distributed.get_backend( + pg) != torch.distributed.Backend.NCCL, ( + "in_the_same_node_as should be tested with a non-NCCL group.") + # local rank inside the group + rank = torch.distributed.get_rank(group=pg) + world_size = torch.distributed.get_world_size(group=pg) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + else: + rank = pg.rank + world_size = pg.world_size + ranks = list(range(world_size)) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + magic_message = b"magic_message" + shm = None + + try: + with contextlib.suppress(OSError): + if rank == source_rank: + # create a shared memory segment + shm = shared_memory.SharedMemory(create=True, size=128) + shm.buf[:len(magic_message)] = magic_message + if isinstance(pg, ProcessGroup): + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg) + else: + pg.broadcast_obj(shm.name, src=source_rank) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + if isinstance(pg, ProcessGroup): + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg) + name = recv[0] + else: + name = pg.broadcast_obj(None, src=source_rank) + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch("multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[:len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as e: + logger.error("Error ignored in is_in_the_same_node: %s", e) + finally: + if shm: + shm.close() + + if isinstance(pg, ProcessGroup): + torch.distributed.barrier(group=pg) + else: + pg.barrier() + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + + if isinstance(pg, ProcessGroup): + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + aggregated_data = is_in_the_same_node + else: + aggregated_data = torch.zeros_like(is_in_the_same_node) + for i in range(world_size): + rank_data = pg.broadcast_obj(is_in_the_same_node, src=i) + aggregated_data += rank_data + + return [x == 1 for x in aggregated_data.tolist()] + + +def initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "") -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % tensor_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = (world_size // + tensor_model_parallel_size) + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size)) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_model_parallel_group(tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "") -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % sequence_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = (world_size // + sequence_model_parallel_size) + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range(i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size)) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_model_parallel_group(sp_group_ranks, + get_world_group().local_rank, + backend, + group_name=group_name) + + return sp_group diff --git a/trainer/distributed/utils.py b/trainer/distributed/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..51e994fcda01fca2b4c5dd713d04a92501b00d49 --- /dev/null +++ b/trainer/distributed/utils.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +import dataclasses +import pickle +import time +from collections import deque +from collections.abc import Sequence +from typing import Any + +import torch +from torch.distributed import TCPStore + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +def ensure_divisibility(numerator, denominator) -> None: + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, "{} is not divisible by {}".format( + numerator, denominator) + + +def divide(numerator: int, denominator: int) -> int: + """Ensure that numerator is divisible by the denominator and return + the division value.""" + ensure_divisibility(numerator, denominator) + return numerator // denominator + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> Sequence[torch.Tensor]: + """ Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = divide(tensor.size()[last_dim], num_partitions) + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # NOTE: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tuple(tensor_list) + + +@dataclasses.dataclass +class StatelessProcessGroup: + """A dataclass to hold a metadata store, and the rank, world_size of the + group. Only use it to communicate metadata between processes. + For data-plane communication, create NCCL-related objects. + """ + rank: int + world_size: int + store: torch._C._distributed_c10d.Store + data_expiration_seconds: int = 3600 # 1 hour + + # dst rank -> counter + send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict) + # src rank -> counter + recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) + broadcast_send_counter: int = 0 + broadcast_recv_src_counter: dict[int, int] = dataclasses.field( + default_factory=dict) + + # A deque to store the data entries, with key and timestamp. + entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque) + + def __post_init__(self): + assert self.rank < self.world_size + self.send_dst_counter = {i: 0 for i in range(self.world_size)} + self.recv_src_counter = {i: 0 for i in range(self.world_size)} + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def send_obj(self, obj: Any, dst: int): + """Send an object to a destination rank.""" + self.expire_data() + key = f"send_to/{dst}/{self.send_dst_counter[dst]}" + self.store.set(key, pickle.dumps(obj)) + self.send_dst_counter[dst] += 1 + self.entries.append((key, time.perf_counter())) + + def expire_data(self) -> None: + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.perf_counter() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def recv_obj(self, src: int) -> Any: + """Receive an object from a source rank.""" + obj = pickle.loads( + self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}")) + self.recv_src_counter[src] += 1 + return obj + + def broadcast_obj(self, obj: Any | None, src: int) -> Any: + """Broadcast an object from a source rank to all other ranks. + It does not clean up after all ranks have received the object. + Use it for limited times, e.g., for initialization. + """ + if self.rank == src: + self.expire_data() + key = (f"broadcast_from/{src}/" + f"{self.broadcast_send_counter}") + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.perf_counter())) + return obj + else: + key = (f"broadcast_from/{src}/" + f"{self.broadcast_recv_src_counter[src]}") + recv_obj = pickle.loads(self.store.get(key)) + self.broadcast_recv_src_counter[src] += 1 + return recv_obj + + def all_gather_obj(self, obj: Any) -> list[Any]: + """All gather an object from all ranks.""" + gathered_objs = [] + for i in range(self.world_size): + if i == self.rank: + gathered_objs.append(obj) + self.broadcast_obj(obj, src=self.rank) + else: + recv_obj = self.broadcast_obj(None, src=i) + gathered_objs.append(recv_obj) + return gathered_objs + + def barrier(self): + """A barrier to synchronize all ranks.""" + for i in range(self.world_size): + if i == self.rank: + self.broadcast_obj(None, src=self.rank) + else: + self.broadcast_obj(None, src=i) + + @staticmethod + def create( + host: str, + port: int, + rank: int, + world_size: int, + data_expiration_seconds: int = 3600, + ) -> "StatelessProcessGroup": + """A replacement for `torch.distributed.init_process_group` that does not + pollute the global state. + + If we have process A and process B called `torch.distributed.init_process_group` + to form a group, and then we want to form another group with process A, B, C, + D, it is not possible in PyTorch, because process A and process B have already + formed a group, and process C and process D cannot join that group. This + function is a workaround for this issue. + + `torch.distributed.init_process_group` is a global call, while this function + is a stateless call. It will return a `StatelessProcessGroup` object that can be + used for exchanging metadata. With this function, process A and process B + can call `StatelessProcessGroup.create` to form a group, and then process A, B, + C, and D can call `StatelessProcessGroup.create` to form another group. + """ # noqa + store = TCPStore( + host_name=host, + port=port, + world_size=world_size, + is_master=(rank == 0), + ) + + return StatelessProcessGroup( + rank=rank, + world_size=world_size, + store=store, + data_expiration_seconds=data_expiration_seconds) diff --git a/trainer/entrypoints/__init__.py b/trainer/entrypoints/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/entrypoints/cli/__init__.py b/trainer/entrypoints/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/entrypoints/cli/cli_types.py b/trainer/entrypoints/cli/cli_types.py new file mode 100644 index 0000000000000000000000000000000000000000..88ccf282673c22a36e7ed0939e8e9ea16b7008fd --- /dev/null +++ b/trainer/entrypoints/cli/cli_types.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py + +import argparse + +from trainer.utils import FlexibleArgumentParser + + +class CLISubcommand: + """Base class for CLI subcommands""" + + name: str + + def cmd(self, args: argparse.Namespace) -> None: + """Execute the command with the given arguments""" + raise NotImplementedError + + def validate(self, args: argparse.Namespace) -> None: + """Validate the arguments for this command""" + pass + + def subparser_init( + self, + subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + """Initialize the subparser for this command""" + raise NotImplementedError diff --git a/trainer/entrypoints/cli/generate.py b/trainer/entrypoints/cli/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..e8965b4039c88125be729f6429e49d7de2d8ff08 --- /dev/null +++ b/trainer/entrypoints/cli/generate.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py + +import argparse +import dataclasses +import os +from typing import cast + +from trainer import VideoGenerator +from trainer.configs.sample.base import SamplingParam +from trainer.entrypoints.cli.cli_types import CLISubcommand +from trainer.entrypoints.cli.utils import RaiseNotImplementedAction +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.utils import FlexibleArgumentParser + +logger = init_logger(__name__) + + +class GenerateSubcommand(CLISubcommand): + """The `generate` subcommand for the Trainer CLI""" + + def __init__(self) -> None: + self.name = "generate" + super().__init__() + self.init_arg_names = self._get_init_arg_names() + self.generation_arg_names = self._get_generation_arg_names() + + def _get_init_arg_names(self) -> list[str]: + """Get names of arguments for VideoGenerator initialization""" + return ["num_gpus", "tp_size", "sp_size", "model_path"] + + def _get_generation_arg_names(self) -> list[str]: + """Get names of arguments for generate_video method""" + return [field.name for field in dataclasses.fields(SamplingParam)] + + def cmd(self, args: argparse.Namespace) -> None: + excluded_args = ['subparser', 'config', 'dispatch_function'] + + provided_args = {} + for k, v in vars(args).items(): + if (k not in excluded_args and v is not None + and hasattr(args, '_provided') and k in args._provided): + provided_args[k] = v + + if 'model_path' in vars(args) and args.model_path is not None: + provided_args['model_path'] = args.model_path + + if 'prompt' in vars(args) and args.prompt is not None: + provided_args['prompt'] = args.prompt + + merged_args = {**provided_args} + + logger.info('CLI Args: %s', merged_args) + + if 'model_path' not in merged_args or not merged_args['model_path']: + raise ValueError( + "model_path must be provided either in config file or via --model-path" + ) + + # Check if either prompt or prompt_txt is provided + has_prompt = 'prompt' in merged_args and merged_args['prompt'] + has_prompt_txt = 'prompt_txt' in merged_args and merged_args[ + 'prompt_txt'] + + if not (has_prompt or has_prompt_txt): + raise ValueError("Either prompt or prompt_txt must be provided") + + if has_prompt and has_prompt_txt: + raise ValueError( + "Cannot provide both 'prompt' and 'prompt_txt'. Use only one of them." + ) + + init_args = { + k: v + for k, v in merged_args.items() + if k not in self.generation_arg_names + } + generation_args = { + k: v + for k, v in merged_args.items() if k in self.generation_arg_names + } + + model_path = init_args.pop('model_path') + prompt = generation_args.pop('prompt', None) + + generator = VideoGenerator.from_pretrained(model_path=model_path, + **init_args) + + # Call generate_video - it handles both single and batch modes + generator.generate_video(prompt=prompt, **generation_args) + + def validate(self, args: argparse.Namespace) -> None: + """Validate the arguments for this command""" + if args.num_gpus is not None and args.num_gpus <= 0: + raise ValueError("Number of gpus must be positive") + + if args.config and not os.path.exists(args.config): + raise ValueError(f"Config file not found: {args.config}") + + def subparser_init( + self, + subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + generate_parser = subparsers.add_parser( + "generate", + help="Run inference on a model", + usage= + "trainer generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]" + ) + + generate_parser.add_argument( + "--config", + type=str, + default='', + required=False, + help= + "Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional." + ) + + generate_parser = TrainerArgs.add_cli_args(generate_parser) + generate_parser = SamplingParam.add_cli_args(generate_parser) + + generate_parser.add_argument( + "--text-encoder-configs", + action=RaiseNotImplementedAction, + help= + "JSON array of text encoder configurations (NOT YET IMPLEMENTED)", + ) + + return cast(FlexibleArgumentParser, generate_parser) + + +def cmd_init() -> list[CLISubcommand]: + return [GenerateSubcommand()] diff --git a/trainer/entrypoints/cli/main.py b/trainer/entrypoints/cli/main.py new file mode 100644 index 0000000000000000000000000000000000000000..d3a80688da5355e6a1a57963efdac4bd267c662b --- /dev/null +++ b/trainer/entrypoints/cli/main.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py + +from trainer.entrypoints.cli.cli_types import CLISubcommand +from trainer.entrypoints.cli.generate import cmd_init as generate_cmd_init +from trainer.utils import FlexibleArgumentParser + + +def cmd_init() -> list[CLISubcommand]: + """Initialize all commands from separate modules""" + commands = [] + commands.extend(generate_cmd_init()) + return commands + + +def main() -> None: + parser = FlexibleArgumentParser(description="Trainer CLI") + parser.add_argument('-v', '--version', action='version', version='0.1.0') + + subparsers = parser.add_subparsers(required=False, dest="subparser") + + cmds = {} + for cmd in cmd_init(): + cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) + cmds[cmd.name] = cmd + args = parser.parse_args() + if args.subparser in cmds: + cmds[args.subparser].validate(args) + + if hasattr(args, "dispatch_function"): + args.dispatch_function(args) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/trainer/entrypoints/cli/utils.py b/trainer/entrypoints/cli/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d3cc1464d2d3933d1327508448fb791b947eb1ba --- /dev/null +++ b/trainer/entrypoints/cli/utils.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import subprocess +import sys + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class RaiseNotImplementedAction(argparse.Action): + + def __call__(self, parser, namespace, values, option_string=None): + raise NotImplementedError( + f"The {option_string} option is not yet implemented") + + +def launch_distributed(num_gpus: int, + args: list[str], + master_port: int | None = None) -> int: + """ + Launch a distributed job with the given arguments + + Args: + num_gpus: Number of GPUs to use + args: Arguments to pass to v1_trainer_inference.py (defaults to sys.argv[1:]) + master_port: Port for the master process (default: random) + """ + + current_env = os.environ.copy() + python_executable = sys.executable + project_root = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../..")) + main_script = os.path.join(project_root, + "trainer/sample/v1_trainer_inference.py") + + cmd = [ + python_executable, "-m", "torch.distributed.run", + f"--nproc_per_node={num_gpus}" + ] + + if master_port is not None: + cmd.append(f"--master_port={master_port}") + + cmd.append(main_script) + cmd.extend(args) + + logger.info("Running inference with %d GPU(s)", num_gpus) + logger.info("Launching command: %s", " ".join(cmd)) + + current_env["PYTHONIOENCODING"] = "utf-8" + process = subprocess.Popen(cmd, + env=current_env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1, + encoding='utf-8', + errors='replace') + + if process.stdout: + for line in iter(process.stdout.readline, ''): + print(line.strip()) + + return process.wait() diff --git a/trainer/entrypoints/video_generator.py b/trainer/entrypoints/video_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..3b35b368a974f73fcca650b2181233a13784c525 --- /dev/null +++ b/trainer/entrypoints/video_generator.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +VideoGenerator module for Trainer. + +This module provides a consolidated interface for generating videos using +diffusion models. +""" + +import math +import os +import time +from copy import deepcopy +from typing import Any + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +from trainer.configs.sample import SamplingParam +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines import ForwardBatch +from trainer.utils import align_to, shallow_asdict +from trainer.worker.executor import Executor + +logger = init_logger(__name__) + + +class VideoGenerator: + """ + A unified class for generating videos using diffusion models. + + This class provides a simple interface for video generation with rich + customization options, similar to popular frameworks like HF Diffusers. + """ + + def __init__(self, trainer_args: TrainerArgs, + executor_class: type[Executor], log_stats: bool): + """ + Initialize the video generator. + + Args: + trainer_args: The inference arguments + executor_class: The executor class to use for inference + """ + self.trainer_args = trainer_args + self.executor = executor_class(trainer_args) + + @classmethod + def from_pretrained(cls, + model_path: str, + device: str | None = None, + torch_dtype: torch.dtype | None = None, + **kwargs) -> "VideoGenerator": + """ + Create a video generator from a pretrained model. + + Args: + model_path: Path or identifier for the pretrained model + device: Device to load the model on (e.g., "cuda", "cuda:0", "cpu") + torch_dtype: Data type for model weights (e.g., torch.float16) + pipeline_config: Pipeline config to use for inference + **kwargs: Additional arguments to customize model loading, set any TrainerArgs or PipelineConfig attributes here. + + Returns: + The created video generator + + Priority level: Default pipeline config < User's pipeline config < User's kwargs + """ + # If users also provide some kwargs, it will override the TrainerArgs and PipelineConfig. + kwargs['model_path'] = model_path + trainer_args = TrainerArgs.from_kwargs(**kwargs) + + return cls.from_trainer_args(trainer_args) + + @classmethod + def from_trainer_args(cls, + trainer_args: TrainerArgs) -> "VideoGenerator": + """ + Create a video generator with the specified arguments. + + Args: + trainer_args: The inference arguments + + Returns: + The created video generator + """ + # Initialize distributed environment if needed + # initialize_distributed_and_parallelism(trainer_args) + + executor_class = Executor.get_class(trainer_args) + return cls( + trainer_args=trainer_args, + executor_class=executor_class, + log_stats=False, # TODO: implement + ) + + def generate_video( + self, + prompt: str | None = None, + sampling_param: SamplingParam | None = None, + **kwargs, + ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]: + """ + Generate a video based on the given prompt. + + Args: + prompt: The prompt to use for generation (optional if prompt_txt is provided) + negative_prompt: The negative prompt to use (overrides the one in trainer_args) + output_path: Path to save the video (overrides the one in trainer_args) + output_video_name: Name of the video file to save. Default is the first 100 characters of the prompt. + save_video: Whether to save the video to disk + return_frames: Whether to return the raw frames + num_inference_steps: Number of denoising steps (overrides trainer_args) + guidance_scale: Classifier-free guidance scale (overrides trainer_args) + num_frames: Number of frames to generate (overrides trainer_args) + height: Height of generated video (overrides trainer_args) + width: Width of generated video (overrides trainer_args) + fps: Frames per second for saved video (overrides trainer_args) + seed: Random seed for generation (overrides trainer_args) + callback: Callback function called after each step + callback_steps: Number of steps between each callback + + Returns: + Either the output dictionary, list of frames, or list of results for batch processing + """ + # Handle batch processing from text file + if self.trainer_args.prompt_txt is not None: + prompt_txt_path = self.trainer_args.prompt_txt + if not os.path.exists(prompt_txt_path): + raise FileNotFoundError( + f"Prompt text file not found: {prompt_txt_path}") + + # Read prompts from file + with open(prompt_txt_path, encoding='utf-8') as f: + prompts = [line.strip() for line in f if line.strip()] + + if not prompts: + raise ValueError(f"No prompts found in file: {prompt_txt_path}") + + logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path) + + if sampling_param is not None: + original_output_video_name = sampling_param.output_video_name + else: + original_output_video_name = None + + results = [] + for i, batch_prompt in enumerate(prompts): + logger.info("Processing prompt %d/%d: %s...", i + 1, + len(prompts), batch_prompt[:100]) + + try: + # Generate video for this prompt using the same logic below + if sampling_param is not None and original_output_video_name is not None: + sampling_param.output_video_name = original_output_video_name + f"_{i}" + result = self._generate_single_video( + batch_prompt, sampling_param, **kwargs) + + # Add prompt info to result + if isinstance(result, dict): + result["prompt_index"] = i + result["prompt"] = batch_prompt + + results.append(result) + logger.info("Successfully generated video for prompt %d", + i + 1) + + except Exception as e: + logger.error("Failed to generate video for prompt %d: %s", + i + 1, e) + continue + + logger.info( + "Completed batch processing. Generated %d videos successfully.", + len(results)) + return results + + # Single prompt generation (original behavior) + if prompt is None: + raise ValueError("Either prompt or prompt_txt must be provided") + + return self._generate_single_video(prompt, sampling_param, **kwargs) + + def _generate_single_video( + self, + prompt: str, + sampling_param: SamplingParam | None = None, + **kwargs, + ) -> dict[str, Any] | list[np.ndarray]: + """Internal method for single video generation""" + # Create a copy of inference args to avoid modifying the original + trainer_args = self.trainer_args + pipeline_config = trainer_args.pipeline_config + + # Validate inputs + if not isinstance(prompt, str): + raise TypeError( + f"`prompt` must be a string, but got {type(prompt)}") + prompt = prompt.strip() + if sampling_param is None: + sampling_param = SamplingParam.from_pretrained( + trainer_args.model_path) + else: + sampling_param = deepcopy(sampling_param) + + kwargs["prompt"] = prompt + sampling_param.update(kwargs) + + # Process negative prompt + if sampling_param.negative_prompt is not None: + sampling_param.negative_prompt = sampling_param.negative_prompt.strip( + ) + + # Validate dimensions + if (sampling_param.height <= 0 or sampling_param.width <= 0 + or sampling_param.num_frames <= 0): + raise ValueError( + f"Height, width, and num_frames must be positive integers, got " + f"height={sampling_param.height}, width={sampling_param.width}, " + f"num_frames={sampling_param.num_frames}") + + temporal_scale_factor = pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = sampling_param.num_frames + num_gpus = trainer_args.num_gpus + use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames + + # Adjust number of frames based on number of GPUs + if use_temporal_scaling_frames: + orig_latent_num_frames = (num_frames - + 1) // temporal_scale_factor + 1 + else: # stepvideo only + orig_latent_num_frames = sampling_param.num_frames // 17 * 3 + + if orig_latent_num_frames % trainer_args.num_gpus != 0: + # Adjust latent frames to be divisible by number of GPUs + if sampling_param.num_frames_round_down: + # Ensure we have at least 1 batch per GPU + new_latent_num_frames = max( + 1, (orig_latent_num_frames // num_gpus)) * num_gpus + else: + new_latent_num_frames = math.ceil( + orig_latent_num_frames / num_gpus) * num_gpus + + if use_temporal_scaling_frames: + # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor + new_num_frames = (new_latent_num_frames - + 1) * temporal_scale_factor + 1 + else: # stepvideo only + # Find the least common multiple of 3 and num_gpus + divisor = math.lcm(3, num_gpus) + # Round up to the nearest multiple of this LCM + new_latent_num_frames = ( + (new_latent_num_frames + divisor - 1) // divisor) * divisor + # Convert back to actual frames using the StepVideo formula + new_num_frames = new_latent_num_frames // 3 * 17 + + logger.info( + "Adjusting number of frames from %s to %s based on number of GPUs (%s)", + sampling_param.num_frames, new_num_frames, + trainer_args.num_gpus) + sampling_param.num_frames = new_num_frames + + # Calculate sizes + target_height = align_to(sampling_param.height, 16) + target_width = align_to(sampling_param.width, 16) + + # Calculate latent sizes + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + # Log parameters + debug_str = f""" + height: {target_height} + width: {target_width} + video_length: {sampling_param.num_frames} + prompt: {prompt} + image_path: {sampling_param.image_path} + neg_prompt: {sampling_param.negative_prompt} + seed: {sampling_param.seed} + infer_steps: {sampling_param.num_inference_steps} + num_videos_per_prompt: {sampling_param.num_videos_per_prompt} + guidance_scale: {sampling_param.guidance_scale} + n_tokens: {n_tokens} + flow_shift: {trainer_args.pipeline_config.flow_shift} + embedded_guidance_scale: {trainer_args.pipeline_config.embedded_cfg_scale} + save_video: {sampling_param.save_video} + output_path: {sampling_param.output_path} + """ # type: ignore[attr-defined] + logger.info(debug_str) + + # Prepare batch + batch = ForwardBatch( + **shallow_asdict(sampling_param), + eta=0.0, + n_tokens=n_tokens, + VSA_sparsity=trainer_args.VSA_sparsity, + extra={}, + ) + + # Use prompt[:100] for video name + if batch.output_video_name is None: + batch.output_video_name = prompt[:100] + + # Run inference + start_time = time.perf_counter() + output_batch = self.executor.execute_forward(batch, trainer_args) + samples = output_batch.output + logging_info = output_batch.logging_info + + gen_time = time.perf_counter() - start_time + logger.info("Generated successfully in %.2f seconds", gen_time) + + # Process outputs + videos = rearrange(samples, "b c t h w -> t b c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + # Save video if requested + if batch.save_video: + output_path = batch.output_path + if output_path: + os.makedirs(output_path, exist_ok=True) + video_path = os.path.join(output_path, + f"{batch.output_video_name}.mp4") + imageio.mimsave(video_path, frames, fps=batch.fps, format="mp4") + logger.info("Saved video to %s", video_path) + else: + logger.warning("No output path provided, video not saved") + + if batch.return_frames: + return frames + else: + return { + "samples": samples, + "frames": frames, + "prompts": prompt, + "size": (target_height, target_width, batch.num_frames), + "generation_time": gen_time, + "logging_info": logging_info, + } + + def set_lora_adapter(self, + lora_nickname: str, + lora_path: str | None = None) -> None: + self.executor.set_lora_adapter(lora_nickname, lora_path) + + def shutdown(self): + """ + Shutdown the video generator. + """ + self.executor.shutdown() + del self.executor diff --git a/trainer/envs.py b/trainer/envs.py new file mode 100644 index 0000000000000000000000000000000000000000..099e8d7cf6e4bb675d31e853d93df14348a1c7dc --- /dev/null +++ b/trainer/envs.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + TRAINER_RINGBUFFER_WARNING_INTERVAL: int = 60 + TRAINER_NCCL_SO_PATH: str | None = None + LD_LIBRARY_PATH: str | None = None + LOCAL_RANK: int = 0 + CUDA_VISIBLE_DEVICES: str | None = None + TRAINER_CACHE_ROOT: str = os.path.expanduser("~/.cache/trainer") + TRAINER_CONFIG_ROOT: str = os.path.expanduser("~/.config/trainer") + TRAINER_CONFIGURE_LOGGING: int = 1 + TRAINER_LOGGING_LEVEL: str = "INFO" + TRAINER_LOGGING_PREFIX: str = "" + TRAINER_LOGGING_CONFIG_PATH: str | None = None + TRAINER_TRACE_FUNCTION: int = 0 + TRAINER_ATTENTION_BACKEND: str | None = None + TRAINER_ATTENTION_CONFIG: str | None = None + TRAINER_WORKER_MULTIPROC_METHOD: str = "fork" + TRAINER_TARGET_DEVICE: str = "cuda" + MAX_JOBS: str | None = None + NVCC_THREADS: str | None = None + CMAKE_BUILD_TYPE: str | None = None + VERBOSE: bool = False + TRAINER_SERVER_DEV_MODE: bool = False + TRAINER_STAGE_LOGGING: bool = False + + +def get_default_cache_root() -> str: + return os.getenv( + "XDG_CACHE_HOME", + os.path.join(os.path.expanduser("~"), ".cache"), + ) + + +def get_default_config_root() -> str: + return os.getenv( + "XDG_CONFIG_HOME", + os.path.join(os.path.expanduser("~"), ".config"), + ) + + +def maybe_convert_int(value: str | None) -> int | None: + if value is None: + return None + return int(value) + + +# The begin-* and end* here are used by the documentation generator +# to extract the used env vars. + +# begin-env-vars-definition + +environment_variables: dict[str, Callable[[], Any]] = { + + # ================== Installation Time Env Vars ================== + + # Target device of Trainer, supporting [cuda (by default), + # rocm, neuron, cpu, openvino] + "TRAINER_TARGET_DEVICE": + lambda: os.getenv("TRAINER_TARGET_DEVICE", "cuda"), + + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": + lambda: os.getenv("MAX_JOBS", None), + + # Number of threads to use for nvcc + # By default this is 1. + # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. + "NVCC_THREADS": + lambda: os.getenv("NVCC_THREADS", None), + + # If set, trainer will use precompiled binaries (*.so) + "TRAINER_USE_PRECOMPILED": + lambda: bool(os.environ.get("TRAINER_USE_PRECOMPILED")) or bool( + os.environ.get("TRAINER_PRECOMPILED_WHEEL_LOCATION")), + + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": + lambda: os.getenv("CMAKE_BUILD_TYPE"), + + # If set, trainer will print verbose logs during installation + "VERBOSE": + lambda: bool(int(os.getenv('VERBOSE', '0'))), + + # Root directory for TRAINER configuration files + # Defaults to `~/.config/trainer` unless `XDG_CONFIG_HOME` is set + # Note that this not only affects how trainer finds its configuration files + # during runtime, but also affects how trainer installs its configuration + # files during **installation**. + "TRAINER_CONFIG_ROOT": + lambda: os.path.expanduser( + os.getenv( + "TRAINER_CONFIG_ROOT", + os.path.join(get_default_config_root(), "trainer"), + )), + + # ================== Runtime Env Vars ================== + + # Root directory for TRAINER cache files + # Defaults to `~/.cache/trainer` unless `XDG_CACHE_HOME` is set + "TRAINER_CACHE_ROOT": + lambda: os.path.expanduser( + os.getenv( + "TRAINER_CACHE_ROOT", + os.path.join(get_default_cache_root(), "trainer"), + )), + + # Interval in seconds to log a warning message when the ring buffer is full + "TRAINER_RINGBUFFER_WARNING_INTERVAL": + lambda: int(os.environ.get("TRAINER_RINGBUFFER_WARNING_INTERVAL", "60")), + + # Path to the NCCL library file. It is needed because nccl>=2.19 brought + # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 + "TRAINER_NCCL_SO_PATH": + lambda: os.environ.get("TRAINER_NCCL_SO_PATH", None), + + # when `TRAINER_NCCL_SO_PATH` is not set, trainer will try to find the nccl + # library file in the locations specified by `LD_LIBRARY_PATH` + "LD_LIBRARY_PATH": + lambda: os.environ.get("LD_LIBRARY_PATH", None), + + # Internal flag to enable Dynamo fullgraph capture + "TRAINER_TEST_DYNAMO_FULLGRAPH_CAPTURE": + lambda: bool( + os.environ.get("TRAINER_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"), + + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": + lambda: int(os.environ.get("LOCAL_RANK", "0")), + + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": + lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + + # timeout for each iteration in the engine + "TRAINER_ENGINE_ITERATION_TIMEOUT_S": + lambda: int(os.environ.get("TRAINER_ENGINE_ITERATION_TIMEOUT_S", "60")), + + # Logging configuration + # If set to 0, trainer will not configure logging + # If set to 1, trainer will configure logging using the default configuration + # or the configuration file specified by TRAINER_LOGGING_CONFIG_PATH + "TRAINER_CONFIGURE_LOGGING": + lambda: int(os.getenv("TRAINER_CONFIGURE_LOGGING", "1")), + "TRAINER_LOGGING_CONFIG_PATH": + lambda: os.getenv("TRAINER_LOGGING_CONFIG_PATH"), + + # this is used for configuring the default logging level + "TRAINER_LOGGING_LEVEL": + lambda: os.getenv("TRAINER_LOGGING_LEVEL", "INFO"), + + # if set, TRAINER_LOGGING_PREFIX will be prepended to all log messages + "TRAINER_LOGGING_PREFIX": + lambda: os.getenv("TRAINER_LOGGING_PREFIX", ""), + + # Trace function calls + # If set to 1, trainer will trace function calls + # Useful for debugging + "TRAINER_TRACE_FUNCTION": + lambda: int(os.getenv("TRAINER_TRACE_FUNCTION", "0")), + + # Backend for attention computation + # Available options: + # - "TORCH_SDPA": use torch.nn.MultiheadAttention + # - "FLASH_ATTN": use FlashAttention + # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention + # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention + # - "SAGE_ATTN": use Sage Attention + "TRAINER_ATTENTION_BACKEND": + lambda: os.getenv("TRAINER_ATTENTION_BACKEND", None), + + # Path to the attention configuration file. Only used for sliding tile + # attention for now. + "TRAINER_ATTENTION_CONFIG": + lambda: (None if os.getenv("TRAINER_ATTENTION_CONFIG", None) is None else + os.path.expanduser(os.getenv("TRAINER_ATTENTION_CONFIG", "."))), + + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "TRAINER_WORKER_MULTIPROC_METHOD": + lambda: os.getenv("TRAINER_WORKER_MULTIPROC_METHOD", "fork"), + + # Enables torch profiler if set. Path to the directory where torch profiler + # traces are saved. Note that it must be an absolute path. + "TRAINER_TORCH_PROFILER_DIR": + lambda: (None + if os.getenv("TRAINER_TORCH_PROFILER_DIR", None) is None else os. + path.expanduser(os.getenv("TRAINER_TORCH_PROFILER_DIR", "."))), + + # If set, trainer will run in development mode, which will enable + # some additional endpoints for developing and debugging, + # e.g. `/reset_prefix_cache` + "TRAINER_SERVER_DEV_MODE": + lambda: bool(int(os.getenv("TRAINER_SERVER_DEV_MODE", "0"))), + + # If set, trainer will enable stage logging, which will print the time + # taken for each stage + "TRAINER_STAGE_LOGGING": + lambda: bool(int(os.getenv("TRAINER_STAGE_LOGGING", "0"))), +} + +# end-env-vars-definition + + +def __getattr__(name: str): + # lazy evaluation of environment variables + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return list(environment_variables.keys()) diff --git a/trainer/forward_context.py b/trainer/forward_context.py new file mode 100644 index 0000000000000000000000000000000000000000..43f653ccddce1531a54fdff11e77d1048461b48c --- /dev/null +++ b/trainer/forward_context.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/forward_context.py + +import time +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from trainer.logger import init_logger + +if TYPE_CHECKING: + from trainer.attention import AttentionMetadata + from trainer.pipelines import ForwardBatch + +logger = init_logger(__name__) + +# TODO(will): check if this is needed +# track_batchsize: bool = envs.TRAINER_LOG_BATCHSIZE_INTERVAL >= 0 +track_batchsize: bool = False +last_logging_time: float = 0 +forward_start_time: float = 0 +# batchsize_logging_interval: float = envs.TRAINER_LOG_BATCHSIZE_INTERVAL +batchsize_logging_interval: float = 1000 +batchsize_forward_time: defaultdict = defaultdict(list) + + +# +@dataclass +class ForwardContext: + current_timestep: int + # TODO(will): check this arg + # copy from vllm_config.compilation_config.static_forward_context + # attn_layers: Dict[str, Any] + # TODO: extend to support per-layer dynamic forward context + attn_metadata: "AttentionMetadata" # set dynamically for each forward pass + forward_batch: Optional["ForwardBatch"] = None + + +_forward_context: Optional["ForwardContext"] = None + + +def get_forward_context() -> "ForwardContext": + """Get the current forward context.""" + assert _forward_context is not None, ( + "Forward context is not set. " + "Please use `set_forward_context` to set the forward context.") + return _forward_context + + +# TODO(will): finalize the interface +@contextmanager +def set_forward_context(current_timestep, + attn_metadata, + forward_batch: Optional["ForwardBatch"] = None): + """A context manager that stores the current forward context, + can be attention metadata, etc. + Here we can inject common logic for every model forward pass. + """ + global forward_start_time + need_to_track_batchsize = track_batchsize and attn_metadata is not None + if need_to_track_batchsize: + forward_start_time = time.perf_counter() + global _forward_context + prev_context = _forward_context + _forward_context = ForwardContext(current_timestep=current_timestep, + attn_metadata=attn_metadata, + forward_batch=forward_batch) + + try: + yield + finally: + global last_logging_time, batchsize_logging_interval + if need_to_track_batchsize: + if hasattr(attn_metadata, "num_prefill_tokens"): + # for v0 attention backends + batchsize = attn_metadata.num_prefill_tokens + \ + attn_metadata.num_decode_tokens + else: + # for v1 attention backends + batchsize = attn_metadata.num_input_tokens + now = time.perf_counter() + # time measurement is in milliseconds + batchsize_forward_time[batchsize].append( + (now - forward_start_time) * 1000) + if now - last_logging_time > batchsize_logging_interval: + last_logging_time = now + forward_stats = [] + for bs, times in batchsize_forward_time.items(): + if len(times) <= 1: + # can be cudagraph / profiling run + continue + medium = torch.quantile(torch.tensor(times), q=0.5).item() + medium = round(medium, 2) + forward_stats.append((bs, len(times), medium)) + forward_stats.sort(key=lambda x: x[1], reverse=True) + if forward_stats: + logger.info(("Batchsize forward time stats " + "(batchsize, count, median_time(ms)): %s"), + forward_stats) + _forward_context = prev_context diff --git a/trainer/layers/__init__.py b/trainer/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/layers/activation.py b/trainer/layers/activation.py new file mode 100644 index 0000000000000000000000000000000000000000..dd4426d469cac013234709571f646601c938c732 --- /dev/null +++ b/trainer/layers/activation.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py +"""Custom activation functions.""" +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# TODO (will): remove this dependency +from trainer.layers.custom_op import CustomOp + + +@CustomOp.register("silu_and_mul") +class SiluAndMul(CustomOp): + """An activation function for SwiGLU. + + The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + def __init__(self) -> None: + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + +@CustomOp.register("gelu_and_mul") +class GeluAndMul(CustomOp): + """An activation function for GeGLU. + + The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d) + return: (batch_size, seq_len, d) or (num_tokens, d) + """ + + def __init__(self, approximate: str = "none"): + super().__init__() + self.approximate = approximate + if approximate not in ("none", "tanh"): + raise ValueError(f"Unknown approximate mode: {approximate}") + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:] + + def extra_repr(self) -> str: + return f'approximate={repr(self.approximate)}' + + +@CustomOp.register("gelu_new") +class NewGELU(CustomOp): + + def __init__(self): + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + c = math.sqrt(2.0 / math.pi) + return 0.5 * x * (1.0 + torch.tanh(c * + (x + 0.044715 * torch.pow(x, 3.0)))) + + +@CustomOp.register("quick_gelu") +class QuickGELU(CustomOp): + # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90 + def __init__(self): + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + return x * torch.sigmoid(1.702 * x) + + +_ACTIVATION_REGISTRY = { + "gelu": nn.GELU, + "gelu_new": NewGELU, + "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"), + "relu": nn.ReLU, + "silu": nn.SiLU, + "quick_gelu": QuickGELU, +} + + +def get_act_fn(act_fn_name: str) -> nn.Module: + """Get an activation function by name.""" + act_fn_name = act_fn_name.lower() + if act_fn_name not in _ACTIVATION_REGISTRY: + raise ValueError( + f"Activation function {act_fn_name!r} is not supported.") + + return _ACTIVATION_REGISTRY[act_fn_name]() + + +_ACTIVATION_AND_MUL_REGISTRY = { + "gelu": GeluAndMul, + "silu": SiluAndMul, +} + + +def get_act_and_mul_fn(act_fn_name: str) -> nn.Module: + """Get an activation-and-mul (i.e. SiluAndMul) function by name.""" + act_fn_name = act_fn_name.lower() + if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY: + raise ValueError( + f"Activation function {act_fn_name!r} is not supported.") + + return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]() diff --git a/trainer/layers/custom_op.py b/trainer/layers/custom_op.py new file mode 100644 index 0000000000000000000000000000000000000000..bab18f9bd3c62cd4f2022307f219e824796d8969 --- /dev/null +++ b/trainer/layers/custom_op.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py + +from collections.abc import Callable +from typing import Any + +import torch.nn as nn + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class CustomOp(nn.Module): + """ + Base class for custom ops. + Dispatches the forward method to the appropriate backend. + """ + + def __init__(self) -> None: + super().__init__() + self._forward_method = self.dispatch_forward() + + def forward(self, *args, **kwargs) -> Any: + return self._forward_method(*args, **kwargs) + + def forward_native(self, *args, **kwargs) -> Any: + """PyTorch-native implementation of the forward method. + This method is optional. If implemented, it can be used with compilers + such as torch.compile or PyTorch XLA. Also, it can be used for testing + purposes. + """ + raise NotImplementedError + + def forward_cuda(self, *args, **kwargs) -> Any: + raise NotImplementedError + + def forward_cpu(self, *args, **kwargs) -> Any: + # By default, we assume that CPU ops are compatible with CUDA ops. + return self.forward_cuda(*args, **kwargs) + + def forward_tpu(self, *args, **kwargs) -> Any: + # By default, we assume that TPU ops are compatible with the + # PyTorch-native implementation. + # NOTE(woosuk): This is a placeholder for future extensions. + return self.forward_native(*args, **kwargs) + + def forward_oot(self, *args, **kwargs) -> Any: + # By default, we assume that OOT ops are compatible with the + # PyTorch-native implementation. + return self.forward_native(*args, **kwargs) + + def dispatch_forward(self) -> Callable: + # FIXME(will): for now, we always use the native implementation, since + # forward_cuda is using vllm's custom ops and it doesn't support + # backwards. We should add our own custom ops that support backwards. + return self.forward_native + # NOTE(woosuk): Here we assume that vLLM was built for only one + # specific backend. Currently, we do not support dynamic dispatching. + enabled = self.enabled() + + if not enabled: + return self.forward_native + + return self.forward_cuda + + @classmethod + def enabled(cls) -> bool: + # since we are not using Inductor, we always return True + return True + + @staticmethod + def default_on() -> bool: + """ + On by default if level < CompilationLevel.PIECEWISE + Specifying 'all' or 'none' in custom_op takes precedence. + """ + raise NotImplementedError + + # Dictionary of all custom ops (classes, indexed by registered name). + # To check if an op with a name is enabled, call .enabled() on the class. + # Examples: + # - MyOp.enabled() + # - op_registry["my_op"].enabled() + op_registry: dict[str, type['CustomOp']] = {} + + # Decorator to register custom ops. + @classmethod + def register(cls, name: str) -> Callable: + + def decorator(op_cls): + assert name not in cls.op_registry, f"Duplicate op name: {name}" + op_cls.name = name + cls.op_registry[name] = op_cls + return op_cls + + return decorator diff --git a/trainer/layers/layernorm.py b/trainer/layers/layernorm.py new file mode 100644 index 0000000000000000000000000000000000000000..e3d3301089fd45d0f043252e1c5f26283fdb9e0b --- /dev/null +++ b/trainer/layers/layernorm.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py +"""Custom normalization layers.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from trainer.layers.custom_op import CustomOp +from trainer.platforms import current_platform + + +@CustomOp.register("rms_norm") +class RMSNorm(CustomOp): + """Root mean square normalization. + + Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. + Refer to https://arxiv.org/abs/1910.07467 + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + var_hidden_size: int | None = None, + has_weight: bool = True, + ) -> None: + super().__init__() + + self.hidden_size = hidden_size + self.variance_epsilon = eps + self.variance_size_override = (None if var_hidden_size == hidden_size + else var_hidden_size) + self.has_weight = has_weight + + self.weight = torch.ones(hidden_size) if current_platform.is_cuda_alike( + ) else torch.ones(hidden_size, dtype=dtype) + if self.has_weight: + self.weight = nn.Parameter(self.weight) + + # if we do fully_shard(model.layer_norm), and we call layer_form.forward_native(input) instead of layer_norm(input), + # we need to call model.layer_norm.register_fsdp_forward_method(model, "forward_native") to make sure fsdp2 hooks are triggered + # for mixed precision and cpu offloading + + # the even better way might be fully_shard(model.layer_norm, mp_policy=, cpu_offloading=), and call model.layer_norm(input). everything should work out of the box + # because fsdp2 hooks will be triggered with model.layer_norm.__call__ + def forward_native( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """PyTorch-native implementation equivalent to forward().""" + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + hidden_size = x.shape[-1] + if hidden_size != self.hidden_size: + raise ValueError("Expected hidden_size to be " + f"{self.hidden_size}, but found: {hidden_size}") + + if self.variance_size_override is None: + x_var = x + else: + if hidden_size < self.variance_size_override: + raise ValueError( + "Expected hidden_size to be at least " + f"{self.variance_size_override}, but found: {hidden_size}") + + x_var = x[:, :, :self.variance_size_override] + + variance = x_var.pow(2).mean(dim=-1, keepdim=True) + + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = x.to(orig_dtype) + if self.has_weight: + x = x * self.weight + if residual is None: + return x + else: + return x, residual + + def extra_repr(self) -> str: + s = f"hidden_size={self.weight.data.size(0)}" + s += f", eps={self.variance_epsilon}" + return s + + +class ScaleResidual(nn.Module): + """ + Applies gated residual connection. + """ + + def __init__(self, prefix: str = ""): + super().__init__() + + def forward(self, residual: torch.Tensor, x: torch.Tensor, + gate: torch.Tensor) -> torch.Tensor: + """Apply gated residual connection.""" + return residual + x * gate + + +# adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py +# NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using +# FSDP's MixedPrecisionPolicy +class FP32LayerNorm(nn.LayerNorm): + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + return F.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +class ScaleResidualLayerNormScaleShift(nn.Module): + """ + Fused operation that combines: + 1. Gated residual connection + 2. LayerNorm + 3. Scale and shift operations + + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + if norm_type == "rms": + self.norm = RMSNorm(hidden_size, + has_weight=elementwise_affine, + eps=eps, + dtype=dtype) + elif norm_type == "layer": + if compute_dtype == torch.float32: + self.norm = FP32LayerNorm(hidden_size, + elementwise_affine=elementwise_affine, + eps=eps) + else: + self.norm = nn.LayerNorm(hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward(self, residual: torch.Tensor, x: torch.Tensor, + gate: torch.Tensor, shift: torch.Tensor, + scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply gated residual connection, followed by layernorm and + scale/shift in a single fused operation. + + Returns: + Tuple containing: + - normalized and modulated output + - residual value (value after residual connection + but before normalization) + """ + # Apply residual connection with gating + residual_output = residual + x * gate + # Apply normalization + normalized = self.norm(residual_output) + # Apply scale and shift + modulated = normalized * (1.0 + scale) + shift + return modulated, residual_output + + +class LayerNormScaleShift(nn.Module): + """ + Fused operation that combines LayerNorm with scale and shift operations. + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.compute_dtype = compute_dtype + if norm_type == "rms": + self.norm = RMSNorm(hidden_size, + has_weight=elementwise_affine, + eps=eps) + elif norm_type == "layer": + if self.compute_dtype == torch.float32: + self.norm = FP32LayerNorm(hidden_size, + elementwise_affine=elementwise_affine, + eps=eps) + else: + self.norm = nn.LayerNorm(hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward(self, x: torch.Tensor, shift: torch.Tensor, + scale: torch.Tensor) -> torch.Tensor: + """Apply ln followed by scale and shift in a single fused operation.""" + normalized = self.norm(x) + if self.compute_dtype == torch.float32: + return (normalized.float() * (1.0 + scale) + shift).to(x.dtype) + else: + return normalized * (1.0 + scale) + shift diff --git a/trainer/layers/linear.py b/trainer/layers/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..c36f7ebbb1396a70230cca90ae36b300e76632d0 --- /dev/null +++ b/trainer/layers/linear.py @@ -0,0 +1,960 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py + +from abc import abstractmethod + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from trainer.distributed import (divide, get_tp_rank, get_tp_world_size, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce) +from trainer.layers.quantization.base_config import (QuantizationConfig, + QuantizeMethodBase) +from trainer.logger import init_logger +# yapf: disable +from trainer.models.parameter import (BasevLLMParameter, + BlockQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, + PerTensorScaleParameter, + RowvLLMParameter) +# yapf: enable +from trainer.models.utils import set_weight_attrs + +logger = init_logger(__name__) + +WEIGHT_LOADER_V2_SUPPORTED = [ + "CompressedTensorsLinearMethod", "AWQMarlinLinearMethod", "AWQLinearMethod", + "GPTQMarlinLinearMethod", "Fp8LinearMethod", "MarlinLinearMethod", + "QQQLinearMethod", "GPTQMarlin24LinearMethod", "TPUInt8LinearMethod", + "GPTQLinearMethod", "FBGEMMFp8LinearMethod", "ModelOptFp8LinearMethod", + "IPEXAWQLinearMethod", "IPEXGPTQLinearMethod", "HQQMarlinMethod", + "QuarkLinearMethod" +] + + +def adjust_scalar_to_fused_array( + param: torch.Tensor, loaded_weight: torch.Tensor, + shard_id: str | int) -> tuple[torch.Tensor, torch.Tensor]: + """For fused modules (QKV and MLP) we have an array of length + N that holds 1 scale for each "logical" matrix. So the param + is an array of length N. The loaded_weight corresponds to + one of the shards on disk. Here, we slice the param based on + the shard_id for loading. + """ + qkv_idxs = {"q": 0, "k": 1, "v": 2} + + if isinstance(shard_id, str): + shard_id = qkv_idxs[shard_id] + elif not isinstance(shard_id, int): + raise ValueError(f"Unknown Shard Id {shard_id}") + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + return param[shard_id], loaded_weight + + +class LinearMethodBase(QuantizeMethodBase): + """Base class for different (maybe quantized) linear methods.""" + + @abstractmethod + def create_weights(self, layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], input_size: int, + output_size: int, params_dtype: torch.dtype, + **extra_weight_attrs) -> None: + """Create weights for a linear layer. + The weights will be set as attributes of the layer. + + Args: + layer: The layer that is using the LinearMethodBase factory. + input_size_per_partition: Size of the weight input dim on rank X. + output_partition_sizes: Sizes of the output dim of each logical + weight on rank X. E.g., output_partition_sizes for QKVLinear + is a list contains the width of Wq, Wk, Wv on rank X. + input_size: Size of the input dim of the weight across all ranks. + output_size: Size of the output dim of the weight across all ranks. + params_dtype: Datatype of the parameters. + """ + raise NotImplementedError + + @abstractmethod + def apply(self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError + + +class UnquantizedLinearMethod(LinearMethodBase): + """Linear method without quantization.""" + + def create_weights(self, layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], input_size: int, + output_size: int, params_dtype: torch.dtype, + **extra_weight_attrs) -> None: + weight = Parameter(torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply(self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None) -> torch.Tensor: + output = F.linear(x, layer.weight, bias) if torch.cuda.is_available( + ) or bias is None else F.linear( + x, layer.weight, bias.to(x.dtype) + ) # NOTE: this line assumes that we are using amp when using cuda and is needed to account for the fact that amp isn't supported in mps + return output + + +class LinearBase(torch.nn.Module): + """Base linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__( + self, + input_size: int, + output_size: int, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + + # Keep input parameters + self.input_size = input_size + self.output_size = output_size + self.skip_bias_add = skip_bias_add + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + self.quant_config = quant_config + self.prefix = prefix + if quant_config is None: + self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod( + ) + else: + self.quant_method = quant_config.get_quant_method(self, + prefix=prefix) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: + raise NotImplementedError + + +class ReplicatedLinear(LinearBase): + """Replicated linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__(self, + input_size: int, + output_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__(input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix=prefix) + + # All the linear layer supports quant method. + assert self.quant_method is not None + self.quant_method.create_weights(self, + self.input_size, [self.output_size], + self.input_size, + self.output_size, + self.params_dtype, + weight_loader=self.weight_loader) + + if bias: + self.bias = Parameter( + torch.empty( + self.output_size, + dtype=self.params_dtype, + )) + set_weight_attrs(self.bias, { + "output_dim": 0, + "weight_loader": self.weight_loader, + }) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, + loaded_weight: torch.Tensor) -> None: + # If the weight on disk does not have a shape, give it one + # (such scales for AutoFp8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param.size() == loaded_weight.size(), ( + f"Tried to load weights of size {loaded_weight.size()}" + f"to a parameter of size {param.size()}") + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: + bias = self.bias if not self.skip_bias_add else None + assert self.quant_method is not None + output = self.quant_method.apply(self, x, bias) + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + return s + + +class ColumnParallelLinear(LinearBase): + """Linear layer with column parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its second dimension as A = [A_1, ..., A_p]. + + Args: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make Y available + to all GPUs, otherwise, every GPU will have its output + which is Y_i = XA_i + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + output_sizes: list of output sizes packed into one output, like for QKV + the list would be size 3. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__(self, + input_size: int, + output_size: int, + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + output_sizes: list[int] | None = None, + prefix: str = ""): + # Divide the weight matrix along the last dimension. + self.tp_size = get_tp_world_size() + self.input_size_per_partition = input_size + self.output_size_per_partition = divide(output_size, self.tp_size) + self.output_partition_sizes = [self.output_size_per_partition] + # If QKV or MergedColumn, use output size of each partition. + if hasattr(self, "output_sizes"): + self.output_partition_sizes = [ + divide(output_size, self.tp_size) + for output_size in self.output_sizes + ] + + super().__init__(input_size, output_size, skip_bias_add, params_dtype, + quant_config, prefix) + + self.gather_output = gather_output + + if output_sizes is None: + output_sizes = [output_size] + + assert self.quant_method is not None + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 if self.quant_method.__class__.__name__ + in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader)) + if bias: + self.bias = Parameter( + torch.empty( + self.output_size_per_partition, + dtype=params_dtype, + )) + set_weight_attrs(self.bias, { + "output_dim": 0, + "weight_loader": self.weight_loader, + }) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, + loaded_weight: torch.Tensor) -> None: + tp_rank = get_tp_rank() + output_dim = getattr(param, "output_dim", None) + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + is_sharded_weight = is_sharded_weight + + param_data = param.data + if output_dim is not None and not is_sharded_weight: + shard_size = param_data.shape[output_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(output_dim, start_idx, + shard_size) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2(self, param: Parameter, + loaded_weight: torch.Tensor) -> None: + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + param.load_column_parallel_weight(loaded_weight=loaded_weight) + + def forward(self, + input_: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: + bias = self.bias if not self.skip_bias_add else None + + # Matrix multiply. + assert self.quant_method is not None + output_parallel = self.quant_method.apply(self, input_, bias) + if self.gather_output: + # All-gather across the partitions. + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size_per_partition}" + s += f", bias={self.bias is not None}" + s += f", tp_size={get_tp_world_size()}" + s += f", gather_output={self.gather_output}" + return s + + +class MergedColumnParallelLinear(ColumnParallelLinear): + """Packed linear layers with column parallelism. + + Similar to ColumnParallelLinear, but the weight matrix is concatenated + along the output dimension. When the weight matrix is loaded, the + different partitions are sharded separately. + + Args: + input_size: input dimension of the linear layer. + output_sizes: list of output dimensions of the linear layer. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make the output + available to all GPUs, otherwise, every GPU will have + its own output. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__(self, + input_size: int, + output_sizes: list[int], + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + self.output_sizes = output_sizes + tp_size = get_tp_world_size() + assert all(output_size % tp_size == 0 for output_size in output_sizes) + super().__init__(input_size=input_size, + output_size=sum(output_sizes), + bias=bias, + gather_output=gather_output, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix) + + def weight_loader(self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None) -> None: + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + # Special case for per-tensor scale to load scalar into fused array. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (mlp). + # (e.g., Phi-3's gate_up_proj). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + current_shard_offset = 0 + shard_offsets: list[tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + for shard_id, shard_offset, shard_size in shard_offsets: + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + assert loaded_shard_id < len(self.output_sizes) + tp_rank = get_tp_rank() + tp_size = get_tp_world_size() + if output_dim is not None: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size + shard_size = self.output_sizes[loaded_shard_id] // tp_size + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + start_idx = tp_rank * shard_size + if not is_sharded_weight: + loaded_weight = loaded_weight.narrow(output_dim, start_idx, + shard_size) + # Special case for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_offset = loaded_shard_id * shard_size + param_data = param_data.narrow(0, shard_offset, shard_size) + + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id) + + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "MergedColumnParallelLinear, assume the weight is " + "the same for all partitions.") + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def _load_fused_module_from_checkpoint(self, param: BasevLLMParameter, + loaded_weight: torch.Tensor) -> None: + """ + Handle special case for models where MLP layers are already + fused on disk. In this case, we have no shard id. This function + determmines the shard id by splitting these layers and then calls + the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + + current_shard_offset = 0 + shard_offsets: list[tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if isinstance(param, PackedColumnParameter | PackedvLLMParameter + ) and param.packed_dim == param.output_dim: + shard_size, shard_offset = \ + param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset) + + loaded_weight_shard = loaded_weight.narrow(param.output_dim, + shard_offset, shard_size) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2(self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None) -> None: + if loaded_shard_id is None: + if isinstance(param, PerTensorScaleParameter): + param.load_merged_column_weight(loaded_weight=loaded_weight, + shard_id=0) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_merged_column_weight(loaded_weight=loaded_weight) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id < len(self.output_sizes) + + tp_size = get_tp_world_size() + + if isinstance(param, BlockQuantScaleParameter): + raise NotImplementedError("FP8 is not implemented yet") + # FIXME(will): add fp8 support + # from vllm.model_executor.layers.quantization.fp8 import ( + # Fp8LinearMethod, Fp8MoEMethod) + # assert self.quant_method is not None + # assert isinstance(self.quant_method, + # (Fp8LinearMethod, Fp8MoEMethod)) + # weight_block_size = self.quant_method.quant_config.weight_block_size + # assert weight_block_size is not None + # block_n, _ = weight_block_size[0], weight_block_size[1] + # shard_offset = ( + # (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // + # block_n) // tp_size + # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) // + # block_n // tp_size) + else: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size + shard_size = self.output_sizes[loaded_shard_id] // tp_size + + param.load_merged_column_weight(loaded_weight=loaded_weight, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size) + + +class QKVParallelLinear(ColumnParallelLinear): + """Linear layers for the attention's QKV transformation. + + Linear layers for the linear transformation of the query, key, and value + vectors in the attention layer. The weight matrix is concatenated along + the output dimension. The layer is parallelized along the head dimension. + When the number of key/value heads is smaller than the number of query + heads (e.g., multi-query/grouped-query attention), the key/value head may + be replicated while the query heads are partitioned. + + Args: + hidden_size: input hidden state size of the transformer. + head_size: size of each attention head. + total_num_heads: total number of attention query heads. + total_num_kv_heads: total number of attention key/value heads. If + None, assume total_num_kv_heads = total_num_heads. + bias: If true, add bias. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__(self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int | None = None, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + self.hidden_size = hidden_size + self.head_size = head_size + self.total_num_heads = total_num_heads + if total_num_kv_heads is None: + total_num_kv_heads = total_num_heads + self.total_num_kv_heads = total_num_kv_heads + # Divide the weight matrix along the last dimension. + tp_size = get_tp_world_size() + self.num_heads = divide(self.total_num_heads, tp_size) + if tp_size >= self.total_num_kv_heads: + self.num_kv_heads = 1 + self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) + else: + self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) + self.num_kv_head_replicas = 1 + input_size = self.hidden_size + output_size = (self.num_heads + + 2 * self.num_kv_heads) * tp_size * self.head_size + self.output_sizes = [ + self.num_heads * self.head_size * tp_size, # q_proj + self.num_kv_heads * self.head_size * tp_size, # k_proj + self.num_kv_heads * self.head_size * tp_size, # v_proj + ] + + super().__init__(input_size=input_size, + output_size=output_size, + bias=bias, + gather_output=False, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix) + + def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: + shard_offset_mapping = { + "q": 0, + "k": self.num_heads * self.head_size, + "v": (self.num_heads + self.num_kv_heads) * self.head_size, + "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size + } + return shard_offset_mapping.get(loaded_shard_id) + + def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None: + shard_size_mapping = { + "q": self.num_heads * self.head_size, + "k": self.num_kv_heads * self.head_size, + "v": self.num_kv_heads * self.head_size, + } + return shard_size_mapping.get(loaded_shard_id) + + def _load_fused_module_from_checkpoint(self, param: BasevLLMParameter, + loaded_weight: torch.Tensor): + """ + Handle special case for models where QKV layers are already + fused on disk. In this case, we have no shard id. This function + determmines the shard id by splitting these layers and then calls + the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ("k", self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size), + ("v", + (self.total_num_heads + self.total_num_kv_heads) * self.head_size, + self.total_num_kv_heads * self.head_size), + ] + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if isinstance(param, PackedColumnParameter | PackedvLLMParameter + ) and param.packed_dim == param.output_dim: + shard_size, shard_offset = \ + param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset) + + loaded_weight_shard = loaded_weight.narrow(param.output_dim, + shard_offset, shard_size) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2(self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None): + if loaded_shard_id is None: # special case for certain models + if isinstance(param, PerTensorScaleParameter): + param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_qkv_weight(loaded_weight=loaded_weight) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id in ["q", "k", "v"] + + shard_offset = self._get_shard_offset_mapping(loaded_shard_id) + shard_size = self._get_shard_size_mapping(loaded_shard_id) + + param.load_qkv_weight(loaded_weight=loaded_weight, + num_heads=self.num_kv_head_replicas, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size) + + def weight_loader(self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None): + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + + # Special case for per-tensor scales in fused case. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (qkv). + # (e.g., Phi-3's qkv_proj). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ("k", self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size), + ("v", (self.total_num_heads + self.total_num_kv_heads) * + self.head_size, self.total_num_kv_heads * self.head_size), + ] + + for shard_id, shard_offset, shard_size in shard_offsets: + + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + tp_rank = get_tp_rank() + assert loaded_shard_id in ["q", "k", "v"] + + # If output dim is defined, use the default loading process. + if output_dim is not None: + if loaded_shard_id == "q": + shard_offset = 0 + shard_size = self.num_heads * self.head_size + elif loaded_shard_id == "k": + shard_offset = self.num_heads * self.head_size + shard_size = self.num_kv_heads * self.head_size + elif loaded_shard_id == "v": + shard_offset = (self.num_heads + + self.num_kv_heads) * self.head_size + shard_size = self.num_kv_heads * self.head_size + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + shard_idx = 0 + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + if loaded_shard_id == "q": + shard_idx = tp_rank + else: + shard_idx = tp_rank // self.num_kv_head_replicas + start_idx = shard_idx * shard_size + + if not is_sharded_weight: + loaded_weight = loaded_weight.narrow(output_dim, start_idx, + shard_size) + + # Special case for for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_index = ["q", "k", "v"].index(loaded_shard_id) + param_data = param_data.narrow(0, shard_index * shard_size, + shard_size) + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id) + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "QKVParallelLinear, assume the weight is the same " + "for all partitions.") + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class RowParallelLinear(LinearBase): + """Linear layer with row parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its first dimension and X along its second dimension as: + - - + | A_1 | + | . | + A = | . | X = [X_1, ..., X_p] + | . | + | A_p | + - - + Arguments: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. Note that bias is not parallelized. + input_is_parallel: If true, we assume that the input is already + split across the GPUs and we do not split + again. + skip_bias_add: This was added to enable performance optimization where + bias can be fused with other element-wise operations. + We skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__(self, + input_size: int, + output_size: int, + bias: bool = True, + input_is_parallel: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + reduce_results: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + # Divide the weight matrix along the first dimension. + self.tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.input_size_per_partition = divide(input_size, self.tp_size) + self.output_size_per_partition = output_size + self.output_partition_sizes = [output_size] + + super().__init__(input_size, output_size, skip_bias_add, params_dtype, + quant_config, prefix) + + self.input_is_parallel = input_is_parallel + self.reduce_results = reduce_results + + assert self.quant_method is not None + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 if self.quant_method.__class__.__name__ + in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader)) + if not reduce_results and (bias and not skip_bias_add): + raise ValueError("When not reduce the results, adding bias to the " + "results can lead to incorrect results") + + if bias: + self.bias = Parameter( + torch.empty(self.output_size, dtype=params_dtype)) + set_weight_attrs(self.bias, { + "output_dim": 0, + "weight_loader": self.weight_loader, + }) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + tp_rank = get_tp_rank() + input_dim = getattr(param, "input_dim", None) + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + param_data = param.data + if input_dim is not None and not is_sharded_weight: + shard_size = param_data.shape[input_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(input_dim, start_idx, + shard_size) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2(self, param: BasevLLMParameter, + loaded_weight: torch.Tensor): + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + + param.load_row_parallel_weight(loaded_weight=loaded_weight) + + def forward(self, input_) -> tuple[torch.Tensor, Parameter | None]: + if self.input_is_parallel: + input_parallel = input_ + else: + tp_rank = get_tp_rank() + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.tp_size) + input_parallel = splitted_input[tp_rank].contiguous() + + # Matrix multiply. + assert self.quant_method is not None + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in TP>1 case) + bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias + output_parallel = self.quant_method.apply(self, + input_parallel, + bias=bias_) + if self.reduce_results and self.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + + output_bias = self.bias if self.skip_bias_add else None + + return output, output_bias + + def extra_repr(self) -> str: + s = f"input_features={self.input_size_per_partition}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + s += f", tp_size={self.tp_size}" + s += f", reduce_results={self.reduce_results}" + return s diff --git a/trainer/layers/lora/linear.py b/trainer/layers/lora/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..f67395f3455642c8377327859e27099e069c9d93 --- /dev/null +++ b/trainer/layers/lora/linear.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py + +import math + +import torch +from torch import nn +from torch.distributed._composable.fsdp import (CPUOffloadPolicy, OffloadPolicy, + fully_shard) +from torch.distributed.tensor import DTensor + +from trainer.distributed import (get_local_torch_device, get_tp_rank, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce) +from trainer.layers.linear import (ColumnParallelLinear, LinearBase, + MergedColumnParallelLinear, + QKVParallelLinear, ReplicatedLinear, + RowParallelLinear) +from trainer.layers.vocab_parallel_embedding import VocabParallelEmbedding +from trainer.utils import get_mixed_precision_state + +torch._dynamo.config.recompile_limit = 16 + + +class BaseLayerWithLoRA(nn.Module): + + def __init__( + self, + base_layer: nn.Module, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ): + super().__init__() + self.base_layer: nn.Module = base_layer + + self.merged: bool = False + self.cpu_weight = base_layer.weight.to("cpu") + # indicates adapter weights don't contain this layer + # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging) + self.disable_lora: bool = False + self.lora_rank = lora_rank + self.lora_alpha = lora_alpha + self.training_mode = training_mode + self.lora_path: str | None = None + + if training_mode: + assert self.lora_rank is not None, "LoRA rank must be set for training mode" + if self.lora_rank is None or self.lora_alpha is None: + self.lora_alpha = lora_rank + self.base_layer.requires_grad_(False) + in_dim = self.base_layer.weight.shape[1] + out_dim = self.base_layer.weight.shape[0] + self.lora_A = nn.Parameter( + torch.zeros(self.lora_rank, + in_dim, + device=self.base_layer.weight.device, + dtype=self.base_layer.weight.dtype)) + self.lora_B = nn.Parameter( + torch.zeros(out_dim, + self.lora_rank, + device=self.base_layer.weight.device, + dtype=self.base_layer.weight.dtype)) + torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + torch.nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5)) + else: + self.lora_A = None + self.lora_B = None + + @torch.compile() + def forward(self, x: torch.Tensor) -> torch.Tensor: + lora_A = self.lora_A + lora_B = self.lora_B + if isinstance(self.lora_B, DTensor): + lora_B = self.lora_B.to_local() + lora_A = self.lora_A.to_local() + + if (self.training_mode or not self.merged) and not self.disable_lora: + delta = x @ ( + self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) + @ self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))) + if self.lora_alpha != self.lora_rank: + delta = delta * ( + self.lora_alpha / self.lora_rank # type: ignore + ) # type: ignore + out, output_bias = self.base_layer(x) + return out + delta, output_bias + else: + out, output_bias = self.base_layer(x) + return out.to(x), output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + return B + + def set_lora_weights(self, + A: torch.Tensor, + B: torch.Tensor, + training_mode: bool = False, + lora_path: str | None = None) -> None: + self.lora_A = torch.nn.Parameter( + A) # share storage with weights in the pipeline + self.lora_B = torch.nn.Parameter(B) + self.disable_lora = False + if not training_mode: + self.merge_lora_weights() + self.lora_path = lora_path + + @torch.no_grad() + def merge_lora_weights(self) -> None: + if self.disable_lora: + return + + if self.merged: + self.unmerge_lora_weights() + assert self.lora_A is not None and self.lora_B is not None, "LoRA weights not set. Please set them first." + if isinstance(self.base_layer.weight, DTensor): + mesh = self.base_layer.weight.data.device_mesh + unsharded_base_layer = ReplicatedLinear( + input_size=self.base_layer.input_size, + output_size=self.base_layer.output_size, + bias=getattr(self.base_layer, "bias", None) is not None, + skip_bias_add=self.base_layer.skip_bias_add, + params_dtype=self.base_layer.params_dtype, + quant_config=self.base_layer.quant_config, + prefix=self.base_layer.prefix, + ) + # Using offload param is on CPU, so current_device is for "CPU -> GPU -> merge -> CPU" + current_device = self.base_layer.weight.data.device + data = self.base_layer.weight.data.to( + get_local_torch_device()).full_tensor() + data += (self.slice_lora_b_weights(self.lora_B).to(data) + @ self.slice_lora_a_weights(self.lora_A).to(data)) + unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + if isinstance(getattr(self.base_layer, "bias", None), DTensor): + unsharded_base_layer.bias = nn.Parameter( + self.base_layer.bias.to( + get_local_torch_device(), + non_blocking=True).full_tensor().to(current_device)) + + offload_policy = CPUOffloadPolicy() if "cpu" in str( + current_device) else OffloadPolicy() + mp_policy = get_mixed_precision_state().mp_policy + + self.base_layer = fully_shard(unsharded_base_layer, + mesh=mesh, + mp_policy=mp_policy, + offload_policy=offload_policy) + else: + current_device = self.base_layer.weight.data.device + data = self.base_layer.weight.data.to(get_local_torch_device()) + data += \ + (self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data))) + self.base_layer.weight.data = data.to(current_device, + non_blocking=True) + + self.merged = True + + @torch.no_grad() + # @torch.compile(dynamic=True) + def unmerge_lora_weights(self) -> None: + if self.disable_lora: + return + + if not self.merged: + raise ValueError( + "LoRA weights not merged. Please merge them first before unmerging." + ) + + # avoid precision loss + if isinstance(self.base_layer.weight, DTensor): + device = self.base_layer.weight.data.device + self.base_layer.weight = nn.Parameter( + self.cpu_weight.to(device, non_blocking=True)) + else: + self.base_layer.weight.data = self.cpu_weight.data.to( + self.base_layer.weight, non_blocking=True) + + self.merged = False + + +class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): + """ + Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation). + + Note: The current version does not yet implement the LoRA functionality. + This class behaves exactly the same as the base VocabParallelEmbedding. + Future versions will integrate LoRA functionality to support efficient parameter fine-tuning. + """ + + def __init__( + self, + base_layer: VocabParallelEmbedding, + ) -> None: + super().__init__(base_layer) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + raise NotImplementedError( + "We don't support VocabParallelEmbeddingWithLoRA yet.") + + +class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): + + def __init__( + self, + base_layer: ColumnParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + # duplicate the logic in ColumnParallelLinear + bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_, bias) + if self.base_layer.gather_output: + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + shard_size = self.base_layer.output_partition_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + B = B[start_idx:end_idx, :] + return B + + +class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + + def __init__( + self, + base_layer: MergedColumnParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A.to(self.base_layer.weight) + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + # Since the outputs for both gate and up are identical, we use a random one. + shard_size = self.base_layer.output_partition_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + return B[:, start_idx:end_idx, :] + + +class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + + def __init__( + self, + base_layer: QKVParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights( + self, B: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + tp_rank = get_tp_rank() + B_q, B_kv = B + base_layer = self.base_layer + q_proj_shard_size = base_layer.q_proj_shard_size + kv_proj_shard_size = base_layer.kv_proj_shard_size + num_kv_head_replicas = base_layer.num_kv_head_replicas + + q_start_idx = q_proj_shard_size * tp_rank + q_end_idx = q_start_idx + q_proj_shard_size + + kv_shard_id = tp_rank // num_kv_head_replicas + kv_start_idx = kv_proj_shard_size * kv_shard_id + kv_end_idx = kv_start_idx + kv_proj_shard_size + + return B_q[q_start_idx:q_end_idx, :], B_kv[:, + kv_start_idx:kv_end_idx, :] + + +class RowParallelLinearWithLoRA(BaseLayerWithLoRA): + + def __init__( + self, + base_layer: RowParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def forward(self, input_: torch.Tensor): + # duplicate the logic in RowParallelLinear + if self.base_layer.input_is_parallel: + input_parallel = input_ + else: + tp_rank = get_tp_rank() + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.base_layer.tp_size) + input_parallel = splitted_input[tp_rank].contiguous() + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_parallel) + + if self.set_lora: + output_parallel = self.apply_lora(output_parallel, input_parallel) + + if self.base_layer.reduce_results and self.base_layer.tp_size > 1: + output_ = tensor_model_parallel_all_reduce(output_parallel) + else: + output_ = output_parallel + + if not self.base_layer.skip_bias_add: + output = (output_ + self.base_layer.bias + if self.base_layer.bias is not None else output_) + output_bias = None + else: + output = output_ + output_bias = self.base_layer.bias + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + shard_size = self.base_layer.input_size_per_partition + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + A = A[:, start_idx:end_idx].contiguous() + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + return B + + +def get_lora_layer(layer: nn.Module, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False) -> BaseLayerWithLoRA | None: + supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = { + # the order matters + # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, + QKVParallelLinear: QKVParallelLinearWithLoRA, + MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA, + ColumnParallelLinear: ColumnParallelLinearWithLoRA, + RowParallelLinear: RowParallelLinearWithLoRA, + ReplicatedLinear: BaseLayerWithLoRA, + } + for src_layer_type, lora_layer_type in supported_layer_types.items(): + if isinstance(layer, src_layer_type): # pylint: disable=unidiomatic-typecheck + ret = lora_layer_type(layer, + lora_rank=lora_rank, + lora_alpha=lora_alpha, + training_mode=training_mode) + return ret + return None + + +# source: https://github.com/vllm-project/vllm/blob/93b38bea5dd03e1b140ca997dfaadef86f8f1855/vllm/lora/utils.py#L9 +def replace_submodule(model: nn.Module, module_name: str, + new_module: nn.Module) -> nn.Module: + """Replace a submodule in a model with a new module.""" + parent = model.get_submodule(".".join(module_name.split(".")[:-1])) + target_name = module_name.split(".")[-1] + setattr(parent, target_name, new_module) + return new_module diff --git a/trainer/layers/mlp.py b/trainer/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..a9f1c4491dd0d5e0f0f7c58d4aebb13209365f17 --- /dev/null +++ b/trainer/layers/mlp.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn + +from trainer.layers.activation import get_act_fn +from trainer.layers.linear import ReplicatedLinear + + +class MLP(nn.Module): + """ + MLP for DiT blocks, NO gated linear units + """ + + def __init__( + self, + input_dim: int, + mlp_hidden_dim: int, + output_dim: int | None = None, + bias: bool = True, + act_type: str = "gelu_pytorch_tanh", + dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.fc_in = ReplicatedLinear( + input_dim, + mlp_hidden_dim, # For activation func like SiLU that need 2x width + bias=bias, + params_dtype=dtype) + + self.act = get_act_fn(act_type) + if output_dim is None: + output_dim = input_dim + self.fc_out = ReplicatedLinear(mlp_hidden_dim, + output_dim, + bias=bias, + params_dtype=dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc_in(x) + x = self.act(x) + x, _ = self.fc_out(x) + return x diff --git a/trainer/layers/quantization/__init__.py b/trainer/layers/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1fa0226c1277e1c1cc75cc88a3c26c82621b94da --- /dev/null +++ b/trainer/layers/quantization/__init__.py @@ -0,0 +1,65 @@ +from typing import Literal, get_args + +from trainer.layers.quantization.base_config import QuantizationConfig + +QuantizationMethods = Literal[None] + +QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) + +# The customized quantization methods which will be added to this dict. +_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {} + + +def register_quantization_config(quantization: str): + """Register a customized vllm quantization config. + + When a quantization method is not supported by vllm, you can register a customized + quantization config to support it. + + Args: + quantization (str): The quantization method name. + + Examples: + >>> from trainer.layers.quantization import register_quantization_config + >>> from trainer.layers.quantization import get_quantization_config + >>> from trainer.layers.quantization.base_config import QuantizationConfig + >>> + >>> @register_quantization_config("my_quant") + ... class MyQuantConfig(QuantizationConfig): + ... pass + >>> + >>> get_quantization_config("my_quant") + + """ # noqa: E501 + + def _wrapper(quant_config_cls): + if quantization in QUANTIZATION_METHODS: + raise ValueError( + f"The quantization method `{quantization}` is already exists.") + if not issubclass(quant_config_cls, QuantizationConfig): + raise ValueError("The quantization config must be a subclass of " + "`QuantizationConfig`.") + _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls + QUANTIZATION_METHODS.append(quantization) + return quant_config_cls + + return _wrapper + + +def get_quantization_config(quantization: str) -> type[QuantizationConfig]: + if quantization not in QUANTIZATION_METHODS: + raise ValueError(f"Invalid quantization method: {quantization}") + + method_to_config: dict[str, type[QuantizationConfig]] = {} + # Update the `method_to_config` with customized quantization methods. + method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) + + return method_to_config[quantization] + + +all = [ + "QuantizationMethods", + "QuantizationConfig", + "get_quantization_config", + "QUANTIZATION_METHODS", +] diff --git a/trainer/layers/quantization/base_config.py b/trainer/layers/quantization/base_config.py new file mode 100644 index 0000000000000000000000000000000000000000..f4ca24c034813563ef5aabb909cd2b1977286b5a --- /dev/null +++ b/trainer/layers/quantization/base_config.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py + +import inspect +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +if TYPE_CHECKING: + from trainer.layers.quantization import QuantizationMethods +else: + QuantizationMethods = str + + +class QuantizeMethodBase(ABC): + """Base class for different quantized methods.""" + + @abstractmethod + def create_weights(self, layer: torch.nn.Module, *weight_args, + **extra_weight_attrs): + """Create weights for a layer. + + The weights will be set as attributes of the layer.""" + raise NotImplementedError + + @abstractmethod + def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError + + # Not required functions + def embedding(self, layer: torch.nn.Module, *args, + **kwargs) -> torch.Tensor: + """Gather embeddings in the layer based on indices in the input tensor. + + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Process the weight after loading. + + This can be used for example, to transpose weights for computation. + """ + return + + +def method_has_implemented_embedding( + method_class: type[QuantizeMethodBase]) -> bool: + """ + Not all quant methods have embedding implemented, so we need to check that + it exists for our given method. We check this by making sure the function + has been changed from the base implementation. + """ + base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", + None) + class_embedding = inspect.getattr_static(method_class, "embedding", None) + + return (class_embedding is not None + and class_embedding is not base_embedding) + + +class QuantizationConfig(ABC): + """Base class for quantization configs.""" + + def __init__(self): + super().__init__() + # mapping is updated by models as they initialize + self.packed_modules_mapping: dict[str, list[str]] = dict() + + @abstractmethod + def get_name(self) -> QuantizationMethods: + """Name of the quantization method.""" + raise NotImplementedError + + @abstractmethod + def get_supported_act_dtypes(self) -> list[torch.dtype]: + """List of supported activation dtypes.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """Minimum GPU capability to support the quantization method. + + E.g., 70 for Volta, 75 for Turing, 80 for Ampere. + This requirement is due to the custom CUDA kernels used by the + quantization method. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_config_filenames() -> list[str]: + """List of filenames to search for in the model directory.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig": + """Create a config class from the model's quantization config.""" + raise NotImplementedError + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, + user_quant) -> QuantizationMethods | None: + """ + Detects if this quantization method can support a given checkpoint + format by overriding the user specified quantization method -- + this method should only be overwritten by subclasses in exceptional + circumstances + """ + return None + + @staticmethod + def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any: + """Get a value from the model's quantization config.""" + for key in keys: + if key in config: + return config[key] + raise ValueError(f"Cannot find any of {keys} in the model's " + "quantization config.") + + @staticmethod + def get_from_keys_or(config: dict[str, Any], keys: list[str], + default: Any) -> Any: + """Get a optional value from the model's quantization config.""" + try: + return QuantizationConfig.get_from_keys(config, keys) + except ValueError: + return default + + @abstractmethod + def get_quant_method(self, layer: torch.nn.Module, + prefix: str) -> QuantizeMethodBase | None: + """Get the quantize method to use for the quantized layer. + + Args: + layer: The layer for the quant method. + prefix: The full name of the layer in the state dict + Returns: + The quantize method. None if the given layer doesn't support quant + method. + """ + raise NotImplementedError + + def get_cache_scale(self, name: str) -> str | None: + return None diff --git a/trainer/layers/rotary_embedding.py b/trainer/layers/rotary_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..893a7398bcdfb8308239da5c866f221a2e56e861 --- /dev/null +++ b/trainer/layers/rotary_embedding.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Rotary Positional Embeddings.""" +from typing import Any + +import torch + +from trainer.distributed.parallel_state import get_sp_group +from trainer.layers.custom_op import CustomOp + + +def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + +def _apply_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + is_neox_style: bool, +) -> torch.Tensor: + """ + Args: + x: [num_tokens, num_heads, head_size] + cos: [num_tokens, head_size // 2] + sin: [num_tokens, head_size // 2] + is_neox_style: Whether to use the Neox-style or GPT-J-style rotary + positional embeddings. + """ + # cos = cos.unsqueeze(-2).to(x.dtype) + # sin = sin.unsqueeze(-2).to(x.dtype) + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + if is_neox_style: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = (x1.float() * cos - x2.float() * sin).type_as(x) + o2 = (x2.float() * cos + x1.float() * sin).type_as(x) + if is_neox_style: + return torch.cat((o1, o2), dim=-1) + else: + return torch.stack((o1, o2), dim=-1).flatten(-2) + + +@CustomOp.register("rotary_embedding") +class RotaryEmbedding(CustomOp): + """Original rotary positional embedding.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int | float, + is_neox_style: bool, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.head_size = head_size + self.rotary_dim = rotary_dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.is_neox_style = is_neox_style + self.dtype = dtype + + cache = self._compute_cos_sin_cache() + cache = cache.to(dtype) + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) + + def _compute_inv_freq(self, base: int | float) -> torch.Tensor: + """Compute the inverse frequency.""" + # NOTE(woosuk): To exactly match the HF implementation, we need to + # use CPU to compute the cache and then move it to GPU. However, we + # create the cache on GPU for faster initialization. This may cause + # a slight numerical difference between the HF implementation and ours. + inv_freq = 1.0 / (base**(torch.arange( + 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim)) + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" + inv_freq = self._compute_inv_freq(self.base) + t = torch.arange(self.max_position_embeddings, dtype=torch.float) + + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """A PyTorch-native implementation of forward().""" + if offsets is not None: + positions = positions + offsets + positions = positions.flatten() + num_tokens = positions.shape[0] + cos_sin = self.cos_sin_cache.index_select(0, positions) + cos, sin = cos_sin.chunk(2, dim=-1) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., :self.rotary_dim] + query_pass = query[..., self.rotary_dim:] + query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., :self.rotary_dim] + key_pass = key[..., self.rotary_dim:] + key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + return s + + +def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]: + if isinstance(x, int): + return (x, ) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start: int | tuple[int, ...], + *args: int | tuple[int, ...], + dim: int = 2) -> torch.Tensor: + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0, ) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = tuple(stop[i] - start[i] for i in range(dim)) + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +def get_1d_rotary_pos_embed( + dim: int, + pos: torch.FloatTensor | int, + theta: float = 10000.0, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0. + + Returns: + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor**(dim / (dim - 2)) + + freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim) + ) # [D/2] + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + freqs_cos = freqs.cos() # [S, D/2] + freqs_sin = freqs.sin() # [S, D/2] + return freqs_cos, freqs_sin + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + theta_rescale_factor: float | list[float] = 1.0, + interpolation_factor: float | list[float] = 1.0, + shard_dim: int = 0, + sp_rank: int = 0, + sp_world_size: int = 1, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + Supports sequence parallelism by allowing sharding of a specific dimension. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + interpolation_factor (float): Factor to scale positions. Defaults to 1.0. + shard_dim (int): Which dimension to shard for sequence parallelism. Defaults to 0. + sp_rank (int): Rank in the sequence parallel group. Defaults to 0. + sp_world_size (int): World size of the sequence parallel group. Defaults to 1. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D/2] + """ + # Get the full grid + full_grid = get_meshgrid_nd( + start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + + # Shard the grid if using sequence parallelism (sp_world_size > 1) + assert shard_dim < len( + rope_dim_list + ), f"shard_dim {shard_dim} must be less than number of dimensions {len(rope_dim_list)}" + if sp_world_size > 1: + # Get the shape of the full grid + grid_shape = list(full_grid.shape[1:]) + + # Ensure the dimension to shard is divisible by sp_world_size + assert grid_shape[shard_dim] % sp_world_size == 0, ( + f"Dimension {shard_dim} with size {grid_shape[shard_dim]} is not divisible " + f"by sequence parallel world size {sp_world_size}") + + # Compute the start and end indices for this rank's shard + shard_size = grid_shape[shard_dim] // sp_world_size + start_idx = sp_rank * shard_size + end_idx = (sp_rank + 1) * shard_size + + # Create slicing indices for each dimension + slice_indices = [slice(None) for _ in range(len(grid_shape))] + slice_indices[shard_dim] = slice(start_idx, end_idx) + + # Shard the grid + # Update grid shape for the sharded dimension + grid_shape[shard_dim] = grid_shape[shard_dim] // sp_world_size + grid = torch.empty((len(rope_dim_list), ) + tuple(grid_shape), + dtype=full_grid.dtype) + for i in range(len(rope_dim_list)): + grid[i] = full_grid[i][tuple(slice_indices)] + else: + grid = full_grid + + if isinstance(theta_rescale_factor, int | float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, + list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int | float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, + list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + dtype=dtype, + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + + +def get_rotary_pos_embed( + rope_sizes, + hidden_size, + heads_num, + rope_dim_list, + rope_theta, + theta_rescale_factor=1.0, + interpolation_factor=1.0, + shard_dim: int = 0, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Generate rotary positional embeddings for the given sizes. + + Args: + rope_sizes: Tuple of dimensions (t, h, w) + hidden_size: Hidden dimension size + heads_num: Number of attention heads + rope_dim_list: List of dimensions for each axis, or None + rope_theta: Base for frequency calculations + theta_rescale_factor: Rescale factor for theta. Defaults to 1.0 + interpolation_factor: Factor to scale positions. Defaults to 1.0 + shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0. + + Returns: + Tuple of (cos, sin) tensors for rotary embeddings + """ + + target_ndim = 3 + head_dim = hidden_size // heads_num + + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + + assert sum( + rope_dim_list + ) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer" + + # Get SP info + sp_group = get_sp_group() + sp_rank = sp_group.rank_in_group + sp_world_size = sp_group.world_size + + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=rope_theta, + theta_rescale_factor=theta_rescale_factor, + interpolation_factor=interpolation_factor, + shard_dim=shard_dim, + sp_rank=sp_rank, + sp_world_size=sp_world_size, + dtype=dtype, + ) + return freqs_cos, freqs_sin + + +_ROPE_DICT: dict[tuple, RotaryEmbedding] = {} + + +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: int | float, + is_neox_style: bool = True, + rope_scaling: dict[str, Any] | None = None, + dtype: torch.dtype | None = None, + partial_rotary_factor: float = 1.0, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = (head_size, rotary_dim, max_position, base, is_neox_style, + rope_scaling_args, dtype) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + if rope_scaling is None: + rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base, + is_neox_style, dtype) + else: + raise ValueError(f"Unknown RoPE scaling {rope_scaling}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb diff --git a/trainer/layers/utils.py b/trainer/layers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e249b616b47e04345d994619473471ac70f3dee8 --- /dev/null +++ b/trainer/layers/utils.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py +"""Utility methods for model layers.""" + +import torch + + +def get_token_bin_counts_and_mask( + tokens: torch.Tensor, + vocab_size: int, + num_seqs: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # Compute the bin counts for the tokens. + # vocab_size + 1 for padding. + bin_counts = torch.zeros((num_seqs, vocab_size + 1), + dtype=torch.long, + device=tokens.device) + bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) + bin_counts = bin_counts[:, :vocab_size] + mask = bin_counts > 0 + + return bin_counts, mask diff --git a/trainer/layers/visual_embedding.py b/trainer/layers/visual_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..016e2e29e33483c21c2d79a8ccb13ed6abfa408a --- /dev/null +++ b/trainer/layers/visual_embedding.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch +import torch.nn as nn + +from trainer.layers.activation import get_act_fn +from trainer.layers.linear import ReplicatedLinear +from trainer.layers.mlp import MLP + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__(self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + prefix: str = ""): + super().__init__() + # Convert patch_size to 2-tuple + if isinstance(patch_size, list | tuple): + if len(patch_size) == 1: + patch_size = (patch_size[0], patch_size[0]) + else: + patch_size = (patch_size, patch_size) + + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d(in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + dtype=dtype) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer="silu", + frequency_embedding_size=256, + max_period=10000, + dtype=None, + freq_dtype=torch.float32, + prefix: str = "", + ): + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + + self.mlp = MLP(frequency_embedding_size, + hidden_size, + hidden_size, + act_type=act_layer, + dtype=dtype) + self.freq_dtype = freq_dtype + + def forward(self, t: torch.Tensor) -> torch.Tensor: + t_freq = timestep_embedding(t, + self.frequency_embedding_size, + self.max_period, + dtype=self.freq_dtype).to( + self.mlp.fc_in.weight.dtype) + # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +def timestep_embedding(t: torch.Tensor, + dim: int, + max_period: int = 10000, + dtype: torch.dtype = torch.float32) -> torch.Tensor: + """ + Create sinusoidal timestep embeddings. + + Args: + t: Tensor of shape [B] with timesteps + dim: Embedding dimension + max_period: Controls the minimum frequency of the embeddings + + Returns: + Tensor of shape [B, dim] with embeddings + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * + torch.arange(start=0, end=half, dtype=dtype) / + half).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class ModulateProjection(nn.Module): + """Modulation layer for DiT blocks.""" + + def __init__( + self, + hidden_size: int, + factor: int = 2, + act_layer: str = "silu", + dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.factor = factor + self.hidden_size = hidden_size + self.linear = ReplicatedLinear(hidden_size, + hidden_size * factor, + bias=True, + params_dtype=dtype) + self.act = get_act_fn(act_layer) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.act(x) + x, _ = self.linear(x) + return x + + +def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor: + """ + Convert patched representation back to image space. + + Args: + x: Tensor of shape [B, T*H*W, C*P_t*P_h*P_w] + t, h, w: Temporal and spatial dimensions + + Returns: + Unpatchified tensor of shape [B, C, T*P_t, H*P_h, W*P_w] + """ + assert x.ndim == 3, f"x.ndim: {x.ndim}" + assert len(patch_size) == 3, f"patch_size: {patch_size}" + assert t * h * w == x.shape[ + 1], f"t * h * w: {t * h * w}, x.shape[1]: {x.shape[1]}" + c = channels + pt, ph, pw = patch_size + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs diff --git a/trainer/layers/vocab_parallel_embedding.py b/trainer/layers/vocab_parallel_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..3d33589802df0ce5a7de2a44d9112e9f62cef16e --- /dev/null +++ b/trainer/layers/vocab_parallel_embedding.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter, UninitializedParameter + +from trainer.distributed import (divide, get_tp_rank, get_tp_world_size, + tensor_model_parallel_all_reduce) +from trainer.layers.quantization.base_config import ( + QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding) +from trainer.models.parameter import BasevLLMParameter +from trainer.models.utils import set_weight_attrs +from trainer.platforms import current_platform + +DEFAULT_VOCAB_PADDING_SIZE = 64 + + +class UnquantizedEmbeddingMethod(QuantizeMethodBase): + """Unquantized method for embeddings.""" + + def create_weights(self, layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], input_size: int, + output_size: int, params_dtype: torch.dtype, + **extra_weight_attrs): + """Create weights for embedding layer.""" + + weight = Parameter(torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply(self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None) -> torch.Tensor: + return F.linear(x, layer.weight, bias) + + def embedding(self, layer: torch.nn.Module, + input_: torch.Tensor) -> torch.Tensor: + return F.embedding(input_, layer.weight) + + +def pad_vocab_size(vocab_size: int, + pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int: + """Pad the vocab size to the given value.""" + return ((vocab_size + pad_to - 1) // pad_to) * pad_to + + +def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size: int, + rank: int, + offset: int = 0) -> Sequence[int]: + index_f = rank * per_partition_vocab_size + index_l = index_f + per_partition_vocab_size + return index_f + offset, index_l + offset + + +def vocab_range_from_global_vocab_size(global_vocab_size: int, + rank: int, + world_size: int, + offset: int = 0) -> Sequence[int]: + per_partition_vocab_size = divide(global_vocab_size, world_size) + return vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, + rank, + offset=offset) + + +@dataclass +class VocabParallelEmbeddingShardIndices: + """Indices for a shard of a vocab parallel embedding.""" + padded_org_vocab_start_index: int + padded_org_vocab_end_index: int + padded_added_vocab_start_index: int + padded_added_vocab_end_index: int + + org_vocab_start_index: int + org_vocab_end_index: int + added_vocab_start_index: int + added_vocab_end_index: int + + @property + def num_org_elements(self) -> int: + return self.org_vocab_end_index - self.org_vocab_start_index + + @property + def num_added_elements(self) -> int: + return self.added_vocab_end_index - self.added_vocab_start_index + + @property + def num_org_elements_padded(self) -> int: + return (self.padded_org_vocab_end_index - + self.padded_org_vocab_start_index) + + @property + def num_added_elements_padded(self) -> int: + return (self.padded_added_vocab_end_index - + self.padded_added_vocab_start_index) + + @property + def num_org_vocab_padding(self) -> int: + return self.num_org_elements_padded - self.num_org_elements + + @property + def num_added_vocab_padding(self) -> int: + return self.num_added_elements_padded - self.num_added_elements + + @property + def num_elements_padded(self) -> int: + return self.num_org_elements_padded + self.num_added_elements_padded + + def __post_init__(self): + # sanity checks + assert (self.padded_org_vocab_start_index + <= self.padded_org_vocab_end_index) + assert (self.padded_added_vocab_start_index + <= self.padded_added_vocab_end_index) + + assert self.org_vocab_start_index <= self.org_vocab_end_index + assert self.added_vocab_start_index <= self.added_vocab_end_index + + assert self.org_vocab_start_index <= self.padded_org_vocab_start_index + assert (self.added_vocab_start_index + <= self.padded_added_vocab_start_index) + assert self.org_vocab_end_index <= self.padded_org_vocab_end_index + assert self.added_vocab_end_index <= self.padded_added_vocab_end_index + + assert self.num_org_elements <= self.num_org_elements_padded + assert self.num_added_elements <= self.num_added_elements_padded + + +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) +def get_masked_input_and_mask( + input_: torch.Tensor, org_vocab_start_index: int, + org_vocab_end_index: int, num_org_vocab_padding: int, + added_vocab_start_index: int, + added_vocab_end_index: int) -> tuple[torch.Tensor, torch.Tensor]: + # torch.compile will fuse all of the pointwise ops below + # into a single kernel, making it very fast + org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ + < org_vocab_end_index) + added_vocab_mask = (input_ >= added_vocab_start_index) & ( + input_ < added_vocab_end_index) + added_offset = added_vocab_start_index - ( + org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding + valid_offset = (org_vocab_start_index * org_vocab_mask) + (added_offset * + added_vocab_mask) + vocab_mask = org_vocab_mask | added_vocab_mask + input_ = vocab_mask * (input_ - valid_offset) + return input_, ~vocab_mask + + +class VocabParallelEmbedding(torch.nn.Module): + """Embedding parallelized in the vocabulary dimension. + + Adapted from torch.nn.Embedding, note that we pad the vocabulary size to + make sure it is divisible by the number of model parallel GPUs. + + In order to support various loading methods, we ensure that LoRA-added + embeddings are always at the end of TP-sharded tensors. In other words, + we shard base embeddings and LoRA embeddings separately (both padded), + and place them in the same tensor. + In this example, we will have the original vocab size = 1010, + added vocab size = 16 and padding to 64. Therefore, the total + vocab size with padding will be 1088 (because we first pad 1010 to + 1024, add 16, and then pad to 1088). + Therefore, the tensor format looks like the following: + TP1, rank 0 (no sharding): + |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| + corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 | + + TP2, rank 0: + |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| + corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | + TP2, rank 1: + |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| + corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 | + + Args: + num_embeddings: vocabulary size. + embedding_dim: size of hidden state. + params_dtype: type of the parameters. + org_num_embeddings: original vocabulary size (without LoRA). + padding_size: padding size for the vocabulary. + quant_config: quant config for the layer + prefix: full name of the layer in the state dict + """ # noqa: E501 + + def __init__(self, + num_embeddings: int, + embedding_dim: int, + params_dtype: torch.dtype | None = None, + org_num_embeddings: int | None = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + + # Keep the input dimensions. + tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.num_embeddings = num_embeddings + self.padding_size = padding_size + self.org_vocab_size = org_num_embeddings or num_embeddings + num_added_embeddings = num_embeddings - self.org_vocab_size + self.org_vocab_size_padded = pad_vocab_size(self.org_vocab_size, + self.padding_size) + self.num_embeddings_padded = pad_vocab_size( + self.org_vocab_size_padded + num_added_embeddings, + self.padding_size) + assert self.org_vocab_size_padded <= self.num_embeddings_padded + + self.shard_indices = self._get_indices(self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, tp_rank, + self.tp_size) + self.embedding_dim = embedding_dim + + quant_method = None + if quant_config is not None: + quant_method = quant_config.get_quant_method(self, prefix=prefix) + if quant_method is None: + quant_method = UnquantizedEmbeddingMethod() + + # If we are making an embedding layer, then our quantization linear + # method must implement the embedding operation. If we are another + # layer type like ParallelLMHead, this is not important. + is_embedding_layer = type(self.__class__) is VocabParallelEmbedding + quant_method_implements_embedding = method_has_implemented_embedding( + type(quant_method)) + if is_embedding_layer and not quant_method_implements_embedding: + raise NotImplementedError( + f"The class {type(quant_method).__name__} must implement " + "the 'embedding' method, see UnquantizedEmbeddingMethod.") + + self.quant_method: QuantizeMethodBase = quant_method + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + # Divide the weight matrix along the vocaburaly dimension. + self.num_added_embeddings = self.num_embeddings - self.org_vocab_size + self.num_embeddings_per_partition = divide(self.num_embeddings_padded, + self.tp_size) + assert (self.shard_indices.num_elements_padded == + self.num_embeddings_per_partition) + self.num_org_embeddings_per_partition = ( + self.shard_indices.org_vocab_end_index - + self.shard_indices.org_vocab_start_index) + self.num_added_embeddings_per_partition = ( + self.shard_indices.added_vocab_end_index - + self.shard_indices.added_vocab_start_index) + + self.quant_method.create_weights(self, + self.embedding_dim, + [self.num_embeddings_per_partition], + self.embedding_dim, + self.num_embeddings_padded, + params_dtype=params_dtype, + weight_loader=self.weight_loader) + + @classmethod + def _get_indices(cls, vocab_size_padded: int, org_vocab_size_padded: int, + vocab_size: int, org_vocab_size: int, tp_rank: int, + tp_size: int) -> VocabParallelEmbeddingShardIndices: + """Get start and end indices for vocab parallel embedding, following the + layout outlined in the class docstring, based on the given tp_rank and + tp_size.""" + num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded + padded_org_vocab_start_index, padded_org_vocab_end_index = ( + vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, + tp_size)) + padded_added_vocab_start_index, padded_added_vocab_end_index = ( + vocab_range_from_global_vocab_size(num_added_embeddings_padded, + tp_rank, + tp_size, + offset=org_vocab_size)) + # remove padding + org_vocab_start_index = min(padded_org_vocab_start_index, + org_vocab_size) + org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size) + added_vocab_start_index = min(padded_added_vocab_start_index, + vocab_size) + added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size) + return VocabParallelEmbeddingShardIndices( + padded_org_vocab_start_index, padded_org_vocab_end_index, + padded_added_vocab_start_index, padded_added_vocab_end_index, + org_vocab_start_index, org_vocab_end_index, added_vocab_start_index, + added_vocab_end_index) + + def get_sharded_to_full_mapping(self) -> list[int] | None: + """Get a mapping that can be used to reindex the gathered + logits for sampling. + + During sampling, we gather logits from all ranks. The relationship + of index->token_id will follow the same format as outlined in the class + docstring. However, after the gather, we want to reindex the final + logits tensor to map index->token_id one-to-one (the index is always + equal the token_id it corresponds to). The indices returned by this + method allow us to do that. + """ + if self.tp_size < 2: + return None + + base_embeddings: list[int] = [] + added_embeddings: list[int] = [] + padding: list[int] = [] + for tp_rank in range(self.tp_size): + shard_indices = self._get_indices(self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, tp_rank, + self.tp_size) + range_start = self.num_embeddings_per_partition * tp_rank + range_end = self.num_embeddings_per_partition * (tp_rank + 1) + base_embeddings.extend( + range(range_start, + range_start + shard_indices.num_org_elements)) + padding.extend( + range(range_start + shard_indices.num_org_elements, + range_start + shard_indices.num_org_elements_padded)) + added_embeddings.extend( + range( + range_start + shard_indices.num_org_elements_padded, + range_start + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements)) + padding.extend( + range( + range_start + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements, + range_start + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded)) + assert (range_start + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded == range_end) + ret = base_embeddings + added_embeddings + padding + assert len(ret) == self.num_embeddings_padded + return ret + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + output_dim = getattr(param, "output_dim", None) + packed_dim = getattr(param, "packed_dim", None) + + # If the parameter is a gguf weight, then load it directly. + if getattr(param, "is_gguf_weight_type", None): + param.data.copy_(loaded_weight) + param.weight_type = loaded_weight.item() + return + elif isinstance(param, UninitializedParameter): + shape = list(loaded_weight.shape) + if output_dim is not None: + shape[output_dim] = self.num_embeddings_per_partition + param.materialize(tuple(shape), dtype=loaded_weight.dtype) + + # If parameter does not have output dim, then it should + # be copied onto all gpus (e.g. g_idx for act_order gptq). + if output_dim is None: + assert param.data.shape == loaded_weight.shape + param.data.copy_(loaded_weight) + return + + # Shard indexes for loading the weight + start_idx = self.shard_indices.org_vocab_start_index + shard_size = self.shard_indices.org_vocab_end_index - start_idx + + # If param packed on the same dim we are sharding on, then + # need to adjust offsets of loaded weight by pack_factor. + if packed_dim is not None and packed_dim == output_dim: + packed_factor = param.packed_factor if isinstance( + param, BasevLLMParameter) else param.pack_factor + assert loaded_weight.shape[output_dim] == (self.org_vocab_size // + param.packed_factor) + start_idx = start_idx // packed_factor + shard_size = shard_size // packed_factor + else: + assert loaded_weight.shape[output_dim] == self.org_vocab_size + + # Copy the data. Select chunk corresponding to current shard. + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) + + param[:loaded_weight.shape[0]].data.copy_(loaded_weight) + param[loaded_weight.shape[0]:].data.fill_(0) + + def forward(self, input_): + if self.tp_size > 1: + # Build the mask. + masked_input, input_mask = get_masked_input_and_mask( + input_, self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + self.shard_indices.num_org_vocab_padding, + self.shard_indices.added_vocab_start_index, + self.shard_indices.added_vocab_end_index) + else: + masked_input = input_ + # Get the embeddings. + output_parallel = self.quant_method.embedding(self, masked_input.long()) + # Mask the output embedding. + if self.tp_size > 1: + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + # Reduce across all the model parallel GPUs. + output = tensor_model_parallel_all_reduce(output_parallel) + return output + + def extra_repr(self) -> str: + s = f"num_embeddings={self.num_embeddings_per_partition}" + s += f", embedding_dim={self.embedding_dim}" + s += f", org_vocab_size={self.org_vocab_size}" + s += f', num_embeddings_padded={self.num_embeddings_padded}' + s += f', tp_size={self.tp_size}' + return s diff --git a/trainer/logger.py b/trainer/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..949dbf70532a6af6b68ca3ece7804819538b5eba --- /dev/null +++ b/trainer/logger.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py +"""Logging configuration for trainer.""" +import datetime +import json +import logging +import os +import sys +from functools import lru_cache, partial +from logging import Logger +from logging.config import dictConfig +from os import path +from types import MethodType +from typing import Any, cast + +import trainer.envs as envs + +TRAINER_CONFIGURE_LOGGING = envs.TRAINER_CONFIGURE_LOGGING +TRAINER_LOGGING_CONFIG_PATH = envs.TRAINER_LOGGING_CONFIG_PATH +TRAINER_LOGGING_LEVEL = envs.TRAINER_LOGGING_LEVEL +TRAINER_LOGGING_PREFIX = envs.TRAINER_LOGGING_PREFIX + +RED = '\033[91m' +GREEN = '\033[92m' +RESET = '\033[0;0m' + +_warned_local_main_process = False +_warned_main_process = False + +_FORMAT = (f"{TRAINER_LOGGING_PREFIX}%(levelname)s %(asctime)s " + "[%(filename)s:%(lineno)d] %(message)s") +_DATE_FORMAT = "%m-%d %H:%M:%S" + +DEFAULT_LOGGING_CONFIG = { + "formatters": { + "trainer": { + "class": "trainer.logging_utils.NewLineFormatter", + "datefmt": _DATE_FORMAT, + "format": _FORMAT, + }, + }, + "handlers": { + "trainer": { + "class": "logging.StreamHandler", + "formatter": "trainer", + "level": TRAINER_LOGGING_LEVEL, + "stream": "ext://sys.stdout", + }, + }, + "loggers": { + "trainer": { + "handlers": ["trainer"], + "level": "DEBUG", + "propagate": False, + }, + }, + "root": { + "handlers": ["trainer"], + "level": "DEBUG", + }, + "version": 1, + "disable_existing_loggers": False +} + + +@lru_cache +def _print_info_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.info(msg, stacklevel=2) + + +@lru_cache +def _print_warning_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.warning(msg, stacklevel=2) + + +# TODO(will): add env variable to control this process-aware logging behavior +def _info(logger: Logger, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any) -> None: + """Process-aware INFO level logging function. + + This function controls logging behavior based on the process rank, allowing for + selective logging from specific processes in a distributed environment. + + Args: + logger: The logger instance to use for logging + msg: The message format string to log + *args: Format string arguments + main_process_only: If True, only log if this is the global main process (RANK=0) + local_main_process_only: If True, only log if this is the local main process (LOCAL_RANK=0) + **kwargs: Additional keyword arguments to pass to the logger.log method + - stacklevel: Defaults to 2 to show the original caller's location + + Note: + - When both main_process_only and local_main_process_only are True, + the message will be logged only if both conditions are met + - When both are False, the message will be logged from all processes + - By default, only logs from processes with LOCAL_RANK=0 + """ + try: + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + except Exception: + local_rank = 0 + rank = 0 + + is_main_process = rank == 0 + is_local_main_process = local_rank == 0 + + if (main_process_only and is_main_process) or (local_main_process_only + and is_local_main_process): + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + global _warned_local_main_process, _warned_main_process + + if not _warned_local_main_process and local_main_process_only: + logger.warning( + '%s is_local_main_process is set to True, logging only from the local main process.%s', + GREEN, + RESET, + ) + _warned_local_main_process = True + if not _warned_main_process and main_process_only: + logger.warning( + '%s is_main_process_only is set to True, logging only from the main process.%s', + GREEN, + RESET, + ) + _warned_main_process = True + + if not main_process_only and not local_main_process_only: + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + +class _TRAINERLogger(Logger): + """ + Note: + This class is just to provide type information. + We actually patch the methods directly on the :class:`logging.Logger` + instance to avoid conflicting with other libraries such as + `intel_extension_for_pytorch.utils._logger`. + """ + + def info_once(self, msg: str) -> None: + """ + As :meth:`info`, but subsequent calls with the same message + are silently dropped. + """ + _print_info_once(self, msg) + + def warning_once(self, msg: str) -> None: + """ + As :meth:`warning`, but subsequent calls with the same message + are silently dropped. + """ + _print_warning_once(self, msg) + + def info( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any) -> None: + _info(self, + msg, + *args, + main_process_only=main_process_only, + local_main_process_only=local_main_process_only, + **kwargs) + + +def _configure_trainer_root_logger() -> None: + logging_config = dict[str, Any]() + + if not TRAINER_CONFIGURE_LOGGING and TRAINER_LOGGING_CONFIG_PATH: + raise RuntimeError( + "TRAINER_CONFIGURE_LOGGING evaluated to false, but " + "TRAINER_LOGGING_CONFIG_PATH was given. TRAINER_LOGGING_CONFIG_PATH " + "implies TRAINER_CONFIGURE_LOGGING. Please enable " + "TRAINER_CONFIGURE_LOGGING or unset TRAINER_LOGGING_CONFIG_PATH." + ) + + if TRAINER_CONFIGURE_LOGGING: + logging_config = DEFAULT_LOGGING_CONFIG + + if TRAINER_LOGGING_CONFIG_PATH: + if not path.exists(TRAINER_LOGGING_CONFIG_PATH): + raise RuntimeError( + "Could not load logging config. File does not exist: %s", + TRAINER_LOGGING_CONFIG_PATH) + with open(TRAINER_LOGGING_CONFIG_PATH, encoding="utf-8") as file: + custom_config = json.loads(file.read()) + + if not isinstance(custom_config, dict): + raise ValueError("Invalid logging config. Expected Dict, got %s.", + type(custom_config).__name__) + logging_config = custom_config + + for formatter in logging_config.get("formatters", {}).values(): + # This provides backwards compatibility after #10134. + if formatter.get("class") == "trainer.logging.NewLineFormatter": + formatter["class"] = "trainer.logging_utils.NewLineFormatter" + + if logging_config: + dictConfig(logging_config) + + +def init_logger(name: str) -> _TRAINERLogger: + """The main purpose of this function is to ensure that loggers are + retrieved in such a way that we can be sure the root trainer logger has + already been configured.""" + + logger = logging.getLogger(name) + + methods_to_patch = { + "info_once": _print_info_once, + "warning_once": _print_warning_once, + "info": _info, + } + + for method_name, method in methods_to_patch.items(): + setattr(logger, method_name, + MethodType(method, logger)) # type: ignore[arg-type] + + return cast(_TRAINERLogger, logger) + + +# The root logger is initialized when the module is imported. +# This is thread-safe as the module is only imported once, +# guaranteed by the Python GIL. +_configure_trainer_root_logger() + +logger = init_logger(__name__) + + +def _trace_calls(log_path, root_dir, frame, event, arg=None): + if event in ['call', 'return']: + # Extract the filename, line number, function name, and the code object + filename = frame.f_code.co_filename + lineno = frame.f_lineno + func_name = frame.f_code.co_name + if not filename.startswith(root_dir): + # only log the functions in the trainer root_dir + return + # Log every function call or return + try: + last_frame = frame.f_back + if last_frame is not None: + last_filename = last_frame.f_code.co_filename + last_lineno = last_frame.f_lineno + last_func_name = last_frame.f_code.co_name + else: + # initial frame + last_filename = "" + last_lineno = 0 + last_func_name = "" + with open(log_path, 'a') as f: + ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") + if event == 'call': + f.write(f"{ts} Call to" + f" {func_name} in {filename}:{lineno}" + f" from {last_func_name} in {last_filename}:" + f"{last_lineno}\n") + else: + f.write(f"{ts} Return from" + f" {func_name} in {filename}:{lineno}" + f" to {last_func_name} in {last_filename}:" + f"{last_lineno}\n") + except NameError: + # modules are deleted during shutdown + pass + return partial(_trace_calls, log_path, root_dir) + + +def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): + """ + Enable tracing of every function call in code under `root_dir`. + This is useful for debugging hangs or crashes. + `log_file_path` is the path to the log file. + `root_dir` is the root directory of the code to trace. If None, it is the + trainer root directory. + + Note that this call is thread-level, any threads calling this function + will have the trace enabled. Other threads will not be affected. + """ + logger.warning( + "TRAINER_TRACE_FUNCTION is enabled. It will record every" + " function executed by Python. This will slow down the code. It " + "is suggested to be used for debugging hang or crashes only.") + logger.info("Trace frame log is saved to %s", log_file_path) + if root_dir is None: + # by default, this is the trainer root directory + root_dir = os.path.dirname(os.path.dirname(__file__)) + sys.settrace(partial(_trace_calls, log_file_path, root_dir)) diff --git a/trainer/logging_utils/__init__.py b/trainer/logging_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dbd5716dc8e765fca1243ca82b4eae7cb9826db1 --- /dev/null +++ b/trainer/logging_utils/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 + +from trainer.logging_utils.formatter import NewLineFormatter +from trainer.logging_utils.formatter import setup_for_distributed + +__all__ = [ + "NewLineFormatter", + "setup_for_distributed", +] diff --git a/trainer/logging_utils/formatter.py b/trainer/logging_utils/formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..2c3e1550f23c40638d136e2f3e0b843761016afb --- /dev/null +++ b/trainer/logging_utils/formatter.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logging_utils/formatter.py + +import logging + + +class NewLineFormatter(logging.Formatter): + """Adds logging prefix to newlines to align multi-line messages.""" + + def __init__(self, fmt, datefmt=None, style="%"): + logging.Formatter.__init__(self, fmt, datefmt, style) + + def format(self, record): + msg = logging.Formatter.format(self, record) + if record.message != "": + parts = msg.split(record.message) + msg = msg.replace("\n", "\r\n" + parts[0]) + return msg + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop('force', False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print diff --git a/trainer/models/__init__.py b/trainer/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/models/encoders/base.py b/trainer/models/encoders/base.py new file mode 100644 index 0000000000000000000000000000000000000000..88dcd124cebc38b57355701c8c0a49d397491e90 --- /dev/null +++ b/trainer/models/encoders/base.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +from abc import ABC, abstractmethod +from dataclasses import field + +import torch +from torch import nn + +from trainer.configs.models.encoders import (BaseEncoderOutput, + ImageEncoderConfig, + TextEncoderConfig) +from trainer.platforms import AttentionBackendEnum + + +class TextEncoder(nn.Module, ABC): + _fsdp_shard_conditions: list = field(default_factory=lambda: []) + _stacked_params_mapping: list[tuple[str, str, + str]] = field(default_factory=list) + _supported_attention_backends: tuple[ + AttentionBackendEnum, + ...] = TextEncoderConfig()._supported_attention_backends + + def __init__(self, config: TextEncoderConfig) -> None: + super().__init__() + self.config = config + self._fsdp_shard_conditions = config._fsdp_shard_conditions + self._stacked_params_mapping = config.arch_config.stacked_params_mapping + if not self.supported_attention_backends: + raise ValueError( + f"Subclass {self.__class__.__name__} must define _supported_attention_backends" + ) + + @abstractmethod + def forward(self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs) -> BaseEncoderOutput: + pass + + @property + def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: + return self._supported_attention_backends + + +class ImageEncoder(nn.Module, ABC): + _supported_attention_backends: tuple[ + AttentionBackendEnum, + ...] = ImageEncoderConfig()._supported_attention_backends + + def __init__(self, config: ImageEncoderConfig) -> None: + super().__init__() + self.config = config + if not self.supported_attention_backends: + raise ValueError( + f"Subclass {self.__class__.__name__} must define _supported_attention_backends" + ) + + @abstractmethod + def forward(self, pixel_values: torch.Tensor, + **kwargs) -> BaseEncoderOutput: + pass + + @property + def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: + return self._supported_attention_backends diff --git a/trainer/models/encoders/bert.py b/trainer/models/encoders/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..33e51f9735e01e812f69f1b340112caec11b5428 --- /dev/null +++ b/trainer/models/encoders/bert.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +import os + +import torch +import torch.nn as nn +from transformers import BertModel, BertTokenizer + + +class HunyuanClip(nn.Module): + """ + Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py + hunyuan's clip used BertModel and BertTokenizer, so we copy it. + """ + + def __init__(self, model_dir, max_length=77): + super().__init__() + + self.max_length = max_length + self.tokenizer = BertTokenizer.from_pretrained( + os.path.join(model_dir, 'tokenizer')) + self.text_encoder = BertModel.from_pretrained( + os.path.join(model_dir, 'clip_text_encoder')) + + @torch.no_grad + def forward(self, prompts, with_mask=True): + self.device = next(self.text_encoder.parameters()).device + text_inputs = self.tokenizer( + prompts, + padding="max_length", + max_length=self.max_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + prompt_embeds = self.text_encoder( + text_inputs.input_ids.to(self.device), + attention_mask=text_inputs.attention_mask.to(self.device) + if with_mask else None, + ) + return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output diff --git a/trainer/models/encoders/clip.py b/trainer/models/encoders/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..6e421c02331435e13da74f6a7c38d69aaadf4209 --- /dev/null +++ b/trainer/models/encoders/clip.py @@ -0,0 +1,633 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/clip.py +# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py +"""Minimal implementation of CLIPVisionModel intended to be only used +within a vision language model.""" +from collections.abc import Iterable + +import torch +import torch.nn as nn + +# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from trainer.attention import LocalAttention +from trainer.configs.models.encoders import (BaseEncoderOutput, + CLIPTextConfig, CLIPVisionConfig) +from trainer.distributed import divide, get_tp_world_size +from trainer.layers.activation import get_act_fn +from trainer.layers.linear import (ColumnParallelLinear, QKVParallelLinear, + RowParallelLinear) +from trainer.layers.quantization import QuantizationConfig +from trainer.logger import init_logger +from trainer.models.encoders.base import ImageEncoder, TextEncoder +from trainer.models.encoders.vision import resolve_visual_encoder_outputs +# TODO: support quantization +# from vllm.model_executor.layers.quantization import QuantizationConfig +from trainer.models.loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + + +# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 # noqa +class CLIPVisionEmbeddings(nn.Module): + + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + assert self.image_size % self.patch_size == 0 + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size)**2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, + self.embed_dim) + self.register_buffer("position_ids", + torch.arange(self.num_positions).expand((1, -1)), + persistent=False) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size = pixel_values.shape[0] + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to( + dtype=target_dtype)) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.position_embedding(self.position_ids) + + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, + embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + if input_ids is not None: + seq_length = input_ids.shape[-1] + elif inputs_embeds is not None: + seq_length = inputs_embeds.shape[-2] + else: + raise ValueError( + "Either input_ids or inputs_embeds must be provided.") + + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +class CLIPAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + "embed_dim must be divisible by num_heads " + f"(got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads}).") + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.embed_dim, + head_size=self.head_dim, + total_num_heads=self.num_heads, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.out_proj = RowParallelLinear( + input_size=self.embed_dim, + output_size=self.embed_dim, + quant_config=quant_config, + prefix=f"{prefix}.out_proj", + ) + + self.tp_size = get_tp_world_size() + self.num_heads_per_partition = divide(self.num_heads, self.tp_size) + + self.attn = LocalAttention( + self.num_heads_per_partition, + self.head_dim, + self.num_heads_per_partition, + softmax_scale=self.scale, + causal=True, + supported_attention_backends=config._supported_attention_backends) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, + self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + ): + """Input shape: Batch x Time x Channel""" + + qkv_states, _ = self.qkv_proj(hidden_states) + query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) + # use flash_attn_func + query_states = query_states.reshape(query_states.shape[0], + query_states.shape[1], + self.num_heads_per_partition, + self.head_dim) + key_states = key_states.reshape(key_states.shape[0], + key_states.shape[1], + self.num_heads_per_partition, + self.head_dim) + value_states = value_states.reshape(value_states.shape[0], + value_states.shape[1], + self.num_heads_per_partition, + self.head_dim) + attn_output = self.attn(query_states, key_states, value_states) + attn_output = attn_output.reshape( + attn_output.shape[0], attn_output.shape[1], + self.num_heads_per_partition * self.head_dim) + attn_output, _ = self.out_proj(attn_output) + + return attn_output, None + + +class CLIPMLP(nn.Module): + + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.activation_fn = get_act_fn(config.hidden_act) + self.fc1 = ColumnParallelLinear(config.hidden_size, + config.intermediate_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc1") + self.fc2 = RowParallelLinear(config.intermediate_size, + config.hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc2") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states, _ = self.fc2(hidden_states) + + return hidden_states + + +class CLIPEncoderLayer(nn.Module): + + def __init__( + self, + config: CLIPTextConfig | CLIPVisionConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = CLIPAttention( + config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.layer_norm1 = nn.LayerNorm(config.hidden_size, + eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config, + quant_config=quant_config, + prefix=f"{prefix}.mlp") + self.layer_norm2 = nn.LayerNorm(config.hidden_size, + eps=config.layer_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn(hidden_states=hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class CLIPEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self + attention layers. Each layer is a [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + + if num_hidden_layers_override is None: + num_hidden_layers = config.num_hidden_layers + else: + num_hidden_layers = num_hidden_layers_override + self.layers = nn.ModuleList([ + CLIPEncoderLayer(config=config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_idx}") + for layer_idx in range(num_hidden_layers) + ]) + + def forward( + self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool + ) -> torch.Tensor | list[torch.Tensor]: + hidden_states_pool = [inputs_embeds] + hidden_states = inputs_embeds + + for encoder_layer in self.layers: + hidden_states = encoder_layer(hidden_states) + if return_all_hidden_states: + hidden_states_pool.append(hidden_states) + # If we have multiple feature sample layers, we return all hidden + # states in order and grab the ones we need by index. + if return_all_hidden_states: + return hidden_states_pool + return [hidden_states] + + +class CLIPTextTransformer(nn.Module): + + def __init__(self, + config: CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + prefix: str = ""): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPTextEmbeddings(config) + + self.encoder = CLIPEncoder( + config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=prefix) + + self.final_layer_norm = nn.LayerNorm(embed_dim, + eps=config.layer_norm_eps) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + ) -> BaseEncoderOutput: + r""" + Returns: + + """ + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else + self.config.output_hidden_states) + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, + position_ids=position_ids) + + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + # causal_attention_mask = _create_4d_causal_attention_mask( + # input_shape, hidden_states.dtype, device=hidden_states.device + # ) + + # # expand attention_mask + # if attention_mask is not None and not self._use_flash_attention_2: + # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") + # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + # attention_mask=attention_mask, + # causal_attention_mask=causal_attention_mask, + # output_attentions=output_attentions, + return_all_hidden_states=output_hidden_states, + # return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[-1] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], + device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device). + argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], + device=last_hidden_state.device), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + (input_ids.to(dtype=torch.int, device=last_hidden_state.device + ) == self.eos_token_id).int().argmax(dim=-1), + ] + + return BaseEncoderOutput( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs, + # attentions=encoder_outputs.attentions, + ) + + +class CLIPTextModel(TextEncoder): + + def __init__( + self, + config: CLIPTextConfig, + ) -> None: + super().__init__(config) + self.text_model = CLIPTextTransformer(config=config, + quant_config=config.quant_config, + prefix=config.prefix) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + + outputs: BaseEncoderOutput = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=output_hidden_states, + ) + return outputs + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + + # Define mapping for stacked parameters + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + # Handle q_proj, k_proj, v_proj -> qkv_proj mapping + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name in name: + # Replace the weight name with the parameter name + model_param_name = name.replace(weight_name, param_name) + + if model_param_name in params_dict: + param = params_dict[model_param_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(model_param_name) + break + else: + # Use default weight loader for all other parameters + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + + +class CLIPVisionTransformer(nn.Module): + + def __init__( + self, + config: CLIPVisionConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + require_post_norm: bool | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + + # NOTE: This typo of "layrnorm" is not fixed on purpose to match + # the original transformers code and name of the model weights. + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.encoder = CLIPEncoder( + config=config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=f"{prefix}.encoder", + ) + + num_hidden_layers = config.num_hidden_layers + if len(self.encoder.layers) > config.num_hidden_layers: + raise ValueError( + f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers.") + + # If possible, skip post_layernorm to conserve memory + if require_post_norm is None: + require_post_norm = len(self.encoder.layers) == num_hidden_layers + + if require_post_norm: + self.post_layernorm = nn.LayerNorm(embed_dim, + eps=config.layer_norm_eps) + else: + self.post_layernorm = None + + def forward( + self, + pixel_values: torch.Tensor, + feature_sample_layers: list[int] | None = None, + ) -> torch.Tensor: + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + return_all_hidden_states = feature_sample_layers is not None + + # Produces either the last layer output or all of the hidden states, + # depending on if we have feature_sample_layers or not + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + return_all_hidden_states=return_all_hidden_states) + + if not return_all_hidden_states: + encoder_outputs = encoder_outputs[0] + + # Handle post-norm (if applicable) and stacks feature layers if needed + encoder_outputs = resolve_visual_encoder_outputs( + encoder_outputs, feature_sample_layers, self.post_layernorm, + self.config.num_hidden_layers) + + return encoder_outputs + + +class CLIPVisionModel(ImageEncoder): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + + def __init__(self, config: CLIPVisionConfig) -> None: + super().__init__(config) + self.vision_model = CLIPVisionTransformer( + config=config, + quant_config=config.quant_config, + num_hidden_layers_override=config.num_hidden_layers_override, + require_post_norm=config.require_post_norm, + prefix=f"{config.prefix}.vision_model") + + def forward( + self, + pixel_values: torch.Tensor, + feature_sample_layers: list[int] | None = None, + **kwargs, + ) -> BaseEncoderOutput: + last_hidden_state = self.vision_model(pixel_values, + feature_sample_layers) + return BaseEncoderOutput(last_hidden_state=last_hidden_state) + + @property + def device(self): + return next(self.parameters()).device + + # (TODO) Add prefix argument for filtering out weights to be loaded + # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986 + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + layer_count = len(self.vision_model.encoder.layers) + + for name, loaded_weight in weights: + if name.startswith("visual_projection"): + continue + # post_layernorm is not needed in CLIPVisionModel + if (name.startswith("vision_model.post_layernorm") + and self.vision_model.post_layernorm is None): + continue + + # omit layers when num_hidden_layers_override is set + if name.startswith("vision_model.encoder.layers"): + layer_idx = int(name.split(".")[3]) + if layer_idx >= layer_count: + continue + + for (param_name, weight_name, + shard_id) in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/trainer/models/encoders/llama.py b/trainer/models/encoders/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..d0dc097666dbbdbe482f7d9d9bd492903b71c0de --- /dev/null +++ b/trainer/models/encoders/llama.py @@ -0,0 +1,431 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/llama.py + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only LLaMA model compatible with HuggingFace weights.""" +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +# from vllm.model_executor.layers.quantization import QuantizationConfig +from trainer.attention import LocalAttention +# from ..utils import (extract_layer_index) +from trainer.configs.models.encoders import BaseEncoderOutput, LlamaConfig +from trainer.distributed import get_tp_world_size +from trainer.layers.activation import SiluAndMul +from trainer.layers.layernorm import RMSNorm +from trainer.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, RowParallelLinear) +from trainer.layers.quantization import QuantizationConfig +from trainer.layers.rotary_embedding import get_rope +from trainer.layers.vocab_parallel_embedding import VocabParallelEmbedding +from trainer.models.encoders.base import TextEncoder +from trainer.models.loader.weight_utils import (default_weight_loader, + maybe_remap_kv_scale_name) + + +class LlamaMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + # output_size=intermediate_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class LlamaAttention(nn.Module): + + def __init__(self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + bias_o_proj: bool = False, + prefix: str = "") -> None: + super().__init__() + # layer_idx = extract_layer_index(prefix) + self.hidden_size = hidden_size + tp_size = get_tp_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr(config, "head_dim", + self.hidden_size // self.total_num_heads) + # Phi models introduced a partial_rotary_factor parameter in the config + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + is_neox_style = True + is_gguf = quant_config and hasattr( + quant_config, "get_name") and quant_config.get_name() == "gguf" + if is_gguf and config.model_type == "llama": + is_neox_style = False + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=int(rope_theta), + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + ) + + self.attn = LocalAttention( + self.num_heads, + self.head_dim, + self.num_kv_heads, + softmax_scale=self.scaling, + causal=True, + supported_attention_backends=config._supported_attention_backends) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + # attn_output = self.attn(q, k, v) + # use flash_attn_func + # TODO (Attn abstraction and backend) + # from flash_attn import flash_attn_func + # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) + batch_size = q.shape[0] + seq_len = q.shape[1] + q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + # import pdb; pdb.set_trace() + # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) + attn_output = self.attn(q, k, v) + attn_output = attn_output.reshape(batch_size, seq_len, + self.num_heads * self.head_dim) + + output, _ = self.o_proj(attn_output) + return output + + +class LlamaDecoderLayer(nn.Module): + + def __init__( + self, + config: LlamaConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings) + max_position_embeddings = getattr(config, "max_position_embeddings", + 8192) + # Support abacusai/Smaug-72B-v0.1 with attention_bias + # Support internlm/internlm-7b with bias + attention_bias = getattr(config, "attention_bias", False) or getattr( + config, "bias", False) + bias_o_proj = attention_bias + # support internlm/internlm3-8b with qkv_bias + if hasattr(config, 'qkv_bias'): + attention_bias = config.qkv_bias + + self.self_attn = LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", + config.num_attention_heads), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + bias_o_proj=bias_o_proj, + prefix=f"{prefix}.self_attn", + ) + self.mlp = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual) + + hidden_states = self.self_attn(positions=positions, + hidden_states=hidden_states) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class LlamaModel(TextEncoder): + + def __init__( + self, + config: LlamaConfig, + ): + super().__init__(config) + + self.config = config + self.quant_config = self.config.quant_config + if config.lora_config is not None: + max_loras = 1 + lora_vocab_size = 1 + if hasattr(config.lora_config, "max_loras"): + max_loras = config.lora_config.max_loras + if hasattr(config.lora_config, "lora_extra_vocab_size"): + lora_vocab_size = config.lora_config.lora_extra_vocab_size + lora_vocab = lora_vocab_size * max_loras + else: + lora_vocab = 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=config.quant_config, + ) + + self.layers = nn.ModuleList([ + LlamaDecoderLayer(config=config, + quant_config=config.quant_config, + prefix=f"{config.prefix}.layers.{i}") + for i in range(config.num_hidden_layers) + ]) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else + self.config.output_hidden_states) + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + + if position_ids is None: + position_ids = torch.arange( + 0, hidden_states.shape[1], + device=hidden_states.device).unsqueeze(0) + + all_hidden_states: tuple[Any, ...] | None = ( + ) if output_hidden_states else None + for layer in self.layers: + if all_hidden_states is not None: + # TODO + all_hidden_states += ( + hidden_states, ) if residual is None else (hidden_states + + residual, ) + hidden_states, residual = layer(position_ids, hidden_states, + residual) + + hidden_states, _ = self.norm(hidden_states, residual) + + # add hidden states from the last decoder layer + if all_hidden_states is not None: + all_hidden_states += (hidden_states, ) + + # TODO(will): maybe unify the output format with other models and use + # our own class + output = BaseEncoderOutput( + last_hidden_state=hidden_states, + # past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + # attentions=all_self_attns, + ) + + return output + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if ("rotary_emb.cos_cached" in name + or "rotary_emb.sin_cached" in name): + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + # if (self.quant_config is not None and + # (scale_name := self.quant_config.get_cache_scale(name))): + # # Loading kv cache quantization scales + # param = params_dict[scale_name] + # weight_loader = getattr(param, "weight_loader", + # default_weight_loader) + # loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else + # loaded_weight[0]) + # weight_loader(param, loaded_weight) + # loaded_params.add(scale_name) + # continue + if "scale" in name: + # Remapping the name of FP8 kv-scale. + kv_scale_name: str | None = maybe_remap_kv_scale_name( + name, params_dict) + if kv_scale_name is None: + continue + else: + name = kv_scale_name + for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/trainer/models/encoders/stepllm.py b/trainer/models/encoders/stepllm.py new file mode 100644 index 0000000000000000000000000000000000000000..76920b665365616299d6d4a6be452e0bcf381007 --- /dev/null +++ b/trainer/models/encoders/stepllm.py @@ -0,0 +1,584 @@ +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import os +from functools import wraps + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.modeling_utils import PretrainedConfig, PreTrainedModel + +from trainer.models.dits.stepvideo import StepVideoRMSNorm + + +class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): + + def __init__(self, device=None): + self.device = device + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if getattr(func, '__module__', None) == 'torch.nn.init': + if 'tensor' in kwargs: + return kwargs['tensor'] + else: + return args[0] + if self.device is not None and func in torch.utils._device._device_constructors( + ) and kwargs.get('device') is None: + kwargs['device'] = self.device + return func(*args, **kwargs) + + +def with_empty_init(func): + + @wraps(func) + def wrapper(*args, **kwargs): + with EmptyInitOnDevice('cpu'): + return func(*args, **kwargs) + + return wrapper + + +class LLaMaEmbedding(nn.Module): + """Language model embeddings. + + Arguments: + hidden_size: hidden size + vocab_size: vocabulary size + max_sequence_length: maximum size of sequence. This + is used for positional embedding + embedding_dropout_prob: dropout probability for embeddings + init_method: weight initialization method + num_tokentypes: size of the token-type embeddings. 0 value + will ignore this embedding + """ + + def __init__( + self, + cfg, + ): + super().__init__() + self.hidden_size = cfg.hidden_size + self.params_dtype = cfg.params_dtype + self.fp32_residual_connection = cfg.fp32_residual_connection + self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 + self.word_embeddings = torch.nn.Embedding( + cfg.padded_vocab_size, + self.hidden_size, + ) + self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) + + def forward(self, input_ids): + # Embeddings. + if self.embedding_weights_in_fp32: + self.word_embeddings = self.word_embeddings.to(torch.float32) + embeddings = self.word_embeddings(input_ids) + if self.embedding_weights_in_fp32: + embeddings = embeddings.to(self.params_dtype) + self.word_embeddings = self.word_embeddings.to(self.params_dtype) + + # Data format change to avoid explicit transposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + + # Dropout. + embeddings = self.embedding_dropout(embeddings) + + return embeddings + + +class StepChatTokenizer: + """Step Chat Tokenizer""" + + def __init__( + self, + model_file, + name="StepChatTokenizer", + bot_token="<|BOT|>", # Begin of Turn + eot_token="<|EOT|>", # End of Turn + call_start_token="<|CALL_START|>", # Call Start + call_end_token="<|CALL_END|>", # Call End + think_start_token="<|THINK_START|>", # Think Start + think_end_token="<|THINK_END|>", # Think End + mask_start_token="<|MASK_1e69f|>", # Mask start + mask_end_token="<|UNMASK_1e69f|>", # Mask end + ): + import sentencepiece + + self._tokenizer = sentencepiece.SentencePieceProcessor( + model_file=model_file) + + self._vocab = {} + self._inv_vocab = {} + + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + for idx in range(self._tokenizer.get_piece_size()): + text = self._tokenizer.id_to_piece(idx) + self._inv_vocab[idx] = text + self._vocab[text] = idx + + if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown( + idx): + self._special_tokens[text] = idx + self._inv_special_tokens[idx] = text + + self._unk_id = self._tokenizer.unk_id() + self._bos_id = self._tokenizer.bos_id() + self._eos_id = self._tokenizer.eos_id() + + for token in [ + bot_token, eot_token, call_start_token, call_end_token, + think_start_token, think_end_token + ]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + assert token in self._special_tokens, f"Token '{token}' is not a special token" + + for token in [mask_start_token, mask_end_token]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + + self._bot_id = self._tokenizer.piece_to_id(bot_token) + self._eot_id = self._tokenizer.piece_to_id(eot_token) + self._call_start_id = self._tokenizer.piece_to_id(call_start_token) + self._call_end_id = self._tokenizer.piece_to_id(call_end_token) + self._think_start_id = self._tokenizer.piece_to_id(think_start_token) + self._think_end_id = self._tokenizer.piece_to_id(think_end_token) + self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) + self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) + + self._underline_id = self._tokenizer.piece_to_id("\u2581") + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def vocab_size(self): + return self._tokenizer.vocab_size() + + def tokenize(self, text: str) -> list[int]: + return self._tokenizer.encode_as_ids(text) + + def detokenize(self, token_ids: list[int]) -> str: + return self._tokenizer.decode_ids(token_ids) + + +class Tokens: + + def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, + max_seq_len) -> None: + self.input_ids = input_ids + self.attention_mask = attention_mask + self.cu_input_ids = cu_input_ids + self.cu_seqlens = cu_seqlens + self.max_seq_len = max_seq_len + + def to(self, device): + self.input_ids = self.input_ids.to(device) + self.attention_mask = self.attention_mask.to(device) + self.cu_input_ids = self.cu_input_ids.to(device) + self.cu_seqlens = self.cu_seqlens.to(device) + return self + + +class Wrapped_StepChatTokenizer(StepChatTokenizer): + + def __call__(self, + text, + max_length=320, + padding="max_length", + truncation=True, + return_tensors="pt"): + # [bos, ..., eos, pad, pad, ..., pad] + self.BOS = 1 + self.EOS = 2 + self.PAD = 2 + out_tokens = [] + attn_mask = [] + if len(text) == 0: + part_tokens = [self.BOS] + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + else: + for part in text: + part_tokens = self.tokenize(part) + part_tokens = part_tokens[:(max_length - + 2)] # leave 2 space for bos and eos + part_tokens = [self.BOS] + part_tokens + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * + (max_length - valid_size)) + + out_tokens = torch.tensor(out_tokens, dtype=torch.long) + attn_mask = torch.tensor(attn_mask, dtype=torch.long) + + # padding y based on tp size + padded_len = 0 + padded_flag = False + if padded_len > 0: + padded_flag = True + if padded_flag: + pad_tokens = torch.tensor([[self.PAD] * max_length], + device=out_tokens.device) + pad_attn_mask = torch.tensor([[1] * padded_len + [0] * + (max_length - padded_len)], + device=attn_mask.device) + out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) + attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) + + # cu_seqlens + cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) + seqlen = attn_mask.sum(dim=1).tolist() + cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), + 0).to(device=out_tokens.device, + dtype=torch.int32) + max_seq_len = max(seqlen) + return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, + max_seq_len) + + +def flash_attn_func(q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=True, + return_attn_probs=False, + tp_group_rank=0, + tp_group_size=1): + softmax_scale = q.size(-1)**( + -0.5) if softmax_scale is None else softmax_scale + return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, + causal, return_attn_probs, None, tp_group_rank, + tp_group_size)[0] + + +class FlashSelfAttention(torch.nn.Module): + + def __init__( + self, + attention_dropout=0.0, + ): + super().__init__() + self.dropout_p = attention_dropout + + def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): + if cu_seqlens is None: + output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) + else: + raise ValueError('cu_seqlens is not supported!') + + return output + + +def safediv(n, d): + q, r = divmod(n, d) + assert r == 0 + return q + + +class MultiQueryAttention(nn.Module): + + def __init__(self, cfg, layer_id=None): + super().__init__() + + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.max_seq_len = cfg.seq_length + self.use_flash_attention = cfg.use_flash_attn + assert self.use_flash_attention, 'FlashAttention is required!' + + self.n_groups = cfg.num_attention_groups + self.tp_size = 1 + self.n_local_heads = cfg.num_attention_heads + self.n_local_groups = self.n_groups + + self.wqkv = nn.Linear( + cfg.hidden_size, + cfg.hidden_size + self.head_dim * 2 * self.n_groups, + bias=False, + ) + self.wo = nn.Linear( + cfg.hidden_size, + cfg.hidden_size, + bias=False, + ) + + # assert self.use_flash_attention, 'non-Flash attention not supported yet.' + self.core_attention = FlashSelfAttention( + attention_dropout=cfg.attention_dropout) + # self.core_attention = LocalAttention( + # num_heads = self.n_local_heads, + # head_size = self.head_dim, + # # num_kv_heads = self.n_local_groups, + # casual = True, + # supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA], # RIVER TODO + # ) + self.layer_id = layer_id + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + seqlen, bsz, dim = x.shape + xqkv = self.wqkv(x) + + xq, xkv = torch.split( + xqkv, + (dim // self.tp_size, + self.head_dim * 2 * self.n_groups // self.tp_size), + dim=-1, + ) + + # gather on 1st dimension + xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) + xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) + xk, xv = xkv.chunk(2, -1) + + # rotary embedding + flash attn + xq = rearrange(xq, "s b h d -> b s h d") + xk = rearrange(xk, "s b h d -> b s h d") + xv = rearrange(xv, "s b h d -> b s h d") + + # q_per_kv = self.n_local_heads // self.n_local_groups + # if q_per_kv > 1: + # b, s, h, d = xk.size() + # if h == 1: + # xk = xk.expand(b, s, q_per_kv, d) + # xv = xv.expand(b, s, q_per_kv, d) + # else: + # ''' To cover the cases where h > 1, we have + # the following implementation, which is equivalent to: + # xk = xk.repeat_interleave(q_per_kv, dim=-2) + # xv = xv.repeat_interleave(q_per_kv, dim=-2) + # but can avoid calling aten::item() that involves cpu. + # ''' + # idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() + # xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + # xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + if self.use_flash_attention: + output = self.core_attention(xq, xk, xv) + # reduce-scatter only support first dimension now + output = rearrange(output, "b s h d -> s b (h d)").contiguous() + else: + xq, xk, xv = [ + rearrange(x, "b s ... -> s b ...").contiguous() + for x in (xq, xk, xv) + ] + output = self.core_attention(xq, xk, xv) #, mask) + output = self.wo(output) + return output + + +class FeedForward(nn.Module): + + def __init__( + self, + cfg, + dim: int, + hidden_dim: int, + layer_id: int, + multiple_of: int = 256, + ): + super().__init__() + + hidden_dim = multiple_of * ( + (hidden_dim + multiple_of - 1) // multiple_of) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.swiglu = swiglu + + self.w1 = nn.Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=False, + ) + + def forward(self, x): + x = self.swiglu(self.w1(x)) + output = self.w2(x) + return output + + +class TransformerBlock(nn.Module): + + def __init__(self, cfg, layer_id: int): + super().__init__() + + self.n_heads = cfg.num_attention_heads + self.dim = cfg.hidden_size + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.attention = MultiQueryAttention( + cfg, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + cfg, + dim=cfg.hidden_size, + hidden_dim=cfg.ffn_hidden_size, + layer_id=layer_id, + ) + self.layer_id = layer_id + self.attention_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + self.ffn_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + residual = self.attention.forward(self.attention_norm(x), mask, + cu_seqlens, max_seq_len) + h = x + residual + ffn_res = self.feed_forward.forward(self.ffn_norm(h)) + out = h + ffn_res + return out + + +class Transformer(nn.Module): + + def __init__( + self, + config, + max_seq_size=8192, + ): + super().__init__() + self.num_layers = config.num_layers + self.layers = self._build_layers(config) + + def _build_layers(self, config): + layers = torch.nn.ModuleList() + for layer_id in range(self.num_layers): + layers.append(TransformerBlock( + config, + layer_id=layer_id + 1, + )) + return layers + + def forward( + self, + hidden_states, + attention_mask, + cu_seqlens=None, + max_seq_len=None, + ): + + if max_seq_len is not None and not isinstance(max_seq_len, + torch.Tensor): + max_seq_len = torch.tensor(max_seq_len, + dtype=torch.int32, + device="cpu") + + for lid, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + attention_mask, + cu_seqlens, + max_seq_len, + ) + return hidden_states + + +class Step1Model(PreTrainedModel): + config_class = PretrainedConfig + + @with_empty_init + def __init__( + self, + config, + ): + super().__init__(config) + self.tok_embeddings = LLaMaEmbedding(config) + self.transformer = Transformer(config) + + def forward( + self, + input_ids=None, + attention_mask=None, + ): + + hidden_states = self.tok_embeddings(input_ids) + + hidden_states = self.transformer( + hidden_states, + attention_mask, + ) + return hidden_states + + +class STEP1TextEncoder(torch.nn.Module): + + def __init__(self, model_dir, max_length=320): + super().__init__() + self.max_length = max_length + self.text_tokenizer = Wrapped_StepChatTokenizer( + os.path.join(model_dir, 'step1_chat_tokenizer.model')) + text_encoder = Step1Model.from_pretrained(model_dir) + self.text_encoder = text_encoder.eval().to(torch.bfloat16) + + @torch.no_grad + def forward(self, prompts, with_mask=True, max_length=None): + self.device = next(self.text_encoder.parameters()).device + + with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16): + if type(prompts) is str: + prompts = [prompts] + txt_tokens = self.text_tokenizer(prompts, + max_length=max_length + or self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + y = self.text_encoder(txt_tokens.input_ids.to(self.device), + attention_mask=txt_tokens.attention_mask.to( + self.device) if with_mask else None) + y_mask = txt_tokens.attention_mask + return y.transpose(0, 1), y_mask diff --git a/trainer/models/encoders/t5.py b/trainer/models/encoders/t5.py new file mode 100644 index 0000000000000000000000000000000000000000..42e13897e0e62fb01e57425935e4d23a5f9ef2c4 --- /dev/null +++ b/trainer/models/encoders/t5.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/t5/modeling_t5.py + +# Derived from T5 implementation posted on HuggingFace; license below: +# +# coding=utf-8 +# Copyright 2018 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch T5 & UMT5 model.""" + +import math +from collections.abc import Iterable +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn + +from trainer.configs.models.encoders import BaseEncoderOutput, T5Config +from trainer.distributed import get_tp_rank, get_tp_world_size +from trainer.layers.activation import get_act_fn +from trainer.layers.layernorm import RMSNorm +from trainer.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, RowParallelLinear) +from trainer.layers.quantization import QuantizationConfig +from trainer.layers.vocab_parallel_embedding import VocabParallelEmbedding +from trainer.models.encoders.base import TextEncoder +from trainer.models.loader.weight_utils import default_weight_loader +from trainer.platforms import current_platform + + +class AttentionType: + """ + Attention type. + Use string to be compatible with `torch.compile`. + """ + # Decoder attention between previous layer Q/K/V + DECODER = "decoder" + # Encoder attention between previous layer Q/K/V for encoder-decoder + ENCODER = "encoder" + # Encoder attention between previous layer Q/K/V + ENCODER_ONLY = "encoder_only" + # Attention between dec. Q and enc. K/V for encoder-decoder + ENCODER_DECODER = "encoder_decoder" + + +@dataclass +class AttentionMetadata: + attn_bias: torch.Tensor + + +class T5DenseActDense(nn.Module): + + def __init__(self, + config: T5Config, + quant_config: QuantizationConfig | None = None): + super().__init__() + self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False) + self.wo = RowParallelLinear(config.d_ff, + config.d_model, + bias=False, + quant_config=quant_config) + self.act = get_act_fn(config.dense_act_fn) + + def forward(self, hidden_states) -> torch.Tensor: + hidden_states, _ = self.wi(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.wo(hidden_states) + return hidden_states + + +class T5DenseGatedActDense(nn.Module): + + def __init__(self, + config: T5Config, + quant_config: QuantizationConfig | None = None): + super().__init__() + self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False, + quant_config=quant_config) + self.wi_1 = MergedColumnParallelLinear(config.d_model, [config.d_ff], + bias=False, + quant_config=quant_config) + # Should not run in fp16 unless mixed-precision is used, + # see https://github.com/huggingface/transformers/issues/20287. + self.wo = RowParallelLinear(config.d_ff, + config.d_model, + bias=False, + quant_config=quant_config) + self.act = get_act_fn(config.dense_act_fn) + + def forward(self, hidden_states) -> torch.Tensor: + hidden_gelu = self.act(self.wi_0(hidden_states)[0]) + hidden_linear, _ = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states, _ = self.wo(hidden_states) + return hidden_states + + +class T5LayerFF(nn.Module): + + def __init__(self, + config: T5Config, + quant_config: QuantizationConfig | None = None): + super().__init__() + if config.is_gated_act: + self.DenseReluDense = T5DenseGatedActDense( + config, quant_config=quant_config) + else: + self.DenseReluDense = T5DenseActDense(config, + quant_config=quant_config) + + self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward(self, hidden_states) -> torch.Tensor: + forwarded_states = self.layer_norm(hidden_states) + forwarded_states = self.DenseReluDense(forwarded_states) + hidden_states = hidden_states + forwarded_states + return hidden_states + + +# T5 has attn_bias and does not use softmax scaling +class T5MultiHeadAttention(nn.Module): + + def __init__(self) -> None: + super().__init__() + + def forward(self, q, k, v, attn_bias=None): + b, _, n, c = q.shape + attn = torch.einsum('binc,bjnc->bnij', q, k) + if attn_bias is not None: + attn += attn_bias + + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + x = x.reshape(b, -1, n * c) + return x + + +class T5Attention(nn.Module): + + def __init__(self, + config: T5Config, + attn_type: str, + has_relative_attention_bias=False, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + self.attn_type = attn_type + # Cross-attention has no relative pos encoding anyway + self.is_decoder = attn_type == AttentionType.DECODER + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = \ + config.relative_attention_num_buckets + self.relative_attention_max_distance = \ + config.relative_attention_max_distance + self.d_model = config.d_model + self.key_value_proj_dim = config.d_kv + self.total_num_heads = self.total_num_kv_heads = config.num_heads + + # Partition heads across multiple tensor parallel GPUs. + tp_world_size = get_tp_world_size() + assert config.num_heads % tp_world_size == 0 + self.n_heads = config.num_heads // tp_world_size + + self.inner_dim = self.n_heads * self.key_value_proj_dim + # No GQA in t5. + # self.n_kv_heads = self.n_heads + + self.qkv_proj = QKVParallelLinear( + self.d_model, + self.d_model // self.total_num_heads, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.attn = T5MultiHeadAttention() + + if self.has_relative_attention_bias: + self.relative_attention_bias = \ + VocabParallelEmbedding(self.relative_attention_num_buckets, + self.total_num_heads, + org_num_embeddings=self.relative_attention_num_buckets, + padding_size=self.relative_attention_num_buckets, + quant_config=quant_config) + self.o = RowParallelLinear( + self.d_model, + self.d_model, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + @staticmethod + def _relative_position_bucket(relative_position, + bidirectional=True, + num_buckets=32, + max_distance=128) -> torch.Tensor: + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + Translate relative position to a bucket number for relative attention. + The relative position is defined as memory_position - query_position, + i.e. the distance in tokens from the attending position to the + attended-to position. If bidirectional=False, then positive relative + positions are invalid. We use smaller buckets for small absolute + relative_position and larger buckets for larger absolute + relative_positions. All relative positions >=max_distance map to the + same bucket. All relative positions <=-max_distance map to the same + bucket. This should allow for more graceful generalization to longer + sequences than the model has been trained on + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + Returns: + a Tensor with the same shape as relative_position, containing int32 + values in the range [0, num_buckets) + """# noqa: E501 + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to( + torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min(relative_position, + torch.zeros_like(relative_position)) + # now relative_position is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + # The other half of the buckets are for logarithmically bigger bins + # in positions up to max_distance + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) / + math.log(max_distance / max_exact) * + (num_buckets - max_exact)).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, + torch.full_like(relative_position_if_large, num_buckets - 1)) + + relative_buckets += torch.where(is_small, relative_position, + relative_position_if_large) + return relative_buckets + + def compute_bias(self, + query_length, + key_length, + device=None) -> torch.Tensor: + """Compute binned relative position bias""" + if device is None: + device = self.relative_attention_bias.weight.device + context_position = torch.arange(query_length, + dtype=torch.long, + device=device)[:, None] + memory_position = torch.arange(key_length, + dtype=torch.long, + device=device)[None, :] + # max_seq_len, nh + relative_position = memory_position - context_position + relative_position_bucket = self._relative_position_bucket( + relative_position, # shape (query_length, key_length) + bidirectional=(not self.is_decoder), + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias( + relative_position_bucket + ) # shape (query_length, key_length, num_heads) + x = values.permute([2, 0, 1]).unsqueeze( + 0) # shape (1, num_heads, query_length, key_length) + return x + + def forward( + self, + hidden_states: torch.Tensor, # (num_tokens, d_model) + attention_mask: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + bs, seq_len, _ = hidden_states.shape + num_seqs = bs + n, c = self.n_heads, self.d_model // self.total_num_heads + qkv, _ = self.qkv_proj(hidden_states) + # Projection of 'own' hidden state (self-attention). No GQA here. + q, k, v = qkv.split(self.inner_dim, dim=-1) + q = q.reshape(bs, seq_len, n, c) + k = k.reshape(bs, seq_len, n, c) + v = v.reshape(bs, seq_len, n, c) + + assert attn_metadata is not None + attn_bias = attn_metadata.attn_bias + # Not compatible with CP here (as all encoder-decoder models), + # as it assumes homogeneous batch (prefills or decodes). + if self.has_relative_attention_bias: + # Self-attention. Compute T5 relative positional encoding. + # The bias term is computed on longest sequence in batch. Biases + # for shorter sequences are slices of the longest. + assert self.attn_type == AttentionType.ENCODER + attn_bias = self.compute_bias(seq_len, + seq_len).repeat(num_seqs, 1, 1, 1) + attn_metadata.attn_bias = attn_bias + else: + # Encoder/Decoder Self-Attention Layer, attn bias already cached. + assert attn_bias is not None + + if attention_mask is not None: + attention_mask = attention_mask.view( + bs, 1, 1, + -1) if attention_mask.ndim == 2 else attention_mask.unsqueeze(1) + mask_val = -1e4 if current_platform.is_mps() else torch.finfo( + q.dtype).min + attn_bias.masked_fill_(attention_mask == 0, mask_val) + + if get_tp_world_size() > 1: + rank = get_tp_rank() + attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) * + self.n_heads, :, :] + attn_output = self.attn(q, k, v, attn_bias) + output, _ = self.o(attn_output) + return output + + +class T5LayerSelfAttention(nn.Module): + + def __init__( + self, + config, + has_relative_attention_bias=False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.SelfAttention = T5Attention( + config, + AttentionType.DECODER + if "decoder" in prefix else AttentionType.ENCODER, + has_relative_attention_bias=has_relative_attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.SelfAttention") + self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.SelfAttention( + hidden_states=normed_hidden_states, + attention_mask=attention_mask, + attn_metadata=attn_metadata, + ) + hidden_states = hidden_states + attention_output + return hidden_states + + +class T5LayerCrossAttention(nn.Module): + + def __init__(self, + config, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + self.EncDecAttention = T5Attention(config, + AttentionType.ENCODER_DECODER, + has_relative_attention_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.EncDecAttention") + self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.EncDecAttention( + hidden_states=normed_hidden_states, + attn_metadata=attn_metadata, + ) + hidden_states = hidden_states + attention_output + return hidden_states + + +class T5Block(nn.Module): + + def __init__(self, + config: T5Config, + is_decoder: bool, + has_relative_attention_bias=False, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): + super().__init__() + self.is_decoder = is_decoder + self.layer = nn.ModuleList() + self.layer.append( + T5LayerSelfAttention( + config, + has_relative_attention_bias=has_relative_attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.self_attn")) + + if self.is_decoder: + self.layer.append( + T5LayerCrossAttention(config, + quant_config=quant_config, + prefix=f"{prefix}.cross_attn")) + + self.layer.append(T5LayerFF(config, quant_config=quant_config)) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + + hidden_states = self.layer[0](hidden_states=hidden_states, + attention_mask=attention_mask, + attn_metadata=attn_metadata) + if self.is_decoder: + hidden_states = self.layer[1](hidden_states=hidden_states, + attn_metadata=attn_metadata) + + # Apply Feed Forward layer + hidden_states = self.layer[2](hidden_states) + else: + hidden_states = self.layer[1](hidden_states) + return hidden_states + + +class T5Stack(nn.Module): + + def __init__(self, + config: T5Config, + is_decoder: bool, + n_layers: int, + embed_tokens=None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_umt5: bool = False): + super().__init__() + self.embed_tokens = embed_tokens + self.is_umt5 = is_umt5 + if is_umt5: + self.block = nn.ModuleList([ + T5Block(config, + is_decoder=is_decoder, + has_relative_attention_bias=True, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{i}") for i in range(n_layers) + ]) + else: + # Only the first block has relative positional encoding. + self.block = nn.ModuleList([ + T5Block(config, + is_decoder=is_decoder, + has_relative_attention_bias=i == 0, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{i}") for i in range(n_layers) + ]) + self.final_layer_norm = RMSNorm(config.d_model, + eps=config.layer_norm_epsilon) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + + for idx, block in enumerate(self.block): + hidden_states = block( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_metadata=attn_metadata, + ) + hidden_states = self.final_layer_norm(hidden_states) + return hidden_states + + +class T5EncoderModel(TextEncoder): + + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__(config) + + quant_config = None + + self.shared = VocabParallelEmbedding( + config.vocab_size, + config.d_model, + org_num_embeddings=config.vocab_size) + + self.encoder = T5Stack(config, + False, + config.num_layers, + self.shared, + quant_config=quant_config, + prefix=f"{prefix}.encoder", + is_umt5=False) + + def get_input_embeddings(self): + return self.shared + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + attn_metadata = AttentionMetadata(None) + hidden_states = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + attn_metadata=attn_metadata, + ) + + return BaseEncoderOutput(last_hidden_state=hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q", "q"), + (".qkv_proj", ".k", "k"), + (".qkv_proj", ".v", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + loaded = False + if "decoder" in name or "lm_head" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded = True + break + if not loaded: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class UMT5EncoderModel(TextEncoder): + + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__(config) + + quant_config = None + + self.shared = VocabParallelEmbedding( + config.vocab_size, + config.d_model, + org_num_embeddings=config.vocab_size) + + self.encoder = T5Stack(config, + False, + config.num_layers, + self.shared, + quant_config=quant_config, + prefix=f"{prefix}.encoder", + is_umt5=True) + + def get_input_embeddings(self): + return self.shared + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + attn_metadata = AttentionMetadata(None) + hidden_states = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + attn_metadata=attn_metadata, + ) + + return BaseEncoderOutput( + last_hidden_state=hidden_states, + attention_mask=attention_mask, + ) + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + loaded = False + if "decoder" in name or "lm_head" in name: + continue + for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded = True + break + if not loaded: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/trainer/models/encoders/vision.py b/trainer/models/encoders/vision.py new file mode 100644 index 0000000000000000000000000000000000000000..3533093931984063049229c40f1f8d6e4697fe46 --- /dev/null +++ b/trainer/models/encoders/vision.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py + +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +import torch +from transformers import PretrainedConfig + +from trainer.logger import init_logger + +logger = init_logger(__name__) + +_C = TypeVar("_C", bound=PretrainedConfig) + + +class VisionEncoderInfo(ABC, Generic[_C]): + + def __init__(self, vision_config: _C) -> None: + super().__init__() + + self.vision_config = vision_config + + @abstractmethod + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + raise NotImplementedError + + @abstractmethod + def get_max_image_tokens(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_image_size(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_patch_size(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_patch_grid_length(self) -> int: + raise NotImplementedError + + +def resolve_visual_encoder_outputs( + encoder_outputs: torch.Tensor | list[torch.Tensor], + feature_sample_layers: list[int] | None, + post_layer_norm: torch.nn.LayerNorm | None, + max_possible_layers: int, +) -> torch.Tensor: + """Given the outputs a visual encoder module that may correspond to the + output of the last layer, or a list of hidden states to be stacked, + handle post normalization and resolve it into a single output tensor. + + Args: + encoder_outputs: Output of encoder's last layer or all hidden states. + feature_sample_layers: Optional layer indices to grab from the encoder + outputs; if provided, encoder outputs must be a list. + post_layer_norm: Post norm to apply to the output of the encoder. + max_possible_layers: Total layers in the fully loaded visual encoder. + + """ + if feature_sample_layers is None: + if post_layer_norm is not None: + return post_layer_norm(encoder_outputs) + return encoder_outputs + + # Get the hidden states corresponding to the layer indices. + # Negative values are relative to the full visual encoder, + # so offset them depending on how many layers were loaded. + # NOTE: this assumes that encoder_outputs is a list containing + # the inputs to the visual encoder, followed by the hidden states + # of each layer. + num_loaded_layers = len(encoder_outputs) - 1 + offset = max_possible_layers - num_loaded_layers + hs_pool = [ + encoder_outputs[layer_idx] + if layer_idx >= 0 else encoder_outputs[layer_idx + offset] + for layer_idx in feature_sample_layers + ] + + # Apply post-norm on the final hidden state if we are using it + uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1) + if post_layer_norm is not None and uses_last_layer: + hs_pool[-1] = post_layer_norm(encoder_outputs) + return torch.cat(hs_pool, dim=-1) \ No newline at end of file diff --git a/trainer/models/hf_transformer_utils.py b/trainer/models/hf_transformer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ccd98976e2869a47d39a8ef98ce04e26de9fa909 --- /dev/null +++ b/trainer/models/hf_transformer_utils.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py + +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for Huggingface Transformers.""" + +import contextlib +import json +import os +from pathlib import Path +from typing import Any + +from huggingface_hub import snapshot_download +from transformers import AutoConfig, PretrainedConfig +from transformers.models.auto.modeling_auto import ( + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES) + +_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { + # ChatGLMConfig.model_type: ChatGLMConfig, + # DbrxConfig.model_type: DbrxConfig, + # ExaoneConfig.model_type: ExaoneConfig, + # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig, +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf(model_path: str): + if os.path.exists(model_path): + return model_path + + return snapshot_download(model_path, + allow_patterns=["*.json", "*.bin", "*.model"]) + + +def get_hf_config( + model: str, + trust_remote_code: bool, + revision: str | None = None, + model_override_args: dict | None = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + raise NotImplementedError("GGUF models are not supported.") + + config = AutoConfig.from_pretrained(model, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs) + if config.model_type in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[config.model_type] + config = config_class.from_pretrained(model, revision=revision) + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + config._name_or_path = model + if model_override_args: + config.update(model_override_args) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError( + f"Can't get gguf config for {config.model_type}.") + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +def get_diffusers_config(model: str, ) -> dict[str, Any]: + """Gets a configuration for the given diffusers model. + + Args: + model: The model name or path. + trainer_args: Optional inference arguments to override in the config. + + Returns: + The loaded configuration. + """ + config_name = "config.json" + if "scheduler" in model: + config_name = "scheduler_config.json" + # Check if the model path exists + if os.path.exists(model): + config_file = os.path.join(model, config_name) + if os.path.exists(config_file): + try: + # Load the config directly from the file + with open(config_file) as f: + config_dict: dict[str, Any] = json.load(f) + if "_diffusers_version" in config_dict: + config_dict.pop("_diffusers_version") + # TODO(will): apply any overrides from inference args + return config_dict + except Exception as e: + raise RuntimeError( + f"Failed to load diffusers config from {config_file}: {e}" + ) from e + raise RuntimeError(f"Config file not found at {config_file}") + else: + raise RuntimeError(f"Diffusers config file not found at {model}") + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def attach_additional_stop_token_ids(tokenizer): + # Special handling for stop token <|eom_id|> generated by llama 3 tool use. + if "<|eom_id|>" in tokenizer.get_added_vocab(): + tokenizer.additional_stop_token_ids = set( + [tokenizer.get_added_vocab()["<|eom_id|>"]]) + else: + tokenizer.additional_stop_token_ids = None + + +def check_gguf_file(model: str | os.PathLike) -> bool: + """Check if the file is a GGUF model.""" + model = Path(model) + if not model.is_file(): + return False + elif model.suffix == ".gguf": + return True + + with open(model, "rb") as f: + header = f.read(4) + return header == b"GGUF" diff --git a/trainer/models/hyvideo/__init__.py b/trainer/models/hyvideo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7be4f2d5d0cf63c0c3c64669e58ca5d2f7b8d205 --- /dev/null +++ b/trainer/models/hyvideo/__init__.py @@ -0,0 +1,44 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import socket +from .commons import get_gpu_memory + + +if 'TOKENIZERS_PARALLELISM' not in os.environ: + os.environ['TOKENIZERS_PARALLELISM'] = 'false' + +def find_free_port(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(('localhost', 0)) + port = sock.getsockname()[1] + sock.close() + return port + +def __initialize_default_distributed_environment(): + if 'RANK' not in os.environ: + os.environ['RANK'] = '0' + if 'WORLD_SIZE' not in os.environ: + os.environ['WORLD_SIZE'] = '1' + if 'LOCAL_RANK' not in os.environ: + os.environ['LOCAL_RANK'] = '0' + if 'MASTER_ADDR' not in os.environ: + os.environ['MASTER_ADDR'] = 'localhost' + if 'MASTER_PORT' not in os.environ: + os.environ['MASTER_PORT'] = str(find_free_port()) + +__initialize_default_distributed_environment() \ No newline at end of file diff --git a/trainer/models/hyvideo/commons/__init__.py b/trainer/models/hyvideo/commons/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4bd3c5a917d3f348518ca3976ea2a8ea16a8ea4 --- /dev/null +++ b/trainer/models/hyvideo/commons/__init__.py @@ -0,0 +1,242 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import torch +from itertools import repeat +from contextlib import contextmanager +from torch import nn +import collections.abc + +def _ntuple(n): + """Create a function that converts input to n-tuple.""" + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + return parse + +# Convenience functions for common tuple sizes +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + +PRECISION_TO_TYPE = { + 'fp32': torch.float32, + 'fp16': torch.float16, + 'bf16': torch.bfloat16, +} + +# Default generation pipeline configurations +PIPELINE_CONFIGS = { + '480p_t2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '480p_i2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '720p_t2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, + '720p_i2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '480p_t2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '480p_i2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '720p_t2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, + '720p_i2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '720p_t2v_distilled_sparse': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '720p_i2v_distilled_sparse': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, +} + +# Default super-resolution pipeline configurations +SR_PIPELINE_CONFIGS = { + '720p_sr_distilled': { + 'flow_shift': 2.0, + 'base_resolution': '480p', + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'num_inference_steps': 6, + }, + '1080p_sr_distilled': { + 'flow_shift': 2.0, + 'base_resolution': '720p', + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'num_inference_steps': 8, + }, +} + +TRANSFORMER_VERSION_TO_SR_VERSION = { + '480p_t2v': '720p_sr_distilled', + '720p_t2v': '1080p_sr_distilled', + '480p_i2v': '720p_sr_distilled', + '720p_i2v': '1080p_sr_distilled', + '480p_t2v_distilled': '720p_sr_distilled', + '720p_t2v_distilled': '1080p_sr_distilled', + '480p_i2v_distilled': '720p_sr_distilled', + '720p_i2v_distilled': '1080p_sr_distilled', + '480p_t2v_distilled_sparse': '720p_sr_distilled', + '720p_t2v_distilled_sparse': '1080p_sr_distilled', + '480p_i2v_distilled_sparse': '720p_sr_distilled', + '720p_i2v_distilled_sparse': '1080p_sr_distilled', +} + +def is_flash2_available(): + try: + from flash_attn import flash_attn_varlen_qkvpacked_func + return True + except Exception: + return False + +def is_flash3_available(): + try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3 # noqa: F401 + return True + except Exception: + return False + +def is_flash_available(): + return is_flash2_available() or is_flash3_available() + +def is_sparse_attn_supported(): + return 'nvidia h' in torch.cuda.get_device_properties(0).name.lower() + +def is_sparse_attn_available(): + if not is_sparse_attn_supported(): + return False + try: + from flex_block_attn import flex_block_attn_func # noqa: F401 + return True + except Exception: + return False + +def maybe_fallback_attn_mode(attn_mode, infer_state=None, block_idx=None): + """ + Determine the final attention mode based on configuration and availability. + + Args: + attn_mode: Requested attention mode + infer_state: Inference configuration object (optional) + block_idx: Current block index (optional) + + Returns: + Final attention mode to use + """ + import warnings + + # Check for sageattn and flex-block-attn conflict + enable_sageattn = False + if infer_state is not None: + enable_sageattn = (infer_state.enable_sageattn and + block_idx in infer_state.sage_blocks_range) + + assert not (enable_sageattn and attn_mode == 'flex-block-attn'), \ + ("SageAttention cannot be used with flex-block-attn mode. " + "Please disable enable_sageattn or use a different attention mode.") + + # Use SageAttention if configured + if enable_sageattn: + attn_mode = 'sageattn' + return attn_mode + + # Handle flash attention modes + if attn_mode == 'flash': + if is_flash3_available(): + attn_mode = 'flash3' + elif is_flash2_available(): + attn_mode = 'flash2' + else: + warnings.warn("flash is not available. Falling back to torch attention.") + attn_mode = 'torch' + elif attn_mode == 'flash3': + if not is_flash3_available(): + warnings.warn("flash3 is not available. Falling back to torch attention.") + attn_mode = 'torch' + elif attn_mode == 'flash2': + if not is_flash2_available(): + warnings.warn("flash2 is not available. Falling back to torch attention.") + attn_mode = 'torch' + if attn_mode in ('flex-block-attn'): + from trainer.models.hyvideo.commons import is_sparse_attn_available + if not is_sparse_attn_available(): + raise ValueError(f"{attn_mode} is not available for your GPU or flex-block-attn is not properly installed.") + return attn_mode + +@contextmanager +def auto_offload_model(models, device, enabled=True): + from diffusers.hooks.group_offloading import _is_group_offload_enabled + if enabled: + if isinstance(models, nn.Module): + models = [models] + for model in models: + if model is not None: + model.to(device) + yield + if enabled: + for model in models: + if model is not None: + model.to(torch.device('cpu')) + +def get_gpu_memory(device=None): + if not torch.cuda.is_available(): + return 0 + device = device if device is not None else torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + if hasattr(torch.cuda, 'get_per_process_memory_fraction'): + memory_fraction = torch.cuda.get_per_process_memory_fraction() + else: + memory_fraction = 1.0 + return props.total_memory * memory_fraction + +def get_rank(): + return int(os.environ.get('RANK', '0')) diff --git a/trainer/models/hyvideo/commons/infer_state.py b/trainer/models/hyvideo/commons/infer_state.py new file mode 100644 index 0000000000000000000000000000000000000000..a991994be05f46493ea8e9f2dd15c9edbf92c739 --- /dev/null +++ b/trainer/models/hyvideo/commons/infer_state.py @@ -0,0 +1,48 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional +from dataclasses import dataclass + +@dataclass +class InferState: + enable_sageattn: bool = False # whether to use SageAttention + sage_blocks_range: Optional[range] = None # block range to use SageAttention + enable_torch_compile: bool = False # whether to use torch compile + +__infer_state = None + +def parse_range(value): + if '-' in value: + start, end = map(int, value.split('-')) + return list(range(start, end + 1)) + else: + return [int(x) for x in value.split(',')] + +def initialize_infer_state(args): + global __infer_state + sage_blocks_range = parse_range(args.sage_blocks_range) + # Map CLI argument use_sageattn to internal enable_sageattn field + use_sageattn = getattr(args, 'use_sageattn', False) + __infer_state = InferState( + enable_sageattn = use_sageattn, + sage_blocks_range = sage_blocks_range, + enable_torch_compile = args.enable_torch_compile, + ) + return __infer_state + +def get_infer_state(): + return __infer_state diff --git a/trainer/models/hyvideo/commons/parallel_states.py b/trainer/models/hyvideo/commons/parallel_states.py new file mode 100644 index 0000000000000000000000000000000000000000..190c5f00520d9332906ff7d4852c145ed7195e10 --- /dev/null +++ b/trainer/models/hyvideo/commons/parallel_states.py @@ -0,0 +1,82 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from dataclasses import dataclass + +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh + +@dataclass +class ParallelDims: + sp: int = 1 + world_size: int = -1 + + def __post_init__(self): + if self.world_size == -1: + if dist.is_initialized(): + self.world_size = dist.get_world_size() + else: + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + self.build_mesh("cuda") + + def build_mesh(self, device_type): + assert self.world_size % self.sp == 0, "world_size must be divisible by sp" + mesh = init_device_mesh( + device_type, + [self.world_size // self.sp, self.sp], + mesh_dim_names=["dp", "sp"] + ) + self.world_mesh = mesh + return mesh + + @property + def sp_enabled(self): + return self.sp > 1 + + @property + def sp_group(self): + return self.world_mesh['sp'].get_group() + + @property + def sp_mesh(self): + return self.world_mesh['sp'] + + @property + def sp_rank(self): + if self.sp_enabled: + return self.world_mesh['sp'].get_local_rank() + else: + return dist.get_rank() + + @property + def dp_enabled(self): + return self.sp > 1 + +__parallel_dims = None + +def initialize_parallel_state( + sp: int = 1, +): + global __parallel_dims + __parallel_dims = ParallelDims(sp=sp) + return __parallel_dims + +def get_parallel_state(): + if __parallel_dims is None: + # create default parallel states (without enabling any parallelism) + initialize_parallel_state() + return __parallel_dims diff --git a/trainer/models/hyvideo/models/__init__.py b/trainer/models/hyvideo/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe721640f8f0188339c70082ecac3ff9b74dcfb --- /dev/null +++ b/trainer/models/hyvideo/models/__init__.py @@ -0,0 +1,16 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + diff --git a/trainer/models/hyvideo/models/autoencoders/__init__.py b/trainer/models/hyvideo/models/autoencoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9108b6ce5a80c2ef04aa8bd269cffdb83830a65 --- /dev/null +++ b/trainer/models/hyvideo/models/autoencoders/__init__.py @@ -0,0 +1,17 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + + diff --git a/trainer/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py b/trainer/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..e97becb39ae75995f982e37e6824619ae0364dc1 --- /dev/null +++ b/trainer/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py @@ -0,0 +1,1149 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import loguru +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, nn +from torch.nn import Conv3d +from torch import distributed as dist + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.autoencoders.vae import BaseOutput, DiagonalGaussianDistribution +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin + +from trainer.models.hyvideo.commons.parallel_states import get_parallel_state + +CACHE_T = 2 + +@dataclass +class DecoderOutput(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +def swish(x: Tensor, inplace=False) -> Tensor: + """Applies the swish activation function (SiLU) with optional inplace support.""" + return F.silu(x, inplace=inplace) + + +def forward_with_checkpointing(module, *inputs, use_checkpointing=False): + """ + Executes the forward pass of a module, optionally using gradient checkpointing. + + Gradient checkpointing trades compute for memory by not storing intermediate + activations during the forward pass, but re-computing them during the backward pass. + + Args: + module (nn.Module): The module to run. + *inputs: The inputs to the module. + use_checkpointing (bool): If True, wraps execution in torch.utils.checkpoint. + """ + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + + if use_checkpointing: + # If checkpointing is enabled, wrap the custom forward function. + # use_reentrant=False is generally recommended for newer PyTorch versions + # to avoid issues with backward hooks and distributed training. + return torch.utils.checkpoint.checkpoint(create_custom_forward(module), *inputs, use_reentrant=False) + else: + # Standard forward pass without checkpointing overhead + return module(*inputs) + +# Optimized implementation of CogVideoXSafeConv3d +# https://github.com/huggingface/diffusers/blob/c9ff360966327ace3faad3807dc871a4e5447501/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py#L38 +class PatchCausalConv3d(nn.Conv3d): + r"""Causal Conv3d with efficient patch processing for large tensors.""" + def find_split_indices(self, seq_len, part_num): + ideal_interval = seq_len / part_num + possible_indices = list(range(0, seq_len, self.stride[0])) + selected_indices = [] + + for i in range(1, part_num): + closest = min(possible_indices, key=lambda x: abs(x - round(i * ideal_interval))) + if closest not in selected_indices: + selected_indices.append(closest) + + merged_indices = [] + prev_idx = 0 + for idx in selected_indices: + if idx - prev_idx >= self.kernel_size[0]: + merged_indices.append(idx) + prev_idx = idx + + return merged_indices + + def forward(self, input): + T = input.shape[2] # input: NCTHW + memory_count = torch.prod(torch.tensor(input.shape)).item() * 2 / 1024**3 + + if T > self.kernel_size[0] and memory_count > 0.6: + kernel_size = self.kernel_size[0] + part_num = int(memory_count / 2) + 1 + split_indices = self.find_split_indices(T, part_num) + input_chunks = torch.tensor_split(input, split_indices, dim=2) if len(split_indices) > 0 else [input] + if kernel_size > 1: + input_chunks = [input_chunks[0]] + [ + torch.cat((input_chunks[i - 1][:, :, -kernel_size + 1:], input_chunks[i]), dim=2) + for i in range(1, len(input_chunks)) + ] + output_chunks = [] + for input_chunk in input_chunks: + output_chunks.append(super().forward(input_chunk)) + output = torch.cat(output_chunks, dim=2) + return output + else: + return super().forward(input) + + +class RMS_norm(nn.Module): + """Root Mean Square Layer Normalization for Channel-First or Last""" + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else -1) + ) * self.scale * self.gamma + self.bias + + + +class CausalConv3d(nn.Module): + """Causal Conv3d with configurable padding for temporal axis.""" + + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode='replicate', + disable_causal=False, + enable_patch_conv=False, + **kwargs + ): + super().__init__() + + self.pad_mode = pad_mode + if disable_causal: + padding = (kernel_size // 2, kernel_size // 2, + kernel_size // 2, kernel_size // 2, + kernel_size // 2, kernel_size // 2) + else: + padding = (kernel_size // 2, kernel_size // 2, + kernel_size // 2, kernel_size // 2, + kernel_size - 1, 0) + self.time_causal_padding = padding + + if enable_patch_conv: + self.conv = PatchCausalConv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + else: + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x, cache_x=None): + padding = list(self.time_causal_padding) + if cache_x is not None and self.time_causal_padding[4] > 0: # not the first frame + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) # concat on temporal dimension (concat one) + padding[4] -= cache_x.shape[2] # from [1,1,1,1,2,0] to [1,1,1,1,1,0] (padding one) + x = F.pad(x, padding, mode=self.pad_mode) + return self.conv(x) + + + +def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None): + """Prepare a causal attention mask for 3D videos. + + Args: + n_frame (int): Number of frames (temporal length). + n_hw (int): Product of height and width. + dtype: Desired mask dtype. + device: Device for the mask. + batch_size (int, optional): If set, expands for batch. + + Returns: + torch.Tensor: Causal attention mask. + """ + seq_len = n_frame * n_hw + mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device) + for i in range(seq_len): + i_frame = i // n_hw + mask[i, : (i_frame + 1) * n_hw] = 0 + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class AttnBlock(nn.Module): + """Self-attention block for 3D video tensors.""" + + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = RMS_norm(in_channels, images=False) + + self.q = Conv3d(in_channels, in_channels, kernel_size=1) + self.k = Conv3d(in_channels, in_channels, kernel_size=1) + self.v = Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, f, h, w = q.shape + q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous() + k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous() + v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous() + attention_mask = prepare_causal_attention_mask(f, h * w, h_.dtype, h_.device, batch_size=b) + h_ = nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask.unsqueeze(1)) + + return rearrange(h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + """ResNet-style block for 3D video tensors.""" + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = RMS_norm(in_channels, images=False) + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = RMS_norm(out_channels, images=False) + self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3) + if self.in_channels != self.out_channels: + self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = x + h = self.norm1(h) + h = swish(h, inplace=True) + + # apply the feature cacheing mechanism + if feat_cache is not None: + # Retrieve the current layer index + idx = feat_idx[0] + + # Clone the last CACHE_T frames from the current input to store for the next step. + cache_x = h[:, :, -CACHE_T:, :, :].clone() + + # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames) + # and we have a previous cache, prepend the last frame of the previous cache. + # This ensures sufficient temporal context for the convolution kernel. + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # actually this means chunk 1 after chunk 0 + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + # Apply the convolution layer using the input 'h' and the *previous* cached state. + h = self.conv1(h, feat_cache[idx]) + + # Update the cache for this layer with the newly prepared context ('cache_x'). + feat_cache[idx] = cache_x + + # Increment the global layer index. + feat_idx[0] += 1 + else: + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv2(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + """Hierarchical downsampling with temporal/ spatial support.""" + + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + self.conv = CausalConv3d(in_channels, out_channels // factor, kernel_size=3) + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward(self, x, feat_cache=None, feat_idx=[0]): + r1 = 2 if self.add_temporal_downsample else 1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv(x) # change the channel, ready for saptial or temporal downsample + + if self.add_temporal_downsample: + if x.shape[2] == 1: + h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + h = torch.cat([h, h], dim=1) + # compute the shortcut part + shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2) + else: + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + # compute the shortcut part + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + else: + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + + return h + shortcut + +class Upsample(nn.Module): + """Hierarchical upsampling with temporal/ spatial support.""" + + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = CausalConv3d(in_channels, out_channels * factor, kernel_size=3) + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + r1 = 2 if self.add_temporal_upsample else 1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv(x) # change the channel, ready for saptial or temporal downsample + + if self.add_temporal_upsample: + if first_chunk: + h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + h = h[:, : h.shape[1] // 2] + # compute the shortcut part + shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1) + elif feat_cache is None and x.shape[2] > 1: + # Multi-frame input without cache: first frame only spatial upsample, rest frames do spatio-temporal upsample + # Separate first frame and remaining frames + h_first = h[:, :, :1, :, :] # first frame + h_rest = h[:, :, 1:, :, :] # remaining frames + x_first = x[:, :, :1, :, :] + x_rest = x[:, :, 1:, :, :] + + # First frame: only spatial upsample + h_first = rearrange(h_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + h_first = h_first[:, : h_first.shape[1] // 2] + shortcut_first = rearrange(x_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + shortcut_first = shortcut_first.repeat_interleave(repeats=self.repeats // 2, dim=1) + out_first = h_first + shortcut_first + + # Remaining frames: spatio-temporal upsample + h_rest = rearrange(h_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut_rest = rearrange(x_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut_rest = shortcut_rest.repeat_interleave(repeats=self.repeats, dim=1) + out_rest = h_rest + shortcut_rest + + return torch.cat([out_first, out_rest], dim=2) + else: + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + # compute the shortcut part + shortcut = rearrange(x, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = shortcut.repeat_interleave(repeats=self.repeats, dim=1) + else: + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + + return h + shortcut + + +class Encoder(nn.Module): + """Hierarchical video encoder with temporal and spatial factorization.""" + + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + downsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[-1] % (2 * z_channels) == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3) + + self.down = nn.ModuleList() + block_in = block_out_channels[0] + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + + add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_downsample = ( + add_spatial_downsample and + bool(i_level >= np.log2(ffactor_spatial // ffactor_temporal)) + ) + if add_spatial_downsample or add_temporal_downsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if downsample_match_channel else block_in + down.downsample = Downsample(block_in, block_out, add_temporal_downsample) + block_in = block_out + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, 2 * z_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + """Forward pass through the encoder.""" + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() # cache x before layers + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv_in(x, feat_cache[idx]) # [B, 160, T, H, W] with padding + feat_cache[idx] = cache_x + feat_idx[0] += 1 # move to next cache + else: + h = self.conv_in(x) + + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks): + h = forward_with_checkpointing( + self.down[i_level].block[i_block], + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing + ) + if hasattr(self.down[i_level], "downsample"): + h = forward_with_checkpointing( + self.down[i_level].downsample, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing + ) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, feat_cache, feat_idx, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, feat_cache, feat_idx, use_checkpointing=use_checkpointing) + + # end + group_size = self.block_out_channels[-1] // (2 * self.z_channels) + shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) + h = self.norm_out(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + h = self.conv_out(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv_out(h) + h += shortcut + return h + + +class Decoder(nn.Module): + """Hierarchical video decoder with upsampling factories.""" + + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + upsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[0] % z_channels == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + block_in = block_out_channels[0] + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal)) + if add_spatial_upsample or add_temporal_upsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if upsample_match_channel else block_in + up.upsample = Upsample(block_in, block_out, add_temporal_upsample) + block_in = block_out + self.up.append(up) + + # end + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, z, feat_cache=None, feat_idx=[0], first_chunk=False): + """Forward pass through the decoder.""" + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + repeats = self.block_out_channels[0] // (self.z_channels) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = z[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv_in(z, feat_cache[idx]) + z.repeat_interleave(repeats=repeats, dim=1) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + # z to block_in + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # middle + h = forward_with_checkpointing( + self.mid.block_1, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing + ) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing( + self.mid.block_2, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing + ) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing(self.up[i_level].block[i_block], + h, feat_cache, feat_idx, use_checkpointing=use_checkpointing) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, feat_cache, + feat_idx, first_chunk, use_checkpointing=use_checkpointing) + + # end + h = self.norm_out(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + h = self.conv_out(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv_out(h) + return h + + +class AutoencoderKLConv3D(ModelMixin, ConfigMixin): + """KL regularized 3D Conv VAE with advanced tiling and slicing strategies.""" + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: Tuple[int, ...], + layers_per_block: int, + ffactor_spatial: int, + ffactor_temporal: int, + sample_size: int, + sample_tsize: int, + scaling_factor: float = None, + shift_factor: Optional[float] = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + ): + super().__init__() + self.ffactor_spatial = ffactor_spatial + self.ffactor_temporal = ffactor_temporal + self.scaling_factor = scaling_factor + self.shift_factor = shift_factor + + self.encoder = Encoder( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + downsample_match_channel=downsample_match_channel, + ) + self.decoder = Decoder( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + upsample_match_channel=upsample_match_channel, + ) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // ffactor_spatial + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // ffactor_temporal + self.tile_overlap_factor = 0.25 + + self._tile_parallelism_enabled = False + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": sum(isinstance(m, CausalConv3d) for m in self.decoder.modules()) + if self.decoder is not None + else 0, + "encoder": sum(isinstance(m, CausalConv3d) for m in self.encoder.modules()) + if self.encoder is not None + else 0, + } + + def set_tile_sample_min_size(self, sample_size: int, tile_overlap_factor: float = 0.2): + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // self.ffactor_spatial + self.tile_overlap_factor = tile_overlap_factor + + assert (self.tile_latent_min_size * self.tile_overlap_factor).is_integer(), \ + "self.tile_latent_min_size multiplied by tile_overlap_factor must be an integer" + + + def _set_gradient_checkpointing(self, module, value=False): + """Enable or disable gradient checkpointing on encoder and decoder.""" + if isinstance(module, (Encoder, Decoder)): + module.gradient_checkpointing = value + + + def enable_temporal_tiling(self, use_tiling: bool = True): + raise RuntimeError('Temporal tiling is not supported for this VAE.') + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + self.enable_spatial_tiling(use_tiling) + + def disable_tiling(self): + self.disable_spatial_tiling() + + def enable_slicing(self): + self.use_slicing = True + + def disable_slicing(self): + self.use_slicing = False + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b horizontally into a at blend_extent region.""" + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = ( + a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + + b[:, :, :, :, x] * (x / blend_extent) + ) + return b + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b vertically into a at blend_extent region.""" + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = ( + a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + + b[:, :, :, y, :] * (y / blend_extent) + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b temporally into a at blend_extent region.""" + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = ( + a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + + b[:, :, x, :, :] * (x / blend_extent) + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor): + """Tiled spatial encoding for large inputs via overlapping.""" + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = x[:, :, :, i: i + self.tile_sample_min_size, j: j + self.tile_sample_min_size] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + return moments + + def temporal_tiled_encode(self, x: torch.Tensor): + """Tiled temporal encoding for large video sequences.""" + raise RuntimeError('Temporal tiling is not supported for this VAE.') + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_latent_min_tsize - blend_extent + + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i: i + self.tile_sample_min_tsize + 1, :, :] + if ( + self.use_spatial_tiling and + (tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size) + ): + tile = self.spatial_tiled_encode(tile) + else: + tile = self.encoder(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + moments = torch.cat(result_row, dim=-3) + return moments + + def enable_tile_parallelism(self): + self._tile_parallelism_enabled = True + + def disable_tile_parallelism(self): + self._tile_parallelism_enabled = False + + def tile_parallel_spatial_tiled_decode(self, z: torch.Tensor): + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + + rank = get_parallel_state().sp_rank + world_size = get_parallel_state().sp + + num_rows = math.ceil(H / overlap_size) + num_cols = math.ceil(W / overlap_size) + total_tiles = num_rows * num_cols + tiles_per_rank = math.ceil(total_tiles / world_size) + + + my_linear_indices = list(range(rank, total_tiles, world_size)) + decoded_tiles = [] + decoded_metas = [] + H_out_std = self.tile_sample_min_size + W_out_std = self.tile_sample_min_size + for lin_idx in my_linear_indices: + ri = lin_idx // num_cols + rj = lin_idx % num_cols + i = ri * overlap_size + j = rj * overlap_size + tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size] + dec = self.decoder(tile) + + pad_h = max(0, H_out_std - dec.shape[-2]) + pad_w = max(0, W_out_std - dec.shape[-1]) + if pad_h > 0 or pad_w > 0: + dec = F.pad(dec, (0, pad_w, 0, pad_h, 0, 0), "constant", 0) + decoded_tiles.append(dec) + decoded_metas.append(torch.tensor([ri, rj, pad_w, pad_h], device=z.device, dtype=torch.int64)) + + while len(decoded_tiles) < tiles_per_rank: + T_out = decoded_tiles[0].shape[2] if len(decoded_tiles) > 0 else (T-1)*self.ffactor_temporal+1 + zero_tile = torch.zeros( + [1, 3, T_out, self.tile_sample_min_size, self.tile_sample_min_size], + device=dec.device, + dtype=dec.dtype + ) + decoded_tiles.append(zero_tile) + meta_tensor = torch.tensor( + [-1, -1, self.tile_sample_min_size, self.tile_sample_min_size], + device=z.device, + dtype=torch.int64 + ) + decoded_metas.append(meta_tensor) + + decoded_tiles = torch.stack(decoded_tiles, dim=0) + decoded_metas = torch.stack(decoded_metas, dim=0) + + tiles_gather_list = [torch.empty_like(decoded_tiles) for _ in range(world_size)] + metas_gather_list = [torch.empty_like(decoded_metas) for _ in range(world_size)] + + dist.all_gather(tiles_gather_list, decoded_tiles, group=get_parallel_state().sp_group) + dist.all_gather(metas_gather_list, decoded_metas, group=get_parallel_state().sp_group) + dist.barrier() + + if rank != 0: + return torch.empty(0, device=z.device) + + rows = [[None for _ in range(num_cols)] for _ in range(num_rows)] + for r in range(world_size): + gathered_tiles_r = tiles_gather_list[r] # [tiles_per_rank, B, C, T, H, W] + gathered_metas_r = metas_gather_list[r] # [tiles_per_rank, 4] + for k in range(gathered_tiles_r.shape[0]): + ri = int(gathered_metas_r[k][0]) + rj = int(gathered_metas_r[k][1]) + if ri < 0 or rj < 0: + continue + if ri < num_rows and rj < num_cols: + # remove padding + pad_w = int(gathered_metas_r[k][2]) + pad_h = int(gathered_metas_r[k][3]) + h_end = None if pad_h == 0 else -pad_h + w_end = None if pad_w == 0 else -pad_w + rows[ri][rj] = gathered_tiles_r[k][:, :, :, :h_end, :w_end] + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if tile is None: + continue + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + return dec + + def spatial_tiled_decode(self, z: torch.Tensor): + if self._tile_parallelism_enabled: + return self.tile_parallel_spatial_tiled_decode(z) + + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = z[:, :, :, i: i + self.tile_latent_min_size, j: j + self.tile_latent_min_size] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + return dec + + def temporal_tiled_decode(self, z: torch.Tensor): + """Tiled temporal decoding for long sequence latents.""" + raise RuntimeError('Temporal tiling is not supported for this VAE.') + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - blend_extent + assert 0 < overlap_size < self.tile_latent_min_tsize + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i: i + self.tile_latent_min_tsize + 1, :, :] + if ( + self.use_spatial_tiling and + (tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size) + ): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self.decoder(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + dec = torch.cat(result_row, dim=-3) + return dec + + def clear_cache(self): + # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + def encode(self, x: Tensor, return_dict: bool = True): + + def _encode(x): + '''slice encoding video latents''' + _, _, num_frame, height, width = x.shape + + if self.use_temporal_tiling and x.shape[-3] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x) + if ( + self.use_spatial_tiling and + (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size) + ): + return self.spatial_tiled_encode(x) + + self.clear_cache() + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) # [B, C, 4, H, W] after patchify + out = torch.cat([out, out_], 2) + + self.clear_cache() + return out + + assert len(x.shape) == 5 # (B, C, T, H, W) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [_encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = _encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: Tensor, return_dict: bool = True, generator=None): + + def _decode(z): + '''slice decoding video latents''' + _, _, num_frame, height, width = z.shape + + if self.use_temporal_tiling and z.shape[-3] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z) + if ( + self.use_spatial_tiling and + (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size) + ): + return self.spatial_tiled_decode(z) + + self.clear_cache() + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + z[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True + ) + else: + out_ = self.decoder(z[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + + self.clear_cache() + + return out + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [_decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = _decode(z) + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_posterior: bool = True, + return_dict: bool = True + ): + """ + Performs the full forward pass of the Autoencoder. + + It encodes the input into a latent distribution, samples (or takes the mode) from it, + and then decodes it back to the image space. + + Args: + sample (torch.Tensor): Input image tensor (B, C, H, W). + sample_posterior (bool): If True, samples from the latent distribution (stochastic). + If False, uses the mode (deterministic). + return_posterior (bool): Whether to include the posterior distribution in the output. + return_dict (bool): If True, returns a structured DecoderOutput object; + otherwise, returns a standard tuple. + """ + # 1. Encode: Map input 'sample' to a latent distribution (Gaussian). + posterior = self.encode(sample).latent_dist + + # 2. Sample: Reparameterization trick. + # Use random sampling for training (variational) or mode for deterministic inference. + z = posterior.sample() if sample_posterior else posterior.mode() + + # 3. Decode: Reconstruct the image from the latent vector 'z'. + dec = self.decode(z).sample + + # 4. Return: Format the output based on 'return_dict' preference. + if return_dict: + return DecoderOutput(sample=dec, posterior=posterior) + else: + return (dec, posterior) \ No newline at end of file diff --git a/trainer/models/hyvideo/models/text_encoders/__init__.py b/trainer/models/hyvideo/models/text_encoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6b78c68ec33e21343905627f2dc18afc4ebeb818 --- /dev/null +++ b/trainer/models/hyvideo/models/text_encoders/__init__.py @@ -0,0 +1,504 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from copy import deepcopy +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from transformers import ( + AutoTokenizer, + AutoModel, +) +from transformers.utils import ModelOutput + + +def use_default(value, default): + """Utility: return value if not None, else default.""" + return value if value is not None else default + +# Prompt templates for different models and tasks + + +__all__ = [ + "C_SCALE", "PROMPT_TEMPLATE", + "MODEL_BASE", +] + +# =================== Constant Values ===================== +# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid +# overflow error when tensorboard logging values. +C_SCALE = 1_000_000_000_000_000 + +PROMPT_TEMPLATE_ENCODE_IMAGE_JSON = [ + {"role": "system", "content": "You are a helpful assistant. Describe the image by detailing the following aspects: \ + 1. The main content and theme of the image. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. The background environment, light, style and atmosphere."}, + {"role": "user", "content": "{}"} +] + +PROMPT_TEMPLATE_ENCODE_VIDEO_JSON = [ + {"role": "system", "content": "You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video."}, + {"role": "user", "content": "{}"} +] + +PROMPT_TEMPLATE = { + "li-dit-encode-image-json": {"template": PROMPT_TEMPLATE_ENCODE_IMAGE_JSON, "crop_start": -1}, # auto-calculate crop_start + "li-dit-encode-video-json": {"template": PROMPT_TEMPLATE_ENCODE_VIDEO_JSON, "crop_start": -1}, # auto-calculate crop_start +} + + +MODEL_BASE = os.getenv("MODEL_BASE", "") +TEXT_ENCODER_PATH = {} +TOKENIZER_PATH = {} + +PRECISION_TO_TYPE = { + 'fp32': torch.float32, + 'fp16': torch.float16, + 'bf16': torch.bfloat16, +} + + +def load_text_encoder( + text_encoder_type, + text_encoder_precision=None, + text_encoder_path=None, + logger=None, + device=None, +): + if text_encoder_path is None: + if text_encoder_type not in TEXT_ENCODER_PATH: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] + + text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True) + + if hasattr(text_encoder, 'language_model'): + text_encoder = text_encoder.language_model + text_encoder.final_layer_norm = text_encoder.norm + + # from_pretrained will ensure that the model is in eval mode. + if text_encoder_precision is not None: + text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision]) + + text_encoder.requires_grad_(False) + + if device is not None: + text_encoder = text_encoder.to(device) + + return text_encoder, text_encoder_path + + +def load_tokenizer( + tokenizer_type, tokenizer_path=None, padding_side="right", logger=None +): + processor = None + if tokenizer_path is None: + if tokenizer_type not in TOKENIZER_PATH: + raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") + tokenizer_path = TOKENIZER_PATH[tokenizer_type] + + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, padding_side=padding_side + ) + + return tokenizer, tokenizer_path, processor + + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + image_features: Optional[list] = None + +class TextEncoder(nn.Module): + def __init__( + self, + text_encoder_type: str, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + tokenizer_type: Optional[str] = None, + tokenizer_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + prompt_template: Optional[dict] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + logger=None, + device=None, + ): + super().__init__() + self.text_encoder_type = text_encoder_type + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.tokenizer_type = ( + tokenizer_type if tokenizer_type is not None else text_encoder_type + ) + self.tokenizer_path = ( + tokenizer_path if tokenizer_path is not None else text_encoder_path + ) + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert ( + use_attention_mask is True + ), "Attention mask is True required when training videos." + self.prompt_template = prompt_template + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + self.logger = logger + + self.use_template = self.prompt_template is not None + if self.use_template: + assert ( + isinstance(self.prompt_template, dict) + and "template" in self.prompt_template + ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" + assert "{}" in str(self.prompt_template["template"]), ( + "`prompt_template['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template['template']}" + ) + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert ( + isinstance(self.prompt_template_video, dict) + and "template" in self.prompt_template_video + ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" + assert "{}" in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}" + ) + + if text_encoder_type != "llm": + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + self.output_key = output_key or "last_hidden_state" + + self.model, self.model_path = load_text_encoder( + text_encoder_type=self.text_encoder_type, + text_encoder_precision=self.precision, + text_encoder_path=self.model_path, + logger=self.logger, + device=device, + ) + + self.tokenizer, self.tokenizer_path, self.processor = load_tokenizer( + tokenizer_type=self.tokenizer_type, + tokenizer_path=self.tokenizer_path, + padding_side="right", + logger=self.logger, + ) + + # pre-calculate crop_start for image and video + if self.use_template and self.prompt_template is not None: + self.text2tokens("a photo of a cat", data_type="image") + if self.use_video_template and self.prompt_template_video is not None: + self.text2tokens("a photo of a cat", data_type="video") + + @property + def dtype(self): + return self.model.dtype + + @property + def device(self): + return self.model.device + + def __repr__(self): + return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + + @staticmethod + def apply_text_to_template(text, template, prevent_empty_text=True): + """ + Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If Ture, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + elif isinstance(template, list): + # For JSON list template format (chat conversation) + # Create a deep copy to avoid modifying the original template + template_copy = deepcopy(template) + for item in template_copy: + if isinstance(item, dict) and "content" in item: + # Replace placeholder with text in the content field + item["content"] = item["content"].format(text if text else (" " if prevent_empty_text else "")) + return template_copy + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def calculate_crop_start(self, tokenized_input): + """ + Automatically calculate the crop_start position based on identifying user tokens. + + Args: + tokenized_input: The output from the tokenizer containing input_ids + + Returns: + int: The position where the actual prompt content begins (after user markers) + """ + input_ids = tokenized_input["input_ids"][0].tolist() # Get the first example's tokens + + marker = "<|im_start|>user\n" + + # Tokenize just the marker to get its token IDs + marker_tokens = self.tokenizer(marker, add_special_tokens=False)["input_ids"] + + # Find the end position of the marker in the input sequence + for i in range(len(input_ids) - len(marker_tokens) + 1): + if input_ids[i:i+len(marker_tokens)] == marker_tokens: + # Return the position after the marker + return i + len(marker_tokens) + + # If marker not found, try to find based on special tokens + if hasattr(self.tokenizer, 'special_tokens_map'): + # Check for user token or any other special token that might indicate user input start + for token_name, token_value in self.tokenizer.special_tokens_map.items(): + if 'user' in token_name.lower(): + user_token_id = self.tokenizer.convert_tokens_to_ids(token_value) + if user_token_id in input_ids: + return input_ids.index(user_token_id) + 1 + + # Default fallback: return 0 (no cropping) + return 0 + + def text2tokens(self, text, data_type="image", max_length=300): + """ + Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = "str" + if self.use_template or self.use_video_template: + if data_type == "image": + prompt_template = self.prompt_template["template"] + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + prompt_template = self.prompt_template_video["template"] + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [ + self.apply_text_to_template(one_text, prompt_template) + for one_text in text + ] + if isinstance(text[0], list): + tokenize_input_type = "list" + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = "list" + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + # First pass: tokenize with arbitrary max_length to find crop_start + if crop_start == -1: + # Use temporary max_length for the first pass (large enough) + temp_kwargs = dict( + truncation=True, + max_length=256, # Temporary large value + padding="max_length", + return_tensors="pt", + ) + + # First tokenization pass to calculate crop_start + if tokenize_input_type == "str": + temp_tokenized = self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **temp_kwargs, + ) + elif tokenize_input_type == "list": + temp_tokenized = self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **temp_kwargs, + ) + + # Calculate the crop_start from this first pass + crop_start = self.calculate_crop_start(temp_tokenized) + + # Store the calculated crop_start for future use + if data_type == "image": + self.prompt_template["crop_start"] = crop_start + else: + self.prompt_template_video["crop_start"] = crop_start + else: + crop_start = 0 + + # Second pass: tokenize with the proper max_length using the found crop_start + kwargs = dict( + truncation=True, + max_length=max_length + (crop_start if crop_start > 0 else 0), + padding="max_length", + return_tensors="pt", + ) + + if tokenize_input_type == "str": + tokenized_output = self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **kwargs, + ) + elif tokenize_input_type == "list": + tokenized_output = self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **kwargs, + ) + else: + raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") + + return tokenized_output + + def encode( + self, + batch_encoding, + use_attention_mask=None, + output_hidden_states=False, + do_sample=None, + hidden_state_skip_layer=None, + return_texts=False, + data_type="image", + device=None, + is_uncond=False + ): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + device = self.model.device if device is None else device + use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) + hidden_state_skip_layer = use_default( + hidden_state_skip_layer, self.hidden_state_skip_layer + ) + do_sample = use_default(do_sample, not self.reproduce) + + attention_mask = ( + batch_encoding["attention_mask"].to(device) if use_attention_mask else None + ) + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + attention_mask=attention_mask, + output_hidden_states=output_hidden_states + or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm(last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", 0) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", 0) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = ( + attention_mask[:, crop_start:] if use_attention_mask else None + ) + + if output_hidden_states: + return TextEncoderModelOutput( + last_hidden_state, attention_mask, outputs.hidden_states + ) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text, max_length=self.max_length) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/trainer/models/hyvideo/models/text_encoders/byT5/__init__.py b/trainer/models/hyvideo/models/text_encoders/byT5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5937758acdd57a8155c4cc98ddacda6d50aeec0a --- /dev/null +++ b/trainer/models/hyvideo/models/text_encoders/byT5/__init__.py @@ -0,0 +1,229 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json +import torch +import torch.nn as nn +from transformers import AutoTokenizer, T5ForConditionalGeneration + + +def load_glyph_byT5_v2(args, device): + """ + Loads ByT5 tokenizer and encoder model for glyph encoding. + + Args: + args (dict): Configuration dictionary containing paths and settings. + device (str or torch.device): Device to load the model onto. + + Returns: + dict: Dictionary with keys 'byt5_tokenizer', 'byt5_model', 'byt5_max_length'. + """ + byt5_tokenizer, byt5_model, byt5_max_length = create_byt5(args, device) + byt5_model = byt5_model.to(device=device) + return { + "byt5_tokenizer": byt5_tokenizer, + "byt5_model": byt5_model, + "byt5_max_length": byt5_max_length, + } + + +def create_byt5(args, device): + """ + Create ByT5 tokenizer and encoder, load weights if provided. + + Args: + args (dict): Configuration dictionary. + device (str or torch.device): Device to load the model onto. + + Returns: + tuple: (byt5_tokenizer, byt5_model, byt5_max_length) + """ + byt5_max_length = args['byt5_max_length'] + byt5_config = dict( + byt5_name=args['byT5_google_path'], + special_token=True, + color_special_token=True, + font_special_token=True, + color_ann_path=args['multilingual_prompt_format_color_path'], + font_ann_path=args['multilingual_prompt_format_font_path'], + multilingual=True, + ) + huggingface_cache_dir = None + byt5_model, byt5_tokenizer = load_byt5_and_byt5_tokenizer( + **byt5_config, + huggingface_cache_dir=huggingface_cache_dir, + device=device, + ) + + # Load custom checkpoint if provided + if args['byT5_ckpt_path'] is not None: + if "cuda" not in str(device): + byt5_state_dict = torch.load(args['byT5_ckpt_path'], map_location=device) + else: + byt5_state_dict = torch.load(args['byT5_ckpt_path'], map_location=device) + if 'state_dict' in byt5_state_dict: + sd = byt5_state_dict["state_dict"] + newsd = {} + for k, v in sd.items(): + if k.startswith('module.text_tower.encoder.'): + newsd[k[len('module.text_tower.encoder.'):]] = v + byt5_state_dict = newsd + byt5_model.load_state_dict(byt5_state_dict) + byt5_model.requires_grad_(False) + return byt5_tokenizer, byt5_model, byt5_max_length + + +def add_special_token( + tokenizer, + text_encoder, + add_color, + add_font, + color_ann_path, + font_ann_path, + multilingual=False, +): + """ + Add special tokens for color and font to tokenizer and text encoder. + + Args: + tokenizer: Huggingface tokenizer. + text_encoder: Huggingface T5 encoder. + add_color (bool): Whether to add color tokens. + add_font (bool): Whether to add font tokens. + color_ann_path (str): Path to color annotation JSON. + font_ann_path (str): Path to font annotation JSON. + multilingual (bool): Whether to use multilingual font tokens. + """ + with open(font_ann_path, 'r') as f: + idx_font_dict = json.load(f) + with open(color_ann_path, 'r') as f: + idx_color_dict = json.load(f) + + if multilingual: + font_token = [f'<{font_code[:2]}-font-{idx_font_dict[font_code]}>' for font_code in idx_font_dict] + else: + font_token = [f'' for i in range(len(idx_font_dict))] + color_token = [f'' for i in range(len(idx_color_dict))] + additional_special_tokens = [] + if add_color: + additional_special_tokens += color_token + if add_font: + additional_special_tokens += font_token + + tokenizer.add_tokens(additional_special_tokens, special_tokens=True) + # Set mean_resizing=False to avoid PyTorch LAPACK dependency + text_encoder.resize_token_embeddings(len(tokenizer), mean_resizing=False) + + +def load_byt5_and_byt5_tokenizer( + byt5_name='google/byt5-small', + special_token=False, + color_special_token=False, + font_special_token=False, + color_ann_path='assets/color_idx.json', + font_ann_path='assets/font_idx_512.json', + huggingface_cache_dir=None, + multilingual=False, + device=None, +): + """ + Load ByT5 encoder and tokenizer from Huggingface, and add special tokens if needed. + + Args: + byt5_name (str): Model name or path. + special_token (bool): Whether to add special tokens. + color_special_token (bool): Whether to add color tokens. + font_special_token (bool): Whether to add font tokens. + color_ann_path (str): Path to color annotation JSON. + font_ann_path (str): Path to font annotation JSON. + huggingface_cache_dir (str): Huggingface cache directory. + multilingual (bool): Whether to use multilingual font tokens. + device (str or torch.device): Device to load the model onto. + + Returns: + tuple: (byt5_text_encoder, byt5_tokenizer) + """ + byt5_tokenizer = AutoTokenizer.from_pretrained( + byt5_name, + cache_dir=huggingface_cache_dir, + ) + byt5_text_encoder = T5ForConditionalGeneration.from_pretrained( + byt5_name, + cache_dir=huggingface_cache_dir, + ).get_encoder() + + if "cuda" not in str(device): + device = torch.device(device) + else: + device = torch.device(device) + byt5_text_encoder = byt5_text_encoder.to(device) + + if special_token: + add_special_token( + byt5_tokenizer, + byt5_text_encoder, + add_color=color_special_token, + add_font=font_special_token, + color_ann_path=color_ann_path, + font_ann_path=font_ann_path, + multilingual=multilingual, + ) + return byt5_text_encoder, byt5_tokenizer + + +class ByT5Mapper(nn.Module): + """ + ByT5Mapper: Maps ByT5 encoder outputs to a new space, with optional residual connection. + + Args: + in_dim (int): Input dimension (must equal out_dim if use_residual). + out_dim (int): Output dimension after second linear layer. + hidden_dim (int): Hidden dimension for intermediate layer. + out_dim1 (int): Final output dimension. + use_residual (bool): Whether to use residual connection (default: True). + """ + + def __init__(self, in_dim, out_dim, hidden_dim, out_dim1, use_residual=True): + super().__init__() + if use_residual: + assert in_dim == out_dim + self.layernorm = nn.LayerNorm(in_dim) + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, out_dim) + self.fc3 = nn.Linear(out_dim, out_dim1) + self.use_residual = use_residual + self.act_fn = nn.GELU() + + def forward(self, x): + """ + Forward pass for ByT5Mapper. + + Args: + x (Tensor): Input tensor of shape (..., in_dim). + + Returns: + Tensor: Output tensor of shape (..., out_dim1). + """ + residual = x + x = self.layernorm(x) + x = self.fc1(x) + x = self.act_fn(x) + x = self.fc2(x) + x2 = self.act_fn(x) + x2 = self.fc3(x2) + if self.use_residual: + x2 = x2 + residual + return x2 diff --git a/trainer/models/hyvideo/models/text_encoders/byT5/format_prompt.py b/trainer/models/hyvideo/models/text_encoders/byT5/format_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..ff7b9c512f14f7c7810d67208164ecc948e9d5e3 --- /dev/null +++ b/trainer/models/hyvideo/models/text_encoders/byT5/format_prompt.py @@ -0,0 +1,86 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json + + +def closest_color(requested_color): + import webcolors + + min_colors = {} + for key, name in webcolors.CSS3_HEX_TO_NAMES.items(): + + r_c, g_c, b_c = webcolors.hex_to_rgb(key) + rd = (r_c - requested_color[0]) ** 2 + gd = (g_c - requested_color[1]) ** 2 + bd = (b_c - requested_color[2]) ** 2 + min_colors[(rd + gd + bd)] = name + return min_colors[min(min_colors.keys())] + + +def convert_rgb_to_names(rgb_tuple): + try: + import webcolors + + color_name = webcolors.rgb_to_name(rgb_tuple) + except ValueError: + color_name = closest_color(rgb_tuple) + return color_name + + +class MultilingualPromptFormat: + + def __init__( + self, + font_path: str = 'assets/glyph_sdxl_assets/multilingual_10-lang_idx.json', + color_path: str = 'assets/glyph_sdxl_assets/color_idx.json', + ): + with open(font_path, 'r') as f: + self.font_dict = json.load(f) + with open(color_path, 'r') as f: + self.color_dict = json.load(f) + + def format_prompt(self, texts, styles): + ''' + Text "{text}" in {color}, {type}. + ''' + + prompt = "" + for text, style in zip(texts, styles): + text_prompt = f'Text "{text}"' + + attr_list = [] + + # format color + if style["color"] is not None: + import webcolors + + hex_color = style["color"] + rgb_color = webcolors.hex_to_rgb(hex_color) + color_name = convert_rgb_to_names(rgb_color) + attr_list.append(f"") + + # format font + if style["font-family"] is not None: + attr_list.append(f"<{style['font-family'][:2]}-font-{self.font_dict[style['font-family']]}>") + attr_suffix = ", ".join(attr_list) + text_prompt += " in " + attr_suffix + text_prompt += ". " + else: + text_prompt += ". " + + prompt = prompt + text_prompt + return prompt diff --git a/trainer/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py b/trainer/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..236fa6f8cdc224b1ab6a84681fcecd1ace2dafee --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py @@ -0,0 +1,998 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Any, List, Tuple, Optional, Union, Dict + +import os +import torch +import torch.nn as nn +from einops import rearrange +from loguru import logger + +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from trainer.distributed import sequence_model_parallel_all_gather +from .modules.activation_layers import get_activation_layer +from .modules.norm_layers import get_norm_layer +from .modules.embed_layers import TimestepEmbedder, PatchEmbed, TextProjection, VisionProjection +from .modules.attention import parallel_attention +from .modules.posemb_layers import apply_rotary_emb, get_nd_rotary_pos_embed +from .modules.mlp_layers import MLP, MLPEmbedder, FinalLayer, LinearWarpforSingle +from .modules.modulate_layers import ModulateDiT, modulate, apply_gate +from .modules.token_refiner import SingleTokenRefiner + +from trainer.models.hyvideo.utils.communications import all_gather +from trainer.models.hyvideo.utils.infer_utils import torch_compile_wrapper +from trainer.models.hyvideo.models.text_encoders.byT5 import ByT5Mapper + +from trainer.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) +from trainer.models.prope.camera_rope import prope_qkv +from trainer.configs.models.dits import HunyuanVideoConfig +from trainer.distributed import (get_world_group) + +def is_blocks(n: str, m) -> bool: + print('is_blocks', n, flush=True) + return "double_blocks" in n and str.isdigit(n.split(".")[-1]) or "single_blocks" in n and str.isdigit(n.split(".")[-1]) + + +class MMDoubleStreamBlock(nn.Module): + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + self.attn_mode = attn_mode + + self.hidden_size = hidden_size + self.qkv_bias = qkv_bias + self.factory_kwargs = factory_kwargs + + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, factor=6, act_layer=get_activation_layer("silu"), **factory_kwargs + ) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_attn_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.img_attn_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.img_attn_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.img_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_mlp = MLP(hidden_size, mlp_hidden_dim, act_layer=get_activation_layer(mlp_act_type), bias=True, **factory_kwargs) + + self.txt_mod = ModulateDiT( + hidden_size, factor=6, act_layer=get_activation_layer("silu"), **factory_kwargs + ) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.txt_attn_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_attn_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_attn_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.txt_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.txt_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, act_layer=get_activation_layer(mlp_act_type), bias=True, **factory_kwargs) + + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + @torch_compile_wrapper() + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec_txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + text_mask=None, + attn_param=None, + is_flash=False, + block_idx=None, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec_txt).chunk(6, dim=-1) + + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) + + img_q = self.img_attn_q(img_modulated) + img_k = self.img_attn_k(img_modulated) + img_v = self.img_attn_v(img_modulated) + img_q = rearrange(img_q, "B L (H D) -> B L H D", H=self.heads_num) + img_k = rearrange(img_k, "B L (H D) -> B L H D", H=self.heads_num) + img_v = rearrange(img_v, "B L (H D) -> B L H D", H=self.heads_num) + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + + # 添加连续的camera pose，通过prope + img_q_prope, img_k_prope, img_v_prope, apply_fn_o = prope_qkv( + img_q.permute(0, 2, 1, 3), + img_k.permute(0, 2, 1, 3), + img_v.permute(0, 2, 1, 3), + viewmats=viewmats, + Ks=Ks, + ) # [batch, num_heads, seqlen, head_dim] + img_q_prope = img_q_prope.permute(0, 2, 1, 3) # [batch, seqlen, num_heads, head_dim] + img_k_prope = img_k_prope.permute(0, 2, 1, 3) # [batch, seqlen, num_heads, head_dim] + img_v_prope = img_v_prope.permute(0, 2, 1, 3) # [batch, seqlen, num_heads, head_dim] + + if freqs_cis is not None: + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) + txt_q = self.txt_attn_q(txt_modulated) + txt_k = self.txt_attn_k(txt_modulated) + txt_v = self.txt_attn_v(txt_modulated) + txt_q = rearrange(txt_q, "B L (H D) -> B L H D", H=self.heads_num) + txt_k = rearrange(txt_k, "B L (H D) -> B L H D", H=self.heads_num) + txt_v = rearrange(txt_v, "B L (H D) -> B L H D", H=self.heads_num) + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + + txt_q_before = txt_q.clone() + + # attn_mode = 'flash' if is_flash else self.attn_mode + attn_mode = 'torch_causal' # for ar model, the default mode is flex_causal + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + block_idx=block_idx, + ) + img_attn, txt_attn = attn[:, :img_q.shape[1]].contiguous(), attn[:, img_q.shape[1]:].contiguous() + + attn_prope = parallel_attention( + (img_q_prope, txt_q), + (img_k_prope, txt_k), + (img_v_prope, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + block_idx=block_idx, + ) + img_attn_prope, _ = attn_prope[:, :img_q_prope.shape[1]].contiguous(), attn[:, img_q_prope.shape[1]:].contiguous() + img_attn_prope = rearrange(img_attn_prope, "B L (H D) -> B H L D", H=self.heads_num) + img_attn_prope = apply_fn_o(img_attn_prope) # [batch, num_heads, seqlen, head_dim] + img_attn_prope = rearrange(img_attn_prope, "B H L D -> B L (H D)") + + img = img + apply_gate(self.img_attn_proj(img_attn) + self.img_attn_prope_proj(img_attn_prope), gate=img_mod1_gate) + img = img + apply_gate( + self.img_mlp( + modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale) + ), + gate=img_mod2_gate, + ) + + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)), + gate=txt_mod2_gate, + ) + return img, txt + + +class MMSingleStreamBlock(nn.Module): + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.attn_mode = attn_mode + + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim ** -0.5 + + self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_k = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_v = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_mlp = nn.Linear(hidden_size, mlp_hidden_dim, **factory_kwargs) + self.linear2 = LinearWarpforSingle(hidden_size + mlp_hidden_dim, hidden_size, bias=True, **factory_kwargs) + self.mlp_act = get_activation_layer(mlp_act_type)() + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.modulation = ModulateDiT(hidden_size, factor=3, act_layer=get_activation_layer("silu"), **factory_kwargs) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec_txt: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask=None, + attn_param=None, + is_flash=False, + ) -> torch.Tensor: + """Forward pass for the single stream block.""" + mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) + txt_mod_shift, txt_mod_scale, txt_mod_gate = self.modulation(vec_txt).chunk(3, dim=-1) + + x_pre_norm = self.pre_norm(x) + img_token, txt_token = x_pre_norm[:, :-txt_len, :], x_pre_norm[:, -txt_len:, :] + + img_mod = modulate(img_token, shift=mod_shift, scale=mod_scale) + txt_mod = modulate(txt_token, shift=txt_mod_shift, scale=txt_mod_scale) + x_mod = torch.cat([img_mod, txt_mod], dim=1) + + q = self.linear1_q(x_mod) + k = self.linear1_k(x_mod) + v = self.linear1_v(x_mod) + + q = rearrange(q, "B L (H D) -> B L H D", H=self.heads_num) + k = rearrange(k, "B L (H D) -> B L H D", H=self.heads_num) + v = rearrange(v, "B L (H D) -> B L H D", H=self.heads_num) + + mlp = self.linear1_mlp(x_mod) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + if is_flash: + attn_mode = 'flash' + else: + attn_mode = self.attn_mode + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + ) + output = self.linear2(attn, self.mlp_act(mlp)) + img_output, txt_output = output[:, :-txt_len, :], output[:, -txt_len:, :] + + img_output = apply_gate(img_output, gate=mod_gate) + txt_output = apply_gate(txt_output, gate=txt_mod_gate) + gate_output = torch.cat([img_output, txt_output], dim=1) + + return x + gate_output + +class ARHunyuanVideo_1_5_DiffusionTransformer(ModelMixin, ConfigMixin): + """ + HunyuanVideo Transformer backbone. + + Args: + patch_size (list): The size of the patch. + in_channels (int): The number of input channels. + out_channels (int): The number of output channels. + hidden_size (int): The hidden size of the transformer backbone. + heads_num (int): The number of attention heads. + mlp_width_ratio (float): Width ratio for the transformer MLPs. + mlp_act_type (str): Activation type for the transformer MLPs. + mm_double_blocks_depth (int): Number of double-stream transformer blocks. + mm_single_blocks_depth (int): Number of single-stream transformer blocks. + rope_dim_list (list): Rotary embedding dim for t, h, w. + qkv_bias (bool): Use bias in qkv projection. + qk_norm (bool): Whether to use qk norm. + qk_norm_type (str): Type of qk norm. + guidance_embed (bool): Use guidance embedding for distillation. + text_projection (str): Text input projection. Default is "single_refiner". + use_attention_mask (bool): If to use attention mask. + text_states_dim (int): Text encoder output dim. + text_states_dim_2 (int): Secondary text encoder output dim. + text_pool_type (str): Type for text pooling. + rope_theta (int): Rotary embedding theta parameter. + attn_mode (str): Attention mode identifier. + attn_param (dict): Attention parameter dictionary. + glyph_byT5_v2 (bool): Use ByT5 glyph module. + vision_projection (str): Vision condition embedding mode. + vision_states_dim (int): Vision encoder states input dim. + is_reshape_temporal_channels (bool): For video VAE adaptation. + use_cond_type_embedding (bool): Use condition type embedding. + """ + + _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 4, + concat_condition: bool = True, + out_channels: int = None, + hidden_size: int = 3072, + heads_num: int = 24, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + mm_double_blocks_depth: int = 20, + mm_single_blocks_depth: int = 40, + rope_dim_list: list = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + guidance_embed: bool = False, + use_meanflow: bool = False, + text_projection: str = "single_refiner", + use_attention_mask: bool = True, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + text_pool_type: str = None, + rope_theta: int = 256, + attn_mode: str = "flash", + attn_param: dict = None, + glyph_byT5_v2: bool = False, + vision_projection: str = "none", + vision_states_dim: int = 1280, + is_reshape_temporal_channels: bool = False, + use_cond_type_embedding: bool = False, + ideal_resolution: str = None, + ideal_task: str = None, + ): + super().__init__() + factory_kwargs = {} + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.use_attention_mask = use_attention_mask + self.text_projection = text_projection + self.attn_mode = attn_mode + self.text_pool_type = text_pool_type + self.text_states_dim = text_states_dim + self.text_states_dim_2 = text_states_dim_2 + self.vision_states_dim = vision_states_dim + + self.glyph_byT5_v2 = glyph_byT5_v2 + if self.glyph_byT5_v2: + self.byt5_in = ByT5Mapper( + in_dim=1472, + out_dim=2048, + hidden_dim=2048, + out_dim1=hidden_size, + use_residual=False + ) + + if hidden_size % heads_num != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}" + ) + pe_dim = hidden_size // heads_num + if sum(rope_dim_list) != pe_dim: + raise ValueError( + f"Got {rope_dim_list} but expected positional dim {pe_dim}" + ) + self.hidden_size = hidden_size + self.heads_num = heads_num + + self.img_in = PatchEmbed( + self.patch_size, self.in_channels, self.hidden_size, + is_reshape_temporal_channels=is_reshape_temporal_channels, concat_condition=concat_condition, **factory_kwargs + ) + + # Vision projection + if vision_projection == "linear": + self.vision_in = VisionProjection( + input_dim=self.vision_states_dim, output_dim=self.hidden_size + ) + else: + self.vision_in = None + + # Text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + text_states_dim, + hidden_size, + heads_num, + depth=2, + **factory_kwargs, + ) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}" + ) + + # time modulation + self.time_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + self.vector_in = ( + MLPEmbedder( + self.config.text_states_dim_2, self.hidden_size, **factory_kwargs + ) if self.text_pool_type is not None else None + ) + self.guidance_in = ( + TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + if guidance_embed + else None + ) + + self.time_r_in = ( + TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + if use_meanflow + else None + ) + + + self.double_blocks = nn.ModuleList( + [ + MMDoubleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode=attn_mode, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + for _ in range(mm_double_blocks_depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + MMSingleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode=attn_mode, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) + for _ in range(mm_single_blocks_depth) + ] + ) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs, + ) + + # STA + if attn_param is None: + self.attn_param = { + # STA + 'win_size': [[3, 3, 3]], + 'win_type': "fixed", + 'win_ratio': 10, + 'tile_size': [6, 8, 8], + # SSTA + 'ssta_topk': 64, + 'ssta_threshold': 0.0, + 'ssta_lambda': 0.7, + 'ssta_sampling_type': 'importance', + 'ssta_adaptive_pool': None, + # flex-block-attn: + "attn_sparse_type": 'ssta', + "attn_pad_type": 'zero', + "attn_use_text_mask": 1, + "attn_mask_share_within_head": 0, + } + else: + self.attn_param = attn_param + + if attn_mode == 'flex-block-attn': + self.register_to_config(attn_param=self.attn_param) + + if use_cond_type_embedding: + self.cond_type_embedding = nn.Embedding(3, self.hidden_size) + self.cond_type_embedding.weight.data.fill_(0) + assert self.glyph_byT5_v2, "text type embedding is only used when glyph_byT5_v2 is True" + assert vision_projection is not None, "text type embedding is only used when vision_projection is not None" + # 0: text_encoder feature + # 1: byt5 feature + # 2: vision_encoder feature + else: + self.cond_type_embedding = None + + self.gradient_checkpointing = False + + def load_hunyuan_state_dict(self, model_path): + load_key = "module" + bare_model = "unknown" + + if model_path.endswith('.safetensors'): + from safetensors.torch import load_file + state_dict = load_file(model_path, device="cpu") + else: + state_dict = torch.load( + model_path, map_location="cpu", weights_only=True + ) + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError( + f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}." + ) + + result = self.load_state_dict(state_dict, strict=False) + + if result.missing_keys: + logger.info("[load.py] Missing keys when loading state_dict:") + for key in result.missing_keys: + logger.info(f"[load.py] Missing key: {key}") + if result.unexpected_keys: + logger.info("[load.py] Unexpected keys when loading state_dict:") + for key in result.unexpected_keys: + logger.info(f"[load.py] Unexpected key: {key}") + if result.missing_keys or result.unexpected_keys: + raise ValueError(f"Missing: {result.missing_keys}, Unexpected: {result.unexpected_keys}") + + return result + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def get_rotary_pos_embed(self, rope_sizes): + target_ndim = 3 + head_dim = self.hidden_size // self.heads_num + rope_dim_list = self.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert ( + sum(rope_dim_list) == head_dim + ), "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + + def reorder_txt_token(self, byt5_txt, txt, byt5_text_mask, text_mask, zero_feat=False, is_reorder=True): + if is_reorder: + reorder_txt = [] + reorder_mask = [] + for i in range(text_mask.shape[0]): + byt5_text_mask_i = byt5_text_mask[i].bool() + text_mask_i = text_mask[i].bool() + + byt5_txt_i = byt5_txt[i] + txt_i = txt[i] + if zero_feat: + # When using block mask with approximate computation, set pad to zero to reduce error + pad_byt5 = torch.zeros_like(byt5_txt_i[~byt5_text_mask_i]) + pad_text = torch.zeros_like(txt_i[~text_mask_i]) + reorder_txt_i = torch.cat( + [byt5_txt_i[byt5_text_mask_i], txt_i[text_mask_i], pad_byt5, pad_text], dim=0 + ) + else: + reorder_txt_i = torch.cat( + [byt5_txt_i[byt5_text_mask_i], txt_i[text_mask_i], byt5_txt_i[~byt5_text_mask_i], txt_i[~text_mask_i]], dim=0 + ) + reorder_mask_i = torch.cat( + [byt5_text_mask_i[byt5_text_mask_i], text_mask_i[text_mask_i], + byt5_text_mask_i[~byt5_text_mask_i], text_mask_i[~text_mask_i]], dim=0 + ) + + reorder_txt.append(reorder_txt_i) + reorder_mask.append(reorder_mask_i) + + reorder_txt = torch.stack(reorder_txt) + reorder_mask = torch.stack(reorder_mask).to(dtype=torch.int64) + else: + reorder_txt = torch.concat([byt5_txt, txt], dim=1) + reorder_mask = torch.concat([byt5_text_mask, text_mask], dim=1).to(dtype=torch.int64) + + return reorder_txt, reorder_mask + + def add_discrete_action_parameters(self): + self.action_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu") + ) + nn.init.zeros_(self.action_in.mlp[2].weight) + # 如果你想初始化偏置为0，执行以下操作： + if self.action_in.mlp[2].bias is not None: + nn.init.zeros_(self.action_in.mlp[2].bias) + + # prope参数,零初始化 + for block in self.double_blocks: + block.img_attn_prope_proj = nn.Linear(block.hidden_size, block.hidden_size, bias=block.qkv_bias, **block.factory_kwargs) + nn.init.zeros_(block.img_attn_prope_proj.weight) + + # 如果你想初始化偏置为0，执行以下操作： + if block.img_attn_prope_proj.bias is not None: + nn.init.zeros_(block.img_attn_prope_proj.bias) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + timestep_txt: torch.Tensor, + text_states: torch.Tensor, + text_states_2: torch.Tensor, + encoder_attention_mask: torch.Tensor, + timestep_r=None, + vision_states: torch.Tensor = None, + output_features=False, + output_features_stride=8, + attention_kwargs: Optional[Dict[str, Any]] = None, + freqs_cos: Optional[torch.Tensor] = None, + freqs_sin: Optional[torch.Tensor] = None, + return_dict: bool = False, + guidance=None, + mask_type="t2v", + extra_kwargs=None, + action: Optional[torch.Tensor] = None, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + + if guidance is None: + guidance = torch.tensor( + [6016.0], device=hidden_states.device, dtype=torch.bfloat16 + ) + + img = x = hidden_states + text_mask = encoder_attention_mask + t = timestep + txt = text_states + bs, _, ot, oh, ow = x.shape + tt, th, tw = ( + ot // self.patch_size[0], + oh // self.patch_size[1], + ow // self.patch_size[2], + ) + self.attn_param['thw'] = [tt, th, tw] + if freqs_cos is None and freqs_sin is None: + freqs_cos, freqs_sin = self.get_rotary_pos_embed((tt, th, tw)) + + img = self.img_in(img) + + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + if sp_world_size > 1: + sp_size = sp_world_size + sp_rank = rank_in_sp_group + if img.shape[1] % sp_size != 0: + n_token = img.shape[1] + assert n_token > (n_token // sp_size + 1) * (sp_size - 1), f'Too short context length for SP {sp_size}' + img = torch.chunk(img, sp_size, dim=1)[sp_rank] + freqs_cos = torch.chunk(freqs_cos, sp_size, dim=0)[sp_rank] + freqs_sin = torch.chunk(freqs_sin, sp_size, dim=0)[sp_rank] + + viewmats = torch.chunk(viewmats, sp_size, dim=1)[sp_rank] + Ks = torch.chunk(Ks, sp_size, dim=1)[sp_rank] + action = torch.chunk(action, sp_size, dim=0)[sp_rank] + t = torch.chunk(t, sp_size, dim=0)[sp_rank] + + # Prepare modulation vectors + vec_txt = self.time_in(timestep_txt) # 为了txt prompt单独计算 + vec = self.time_in(t) + + if text_states_2 is not None: + vec_2 = self.vector_in(text_states_2) + vec = vec + vec_2 + + if self.guidance_embed: + if guidance is None: + raise ValueError( + "Didn't get guidance strength for guidance distilled model." + ) + vec = vec + self.guidance_in(guidance) + + if timestep_r is not None: + vec = vec + self.time_r_in(timestep_r) + + vec = vec + self.action_in(action) # 添加离散的action + + # Embed text tokens + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in(txt, timestep_txt, text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}" + ) + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + torch.zeros_like(txt[:, :, 0], device=text_mask.device, dtype=torch.long) + ) + txt = txt + cond_emb + + if self.glyph_byT5_v2: + byt5_text_states = extra_kwargs["byt5_text_states"] + byt5_text_mask = extra_kwargs["byt5_text_mask"] + byt5_txt = self.byt5_in(byt5_text_states) + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + torch.ones_like(byt5_txt[:, :, 0], device=byt5_txt.device, dtype=torch.long) + ) + byt5_txt = byt5_txt + cond_emb + txt, text_mask = self.reorder_txt_token( + byt5_txt, txt, byt5_text_mask, text_mask, zero_feat=True + ) + + if self.vision_in is not None and vision_states is not None: + extra_encoder_hidden_states = self.vision_in(vision_states) + # If t2v, set extra_attention_mask to 0 to avoid attention to semantic tokens + if mask_type == "t2v" and torch.all(vision_states == 0): + extra_attention_mask = torch.zeros( + (bs, extra_encoder_hidden_states.shape[1]), + dtype=text_mask.dtype, + device=text_mask.device, + ) + # Set vision tokens to zero to mitigate potential block mask error in SSTA + extra_encoder_hidden_states = extra_encoder_hidden_states * 0.0 + else: + extra_attention_mask = torch.ones( + (bs, extra_encoder_hidden_states.shape[1]), + dtype=text_mask.dtype, + device=text_mask.device, + ) + # Ensure valid tokens precede padding tokens + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + 2 * torch.ones_like( + extra_encoder_hidden_states[:, :, 0], + dtype=torch.long, + device=extra_encoder_hidden_states.device, + ) + ) + extra_encoder_hidden_states = extra_encoder_hidden_states + cond_emb + + txt, text_mask = self.reorder_txt_token( + extra_encoder_hidden_states, txt, extra_attention_mask, text_mask + ) + + # print('text embedding shape:', txt.shape, text_mask.shape) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + # mask the txt tokens based on the text_mask + txt = txt[text_mask.bool().to(txt.device)].unsqueeze(0) + + # Pass through double-stream blocks + for index, block in enumerate(self.double_blocks): + force_full_attn = ( + self.attn_mode in ["flex-block-attn"] + and self.attn_param["win_type"] == "hybrid" + and self.attn_param["win_ratio"] > 0 + and ( + (index + 1) % self.attn_param["win_ratio"] == 0 + or (index + 1) == len(self.double_blocks) + ) + ) + self.attn_param["layer-name"] = f"double_block_{index+1}" + + img, txt = block( + img=img, + txt=txt, + vec_txt=vec_txt, + vec=vec, + freqs_cis=freqs_cis, + text_mask=None, # we have masked txt tokens already, set None here + attn_param=self.attn_param, + is_flash=force_full_attn, + block_idx=index, + viewmats=viewmats, + Ks=Ks, + ) + + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + # Merge image and text for single-stream blocks + x = torch.cat((img, txt), 1) + features_list = [] if output_features else None + if len(self.single_blocks) > 0: + for index, block in enumerate(self.single_blocks): + force_full_attn = ( + self.attn_mode in ["flex-block-attn"] + and self.attn_param["win_type"] == "hybrid" + and self.attn_param["win_ratio"] > 0 + and ( + (index + 1) % self.attn_param["win_ratio"] == 0 + or (index + 1) == len(self.single_blocks) + ) + ) + self.attn_param["layer-name"] = f"single_block_{index+1}" + x = block( + x=x, + vec_txt=vec_txt, + vec=vec, + txt_len=txt_seq_len, + freqs_cis=(freqs_cos, freqs_sin), + text_mask=text_mask, + attn_param=self.attn_param, + is_flash=force_full_attn, + ) + + if output_features and index % output_features_stride == 0: + features_list.append(x[:, :img_seq_len, ...]) + img = x[:, :img_seq_len, ...] + + # Final Layer + img = self.final_layer(img, vec) + if sp_world_size > 1: + img = sequence_model_parallel_all_gather(img, dim=1) + img = self.unpatchify(img, tt, th, tw) + assert return_dict is False, "return_dict is not supported." + if output_features: + features_list = torch.stack(features_list, dim=0) + if sp_world_size > 1: + features_list = sequence_model_parallel_all_gather(features_list, dim=2) + else: + features_list = None + return (img, features_list) + + def unpatchify(self, x, t, h, w): + """ + Unpatchify a tensorized input back to frame format. + + Args: + x (Tensor): Input tensor of shape (N, T, patch_size**2 * C) + t (int): Number of time steps + h (int): Height in patch units + w (int): Width in patch units + + Returns: + Tensor: Output tensor of shape (N, C, t * pt, h * ph, w * pw) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + return imgs + + def set_attn_mode(self, attn_mode: str): + self.attn_mode = attn_mode + for block in self.double_blocks: + block.attn_mode = attn_mode + for block in self.single_blocks: + block.attn_mode = attn_mode diff --git a/trainer/models/hyvideo/models/transformers/modules/activation_layers.py b/trainer/models/hyvideo/models/transformers/modules/activation_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..b3f3fa95d91191c85bd24180aaec3e6b5c7d590f --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/activation_layers.py @@ -0,0 +1,39 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/trainer/models/hyvideo/models/transformers/modules/attention.py b/trainer/models/hyvideo/models/transformers/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..6f29fd850808061f26d864ce03066f9cbc069f5e --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/attention.py @@ -0,0 +1,351 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import einops +import torch +from typing import Optional +from loguru import logger +import numpy as np +import torch.nn.functional as F + +from trainer.models.hyvideo.commons.parallel_states import get_parallel_state +from trainer.models.hyvideo.utils.flash_attn_no_pad import ( + flash_attn_no_pad, + flash_attn_no_pad_v3, + flex_attn_no_pad +) +from trainer.models.hyvideo.commons import maybe_fallback_attn_mode + + +from trainer.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) +from trainer.distributed.communication_op import ( + sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D) + +try: + from torch.nn.attention.flex_attention import flex_attention + + flex_attention = torch.compile(flex_attention, dynamic=False) + torch._dynamo.config.cache_size_limit = 192 + torch._dynamo.config.accumulated_cache_size_limit = 192 + flex_mask_cache = {} +except Exception: + logger.warning("Could not load Sliding Tile Attention of FlexAttn.") + +from trainer.models.hyvideo.models.transformers.modules.ssta_attention import ssta_3d_attention +from trainer.models.hyvideo.commons.infer_state import get_infer_state + + + +@torch.compiler.disable +def attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + drop_rate: float = 0.0, + attn_mask: Optional[torch.Tensor] = None, + causal: bool = False, + attn_mode: str = "flash", +) -> torch.Tensor: + """ + Compute attention using flash_attn_no_pad or torch scaled_dot_product_attention. + + Args: + q: Query tensor of shape [B, L, H, D] + k: Key tensor of shape [B, L, H, D] + v: Value tensor of shape [B, L, H, D] + drop_rate: Dropout rate for attention weights. + attn_mask: Optional attention mask of shape [B, L]. + causal: Whether to apply causal masking. + attn_mode: Attention mode, either "flash" or "torch". Defaults to "flash". + + Returns: + Output tensor after attention of shape [B, L, H*D] + """ + attn_mode = maybe_fallback_attn_mode(attn_mode) + + if attn_mode == "torch": + # transpose q,k,v dim to fit scaled_dot_product_attention + query = q.transpose(1, 2) # B * H * L * D + key = k.transpose(1, 2) # B * H * L * D + value = v.transpose(1, 2) # B * H * L * D + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [torch.int64, torch.int32]: + assert attn_mask.max() <= 1 and attn_mask.min() >= 0, f'Integer attention mask must be between 0 and 1 for torch attention.' + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError(f'Float attention mask is not implemented for torch attention.') + attn_mask1 = einops.rearrange(attn_mask, 'b l -> b 1 l 1') + attn_mask2 = einops.rearrange(attn_mask1, 'b 1 l 1 -> b 1 1 l') + attn_mask = attn_mask1 & attn_mask2 + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) + + # transpose back + x = x.transpose(1, 2) # B * L * H * D + b, s, h, d = x.shape + out = x.reshape(b, s, -1) + return out + else: + # flash mode (default) + qkv = torch.stack([q, k, v], dim=2) + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +@torch.compiler.disable +def parallel_attention(q, k, v, img_q_len, img_kv_len, + attn_mode=None, text_mask=None, + attn_param=None, + block_idx=None, + ): + return sequence_parallel_attention(q, k, v, img_q_len, img_kv_len, attn_mode, text_mask, attn_param=attn_param, block_idx=block_idx) + + +def sequence_parallel_attention(q, k, v, + img_q_len, img_kv_len, + attn_mode=None, text_mask=None, + attn_param=None, + block_idx=None, + ): + assert attn_mode is not None + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + + if sp_world_size > 1: + sp_size = sp_world_size + sp_rank = rank_in_sp_group + + # batch_size, seq_len, attn_heads, head_dim + query = sequence_model_parallel_all_to_all_4D(query, scatter_dim=2, gather_dim=1) + key = sequence_model_parallel_all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = sequence_model_parallel_all_to_all_4D(value, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // sp_size + return encoder_state.narrow( + dim, sp_rank * local_heads, local_heads + ) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + attn_mode = maybe_fallback_attn_mode(attn_mode, get_infer_state(), block_idx) + + if attn_mode == "sageattn": + from sageattention import sageattn + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + hidden_states = sageattn(query, key, value, tensor_layout="NHD", is_causal=False) + elif attn_mode == "torch": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + if text_mask is not None: + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + else: + attn_mask = None + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [torch.int64, torch.int32]: + assert attn_mask.max() <= 1 and attn_mask.min() >= 0, f'Integer attention mask must be between 0 and 1 for torch attention.' + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError(f'Float attention mask is not implemented for torch attention.') + + # transpose q,k,v dim to fit scaled_dot_product_attention + query = query.transpose(1, 2) # B * Head_num * length * dim + key = key.transpose(1, 2) # B * Head_num * length * dim + value = value.transpose(1, 2) # B * Head_num * length * dim + if attn_mask is not None: + attn_mask1 = einops.rearrange(attn_mask, 'b l -> b 1 l 1') + attn_mask2 = einops.rearrange(attn_mask1, 'b 1 l 1 -> b 1 1 l') + attn_mask = attn_mask1 & attn_mask2 + hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) + + # add new attn for chunk-wise attn + elif attn_mode == "torch_causal": # now: we set text_mask = None + # attention: here we concat the encoder text sequence first, then apply causal attention + vision_seq_length = query.shape[1] + text_seq_length = encoder_query.shape[1] + total_seq_length = vision_seq_length + text_seq_length + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + # prepare causal mask for chunk-wise attention + latent_seq_length = 1560 # set for hunyuanvideo 1.5, which is for 480 * 832 resolution + chunk_seq_length = 1560 * 4 + chunk_num = (vision_seq_length) // chunk_seq_length + causal_mask = torch.zeros((total_seq_length, total_seq_length), device=query.device) + causal_mask[:, :text_seq_length] = 1 # no attention for the rest + for i in range(chunk_num): + start_i = text_seq_length + i * chunk_seq_length + end_i = min(start_i + chunk_seq_length, total_seq_length) + for j in range(i + 1): + start_j = text_seq_length + j * chunk_seq_length + end_j = min(start_j + chunk_seq_length, total_seq_length) + # full attention within chunk i for j == i, causal for j < i + causal_mask[start_i:end_i, start_j:end_j] = 1 + + causal_mask = causal_mask.unsqueeze(0).unsqueeze(1) # 1, 1, S, S + causal_mask = causal_mask.expand(query.shape[0], 1, -1, -1) + causal_mask = causal_mask.to(torch.bool) # Force bool dtype + + query = query.transpose(1, 2) # B * H * L * D + key = key.transpose(1, 2) # B * H * L * D + value = value.transpose(1, 2) # B * H * L * D + + hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=causal_mask, dropout_p=0.0, is_causal=False) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) # [B, S, H, D] + + # return back to the original order: [query, encoder_query] + hidden_states, encoder_hidden_states = hidden_states[:, text_seq_length:, :, :], hidden_states[:, :text_seq_length, :, :] + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + elif attn_mode == "flash2": + query = torch.cat([query, encoder_query], dim=1) # vision token first + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + elif attn_mode == "flex_causal": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flex_attn_no_pad(qkv, attn_mask, causal=True, dropout_p=0.0, softmax_scale=None) + + elif attn_mode == "flash3": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad_v3(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + elif attn_mode == "flex-block-attn": + sparse_type = attn_param["attn_sparse_type"] # sta/block_attn/ssta + ssta_threshold = attn_param["ssta_threshold"] + ssta_lambda = attn_param["ssta_lambda"] + ssta_sampling_type = attn_param["ssta_sampling_type"] + ssta_adaptive_pool = attn_param["ssta_adaptive_pool"] + + attn_pad_type = attn_param["attn_pad_type"] # repeat/zero + attn_use_text_mask = attn_param["attn_use_text_mask"] + attn_mask_share_within_head = attn_param["attn_mask_share_within_head"] + + ssta_topk = attn_param["ssta_topk"] + thw = attn_param["thw"] + tile_size = attn_param["tile_size"] + win_size = attn_param["win_size"][0].copy() + + def get_image_tile(tile_size): + block_size = np.prod(tile_size) + if block_size == 384: + tile_size = (1, 16, 24) + elif block_size == 128: + tile_size = (1, 16, 8) + elif block_size == 64: + tile_size = (1, 8, 8) + elif block_size == 16: + tile_size = (1, 4, 4) + else: + raise ValueError(f"Error tile_size {tile_size}, only support in [16, 64, 128, 384]") + return tile_size + + if thw[0] == 1: + tile_size = get_image_tile(tile_size) + win_size = [1, 1, 1] + elif thw[0] <= 31: # 16fps: 5 * 16 / 4 + 1 = 21; 24fps: 5 * 24 / 4 + 1 = 31 + ssta_topk = ssta_topk // 2 + + # Concatenate and permute query, key, value to (B, H, S, D) + query = torch.cat([query, encoder_query], dim=1).permute(0, 2, 1, 3) + key = torch.cat([key, encoder_key], dim=1).permute(0, 2, 1, 3) + value = torch.cat([value, encoder_value], dim=1).permute(0, 2, 1, 3) + + assert ( + query.shape[-1] == 128 + ), "The last dimension of query, key and value must be 128 for flex-block-attn." + + hidden_states = ssta_3d_attention( + query, + key, + value, + thw, + topk=ssta_topk, + tile_thw=tile_size, + kernel_thw=win_size, + text_len=encoder_sequence_length, + sparse_type=sparse_type, + threshold=ssta_threshold, + lambda_=ssta_lambda, + pad_type=attn_pad_type, + text_mask=text_mask if attn_use_text_mask else None, + sampling_type=ssta_sampling_type, + adaptive_pool=ssta_adaptive_pool, + mask_share_within_head=attn_mask_share_within_head, + ) + hidden_states, sparse_ratio = hidden_states + hidden_states = hidden_states.permute(0, 2, 1, 3) + + else: + raise NotImplementedError( + f'Unsupported attention mode: {attn_mode}. Only torch, flash, flash3, sageattn and flex-block-attn are supported.' + ) + + if sp_world_size > 1: + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), dim=1) + hidden_states = sequence_model_parallel_all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = sequence_model_parallel_all_gather(encoder_hidden_states, dim=2).contiguous() + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = hidden_states.shape + hidden_states = hidden_states.reshape(b, s, -1) + + return hidden_states diff --git a/trainer/models/hyvideo/models/transformers/modules/embed_layers.py b/trainer/models/hyvideo/models/transformers/modules/embed_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..36acd660d85d4dd82603a8d57a9b8930155b03d5 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/embed_layers.py @@ -0,0 +1,215 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import torch.nn as nn + +from trainer.models.hyvideo.commons import to_2tuple + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + is_reshape_temporal_channels=False, + concat_condition=True, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + # Only support concat mode (multitask mask training) + orig_in_chans = in_chans + if concat_condition: + if is_reshape_temporal_channels: + in_chans = in_chans + in_chans//2 + 1 + else: + in_chans = in_chans * 2 + 1 + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + + nn.init.xavier_uniform_( + self.proj.weight[:, :orig_in_chans].view(self.proj.weight[:, :orig_in_chans].size(0), -1) + ) + # Special initialization for concat mode + nn.init.zeros_(self.proj.weight[:, orig_in_chans:].view(self.proj.weight[:, orig_in_chans:].size(0), -1)) + + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + + +class VisionProjection(torch.nn.Module): + + def __init__(self, input_dim, output_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + torch.nn.Linear(input_dim, input_dim), + torch.nn.GELU(), + torch.nn.Linear(input_dim, output_dim), + torch.nn.LayerNorm(output_dim) + ) + + + def forward(self, vision_embeds): + return self.proj(vision_embeds) + +class ClipVisionProjection(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.up = nn.Linear(in_channels, out_channels * 3) + self.down = nn.Linear(out_channels * 3, out_channels) + torch.nn.init.zeros_(self.down.weight) + torch.nn.init.zeros_(self.down.bias) + + def forward(self, x): + projected_x = self.down(nn.functional.silu(self.up(x))) + return projected_x + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding( + t, self.frequency_embedding_size, self.max_period + ).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/trainer/models/hyvideo/models/transformers/modules/mlp_layers.py b/trainer/models/hyvideo/models/transformers/modules/mlp_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..a89ec8d8066c822d3bb910a92f59ad775372491f --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/mlp_layers.py @@ -0,0 +1,137 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from trainer.models.hyvideo.commons import to_2tuple +from .modulate_layers import modulate + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity() + self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class LinearWarpforSingle(nn.Module): + def __init__(self, in_dim: int, out_dim: int, bias=False, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.fc = nn.Linear(in_dim, out_dim, bias=bias, **factory_kwargs) + + def forward(self, x, y): + input = torch.cat([x.contiguous(), y.contiguous()], dim=2).contiguous() + return self.fc(input) + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs, + ) + else: + out_size = ( + patch_size[0] * patch_size[1] * patch_size[2] if len(patch_size) == 3 else patch_size[0] * patch_size[1] + ) * out_channels + self.linear = nn.Linear( + hidden_size, + out_size, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/trainer/models/hyvideo/models/transformers/modules/modulate_layers.py b/trainer/models/hyvideo/models/transformers/modules/modulate_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..3c55c5ef8a2b4e383884b6a5a3cb9269b6b2a3af --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/modulate_layers.py @@ -0,0 +1,114 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + scale = scale.unsqueeze(0) + latent_length = scale.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + # operate on the hidden_states + scale = scale.repeat_interleave(token_length, dim=1).type_as(x) + + return x * (1 + scale) + elif scale is None: + shift = shift.unsqueeze(0) + latent_length = shift.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + # operate on the hidden_states + shift = shift.repeat_interleave(token_length, dim=1).type_as(x) + return x + shift + else: + shift = shift.unsqueeze(0) + scale = scale.unsqueeze(0) + latent_length = shift.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + + scale = scale.repeat_interleave(token_length, dim=1).type_as(x) + shift = shift.repeat_interleave(token_length, dim=1).type_as(x) + return x * (1 + scale) + shift + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + gate = gate.unsqueeze(0) + latent_length = gate.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + gate = gate.repeat_interleave(token_length, dim=1).type_as(x) + if tanh: + return x * gate.tanh() + else: + return x * gate + + +def ckpt_wrapper(module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + diff --git a/trainer/models/hyvideo/models/transformers/modules/norm_layers.py b/trainer/models/hyvideo/models/transformers/modules/norm_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..1d02aa8d220c225581e4028062784f0104ec8598 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/norm_layers.py @@ -0,0 +1,97 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def reset_parameters(self): + if hasattr(self, "weight"): + self.weight.fill_(1) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/trainer/models/hyvideo/models/transformers/modules/posemb_layers.py b/trainer/models/hyvideo/models/transformers/modules/posemb_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..5c272bcc5e151fa028ffa6b6aa2c638964ec5364 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/posemb_layers.py @@ -0,0 +1,326 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from typing import Union, Tuple, List + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = ( + x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) + ) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex( + xq.float().reshape(*xq.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to( + xq.device + ) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex( + xk.float().reshape(*xk.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd( + start, *args, dim=len(rope_dim_list) + ) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor ** (dim / (dim - 2)) + + freqs = 1.0 / ( + theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim) + ) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar( + torch.ones_like(freqs), freqs + ) # complex64 # [S, D/2] + return freqs_cis diff --git a/trainer/models/hyvideo/models/transformers/modules/ssta_attention.py b/trainer/models/hyvideo/models/transformers/modules/ssta_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..3a59ad18bbe4a70dbf2323ffb7264b7067b2ccd0 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/ssta_attention.py @@ -0,0 +1,726 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import numpy as np +from einops import rearrange +from functools import lru_cache + +def tile(x, canvas_thw, tile_thw, sp_size=1): + r"""Rearrange tensor into tiles for block-based attention. + + Args: + x: Input tensor with shape (b, head, s, d) where s = t * h * w + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Rearranged tensor organized by tiles + """ + b, h, s, d = x.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} == s:{s}" + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t/tile_t_dim) + n_h = int(h/tile_h_dim) + n_w = int(w/tile_w_dim) + x = rearrange(x, "b head (sp t h w) d -> b head (t sp h w) d", sp=sp_size, t=t // sp_size, h=h, w=w) + return rearrange(x, + "b h (n_t ts_t n_h ts_h n_w ts_w) d -> b h (n_t n_h n_w ts_t ts_h ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim) + +def untile(x, canvas_thw, tile_thw, sp_size=1): + r"""Reverse the tiling operation to restore original tensor layout. + + Args: + x: Tiled tensor + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Restored tensor with original layout + """ + t, h, w = canvas_thw + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t/tile_t_dim) + n_h = int(h/tile_h_dim) + n_w = int(w/tile_w_dim) + + x = rearrange(x, + "b h (n_t n_h n_w ts_t ts_h ts_w) d -> b h (n_t ts_t n_h ts_h n_w ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim) + return rearrange(x, "b head (t sp h w) d -> b head (sp t h w) d", sp=sp_size, t=t // sp_size, h=h, w=w) + +def get_tile_t_h_w(tile_id, tile_thw_dim): + """Extract temporal, height, and width indices from a flattened tile ID.""" + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw_dim + tile_t = tile_id // (tile_h_dim * tile_w_dim) + tile_h = (tile_id % (tile_h_dim * tile_w_dim)) // tile_w_dim + tile_w = tile_id % tile_w_dim + + return tile_t, tile_h, tile_w +def importance_sampling(q, k, topk, threshold=0.0, lambda_=0.9, adaptive_pool=None): + r"""Select top-k blocks based on importance scores considering both similarity and redundancy. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Threshold parameter (not implemented) + lambda_: Weight factor balancing similarity and redundancy + adaptive_pool: Adaptive pooling parameter (unused) + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + raise NotImplementedError("importance_sampling with threshold not implemented") + + q = q / q.norm(dim=-1, keepdim=True) + k = k / k.norm(dim=-1, keepdim=True) + gate_similarity = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate_similarity = (gate_similarity + 1.0) / 2.0 + gate_unique = torch.einsum("bhsd,bhkd->bhsk", k, k) + gate_unique = (gate_unique + 1.0) / 2.0 + + B, H, K_num, D = k.shape + diag_indices = torch.arange(K_num, device=k.device) + gate_unique[:, :, diag_indices, diag_indices] = torch.nan + + mean_redundancy = torch.nanmean(gate_unique, dim=-2, keepdim=True) + + importance_scores = lambda_ * gate_similarity - (1 - lambda_) * mean_redundancy + + topk = min(topk, importance_scores.size(-1)) + _, top_block_indices = importance_scores.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + +def similarity_sampling(q, k, topk, + threshold=0.0, + block_num=None, + adaptive_pool=None, + temperature=0.01): + r"""Select top-k blocks based on similarity scores between query and key averages. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Cumulative score threshold for dynamic topk selection + block_num: Total number of blocks (unused) + adaptive_pool: Adaptive pooling parameter (unused) + temperature: Temperature scaling for softmax + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate = gate / temperature + gate_ = torch.softmax(gate, dim=-1) + sorted_gate, sorted_indices = torch.sort(gate_, dim=-1, descending=True) + cum_scores = torch.cumsum(sorted_gate, dim=-1) + above_threshold = cum_scores >= threshold + has_any_above = above_threshold.any(dim=-1, keepdim=True) + # Find the first position exceeding threshold, add 1 to get the number of blocks to select + dynamic_topk = above_threshold.int().argmax(dim=-1, keepdim=True) + 1 + dynamic_topk = torch.where(has_any_above, dynamic_topk, + torch.full_like(dynamic_topk, topk)) + # Generate fixed-shape indices and pad with top1 indices for parts not meeting threshold + # Limit minimum k value + dynamic_topk = torch.clamp(dynamic_topk, min=8, max=topk) + indices = torch.arange(gate.size(-1), device=gate.device).expand(gate.size()) + mask = (indices < dynamic_topk).int() + # Use top1 to pad to topk to ensure same size + top_block_indices = torch.gather(sorted_indices, -1, indices) * mask + (1 - mask) * sorted_indices[..., 0:1] + else: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + topk = min(topk, gate.size(-1)) + _, top_block_indices = gate.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + +def create_moba_3d_mask(q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + add_text_mask=False, + threshold=0.0, + lambda_=None, + mask_share_within_head=True, + q_block_avg_pool=True, + adaptive_pool=None, + sampling_type=None): + r"""Create MOBA (Mixture of Block Attention) 3D mask for sparse attention. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + add_text_mask: Whether to add text mask + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + mask_share_within_head: Whether to share mask across heads + q_block_avg_pool: Whether to apply average pooling to query blocks + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + moba_3d_mask: 3D attention mask with shape (num_heads, block_num, block_num) + """ + seq_len = q.size(2) + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + image_shape = canvas_thw + block_shape = tile_thw + batch_size, num_heads, seq_len, head_dim = k.shape + num_blocks_t = np.ceil(image_shape[0] / block_shape[0]).astype(int) + num_blocks_h = np.ceil(image_shape[1] / block_shape[1]).astype(int) + num_blocks_w = np.ceil(image_shape[2] / block_shape[2]).astype(int) + + def get_block_avg_feat(x: torch.Tensor, adaptive_pool=None, pooling_type="avg") -> torch.Tensor: + x_block_means = x.view(batch_size, num_heads, num_blocks_t, num_blocks_h, num_blocks_w, + block_shape[0], block_shape[1], block_shape[2], head_dim) + if adaptive_pool is not None: + x_block_means = x_block_means.view(-1, block_shape[0], block_shape[1], block_shape[2], head_dim) + x_block_means = x_block_means.permute(0, 4, 1, 2, 3) + if pooling_type == "avg": + global_avg = torch.nn.functional.adaptive_avg_pool3d(x_block_means, adaptive_pool) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape( + batch_size, + num_heads, + -1, + head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2] + ) + x_block_means = global_avg + elif pooling_type == "max": + max_pool = torch.nn.functional.adaptive_max_pool3d(x_block_means, adaptive_pool) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape( + batch_size, + num_heads, + -1, + head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2] + ) + x_block_means = max_pool + elif pooling_type == "mix": + global_avg = torch.nn.functional.adaptive_avg_pool3d(x_block_means, (1, 1, 1)) + global_avg = global_avg / (global_avg.norm(dim=1, keepdim=True) + 1e-8) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape(batch_size, num_heads, -1, head_dim) + + max_pool = torch.nn.functional.adaptive_max_pool3d(x_block_means, adaptive_pool) + max_pool = max_pool / (max_pool.norm(dim=1, keepdim=True) + 1e-8) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape( + batch_size, + num_heads, + -1, + head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2] + ) + x_block_means = torch.cat([max_pool, global_avg], dim=-1) + else: + raise ValueError(f"pooling_type={pooling_type} is not Supported") + else: + x_block_means = x_block_means.mean(dim=(-2, -3, -4)).view(batch_size, num_heads, -1, head_dim) + return x_block_means + + if (sampling_type == "similarity" and threshold > 0.0) and adaptive_pool is None: + adaptive_pool = (2,2,2) + + k_block_means = get_block_avg_feat(k, adaptive_pool) + + if q_block_avg_pool: + q = get_block_avg_feat(q, adaptive_pool) + q = q.type(torch.float32) + k_block_means = k_block_means.type(torch.float32) # float logit for better gate logit perception + + if mask_share_within_head: + q = q.mean(dim=1, keepdim=True) + k_block_means = k_block_means.mean(dim=1, keepdim=True) + + if sampling_type == "similarity": + top_block_indices = similarity_sampling(q, k_block_means, topk, threshold, + block_num=block_num, + adaptive_pool=adaptive_pool) + elif sampling_type == "importance": + top_block_indices = importance_sampling(q, + k_block_means, + topk, + threshold, + lambda_=lambda_, + adaptive_pool=adaptive_pool) + else: + raise NotImplementedError(f"sampling_type={sampling_type} is not Supported") + + q = q.type_as(k) + # Keep head dimension to dynamically build mask + assert top_block_indices.size(0) == 1, "top_block_indices batch size must be 1" + top_block_indices = top_block_indices.squeeze(0) # Remove batch dimension [h, s_block, topk] + + # Create 3D mask template (heads, block_num, block_num) + gate_idx_mask = torch.zeros( + (top_block_indices.size(0), block_num, block_num), + dtype=torch.bool, + device=q.device + ) + + # Iterate over each attention head using dim=0 + for head_idx in range(top_block_indices.size(0)): + # Build index matrix for each head + head_mask = torch.zeros( + (block_num, block_num), + dtype=torch.bool, + device=q.device + ) + # Fill indices for current head + head_mask.scatter_( + dim=-1, + index=top_block_indices[head_idx], + value=True + ) + gate_idx_mask[head_idx] = head_mask + + if text_block_num > 0: + pad_block_num = block_num + text_block_num + moba_3d_mask = torch.full( + (gate_idx_mask.size(0), pad_block_num, pad_block_num), + False, + dtype=torch.bool, + device=gate_idx_mask.device + ) + moba_3d_mask[:, :block_num, :block_num] = gate_idx_mask + if add_text_mask: + moba_3d_mask[:, :, -text_block_num:] = True + moba_3d_mask[:, -text_block_num:, :] = True + else: + moba_3d_mask = gate_idx_mask + + return moba_3d_mask + +@lru_cache(maxsize=4096) +def create_sta_3d_mask_optimize(canvas_thw, tile_thw, kernel_thw): + r"""Create optimized STA (Spatio-Temporal Attention) 3D mask using vectorized operations. + + Args: + canvas_thw: String representation of canvas dimensions "t_h_w" + tile_thw: String representation of tile dimensions "t_h_w" + kernel_thw: String representation of kernel dimensions "t_h_w" + + Returns: + block_mask: Boolean mask tensor with shape (block_num, block_num) + """ + canvas_thw = tuple(map(int, canvas_thw.split('_'))) + seq_len = np.prod(canvas_thw) + tile_thw = tuple(map(int, tile_thw.split('_'))) + kernel_thw = tuple(map(int, kernel_thw.split('_'))) + + kernel_t, kernel_h, kernel_w = kernel_thw + + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + block_mask = np.full((block_num, block_num), False, dtype=bool) + tile_thw_num=(canvas_thw[0] // tile_thw[0], canvas_thw[1] // tile_thw[1], canvas_thw[2] // tile_thw[2]) + + i_indices = np.arange(block_num) + j_indices = np.arange(block_num) + + i_grid, j_grid = np.meshgrid(i_indices, j_indices, indexing='ij') + + q_t_tile = i_grid // (tile_thw_num[1] * tile_thw_num[2]) + q_h_tile = (i_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[2] + q_w_tile = i_grid % tile_thw_num[2] + + kv_t_tile = j_grid // (tile_thw_num[1] * tile_thw_num[2]) + kv_h_tile = (j_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[2] + kv_w_tile = j_grid % tile_thw_num[2] + + kernel_center_t = np.clip(q_t_tile, kernel_t // 2, (tile_thw_num[0] - 1) - kernel_t // 2) + kernel_center_h = np.clip(q_h_tile, kernel_h // 2, (tile_thw_num[1] - 1) - kernel_h // 2) + kernel_center_w = np.clip(q_w_tile, kernel_w // 2, (tile_thw_num[2] - 1) - kernel_w // 2) + + time_mask = np.abs(kernel_center_t - kv_t_tile) <= kernel_t // 2 + hori_mask = np.abs(kernel_center_h - kv_h_tile) <= kernel_h // 2 + vert_mask = np.abs(kernel_center_w - kv_w_tile) <= kernel_w // 2 + + block_mask = time_mask & hori_mask & vert_mask + + block_mask = torch.tensor(block_mask, dtype=torch.bool) + return block_mask + +@torch.no_grad() +def create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num=0): + r"""Create STA (Spatio-Temporal Attention) 3D mask. + + Args: + canvas_thw: Canvas dimensions (t, h, w) + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks to pad + + Returns: + sta_mask: Boolean mask tensor with optional text block padding + """ + block_mask = create_sta_3d_mask_optimize("_".join([str(x) for x in canvas_thw]), + "_".join([str(x) for x in tile_thw]), + "_".join([str(x) for x in kernel_thw])) + sta_mask = None + block_num = block_mask.size(0) + if text_block_num > 0: + pad_block_num = block_num + text_block_num + sta_mask = torch.full( + (pad_block_num, pad_block_num), + False, + dtype=torch.bool, + ) + sta_mask[:block_num, :block_num] = block_mask + sta_mask[:, -text_block_num:] = True + sta_mask[-text_block_num:, :] = True + else: + sta_mask = block_mask + return sta_mask + +@torch.no_grad() +def create_ssta_3d_mask(q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + threshold=0.0, + lambda_=None, + text_mask=None, + mask_share_within_head=True, + adaptive_pool=None, + sampling_type=None): + r"""Create SSTA (Sparse Spatio-Temporal Attention) 3D mask combining STA and MOBA masks. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + ssta_3d_mask: Combined 3D attention mask + """ + sta_3d_mask = create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num) + sta_3d_mask = sta_3d_mask.to(q.device) + + moba_3d_mask = create_moba_3d_mask(q, k, canvas_thw, topk, tile_thw, kernel_thw, text_block_num, + threshold=threshold, + lambda_=lambda_, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type) + ssta_3d_mask = torch.logical_or(sta_3d_mask.unsqueeze(0), moba_3d_mask) + assert len(ssta_3d_mask.size()) == 3, "ssta_3d_mask should be 3D" + + # set text block mask + if text_mask is not None: + block_size = np.prod(tile_thw) + seq_len = q.size(2) + block_num = int(seq_len / block_size) + text_mask_index = torch.ceil(text_mask.sum() / block_size).long() + text_mask_index = text_mask_index.clamp(min=1).item() + assert ssta_3d_mask.shape[-1] == ssta_3d_mask.shape[-2], "ssta_3d_mask should be square" + + pad_start_index = block_num + text_mask_index + ssta_3d_mask[:, pad_start_index:, :] = False + ssta_3d_mask[:, :, pad_start_index:] = False + eye_mask = torch.eye( + ssta_3d_mask.shape[1]-pad_start_index, + dtype=torch.bool, + device=ssta_3d_mask.device + ).unsqueeze(0) + ssta_3d_mask[:, pad_start_index:, pad_start_index:] = ( + ssta_3d_mask[:, pad_start_index:, pad_start_index:] | eye_mask + ) + return ssta_3d_mask + +def ssta_3d_attention(all_q, + all_k, + all_v, + canvas_thw, + topk=1, + tile_thw=(6, 8, 8), + kernel_thw=(1, 1, 1), + text_len=0, + sparse_type='ssta', + threshold=0.0, + lambda_=None, + pad_type="zero", + text_mask=None, + mask_share_within_head=True, + sampling_type=None, + adaptive_pool=None): + r"""Sparse Spatio-Temporal Attention (SSTA) 3D attention mechanism. + + Args: + all_q: Query tensor with shape (B, H, S, D) + all_k: Key tensor with shape (B, H, S, D) + all_v: Value tensor with shape (B, H, S, D) + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_len: Length of text sequence + sparse_type: Type of sparse attention ('sta', 'block_attn', or 'ssta') + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + pad_type: Padding type ("zero" or "repeat") + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + sampling_type: Type of sampling ("similarity" or "importance") + adaptive_pool: Adaptive pooling size + + Returns: + tuple: (output tensor, sparse_ratio) + - output: Attention output with shape (B, H, S, D) + - sparse_ratio: Ratio of non-zero attention weights + """ + try: + from flex_block_attn import flex_block_attn_func + except ImportError as e: + raise Exception( + "Could not load flex-block-attn. Please install the flex-block-attn package. " + "You can install it via 'pip install flex-block-attn'." + ) from e + assert pad_type in ["zero", "repeat"] + assert sampling_type in ["similarity", "importance"] + assert (lambda_ is not None and sampling_type=="importance") or sampling_type == "similarity" + + if text_len > 0: + image_q = all_q[:, :, :-text_len, :] + image_k = all_k[:, :, :-text_len, :] + image_v = all_v[:, :, :-text_len, :] + + text_q = all_q[:, :, -text_len:, :] + text_k = all_k[:, :, -text_len:, :] + text_v = all_v[:, :, -text_len:, :] + else: + image_q = all_q + image_k = all_k + image_v = all_v + + b, hd, s, d = image_q.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} != s:{s}" + tile_t, tile_h, tile_w = tile_thw + block_size = np.prod(tile_thw) + + need_pad = False + if t % tile_t != 0 or h % tile_h !=0 or w % tile_w != 0: + need_pad = True + pad_image_q = image_q.reshape(b, hd, t, h, w, d) + pad_image_k = image_k.reshape(b, hd, t, h, w, d) + pad_image_v = image_v.reshape(b, hd, t, h, w, d) + + pad_t = 0 if t % tile_t == 0 else tile_t - t % tile_t + if pad_t > 0: + t = t + pad_t + repeat_q = pad_image_q[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + repeat_k = pad_image_k[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + repeat_v = pad_image_v[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=2) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=2) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=2) + + pad_h = 0 if h % tile_h == 0 else tile_h - h % tile_h + if pad_h > 0: + h = h + pad_h + repeat_q = pad_image_q[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + repeat_k = pad_image_k[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + repeat_v = pad_image_v[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=3) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=3) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=3) + + pad_w = 0 if w % tile_w == 0 else tile_w - w % tile_w + if pad_w > 0: + w = w + pad_w + repeat_q = pad_image_q[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + repeat_k = pad_image_k[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + repeat_v = pad_image_v[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=4) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=4) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=4) + + image_q = pad_image_q.reshape(b, hd, -1, d) + image_k = pad_image_k.reshape(b, hd, -1, d) + image_v = pad_image_v.reshape(b, hd, -1, d) + + canvas_thw = (t, h, w) + + need_pad_text = False + text_block_num = math.ceil(text_len / block_size) + text_target_size = text_block_num * block_size + if text_len % block_size > 0: + need_pad_text = True + text_pad_size = text_target_size - text_len + + pad_text_q = text_q[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + pad_text_k = text_k[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + pad_text_v = text_v[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + + text_q = torch.cat([text_q, pad_text_q], dim=2) + text_k = torch.cat([text_k, pad_text_k], dim=2) + text_v = torch.cat([text_v, pad_text_v], dim=2) + + image_q = tile(image_q, canvas_thw, tile_thw) + image_k = tile(image_k, canvas_thw, tile_thw) + image_v = tile(image_v, canvas_thw, tile_thw) + + if text_len > 0: + q = torch.cat([image_q, text_q], dim=2) + k = torch.cat([image_k, text_k], dim=2) + v = torch.cat([image_v, text_v], dim=2) + else: + q = image_q + k = image_k + v = image_v + + if sparse_type == 'sta': + assert text_mask is None, "text_mask do not support in sta sparse_type" + block_mask = create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num).to(q.device) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + elif sparse_type == 'block_attn': + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_moba_3d_mask(image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + topk=topk, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + add_text_mask=True, + lambda_=lambda_, + threshold=threshold, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + + elif sparse_type == 'ssta': + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_ssta_3d_mask( + image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + topk=topk, + threshold=threshold, + lambda_=lambda_, + text_mask=text_mask[i] if text_mask is not None else None, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type + ) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + if mask_share_within_head: + block_mask = block_mask.unsqueeze(1) # [b, 1, s_block, s_block] + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + else: + raise Exception(f"unsupported sparse_type:{sparse_type}") + sparse_ratio = block_mask.float().mean().cpu().item() + + if text_len > 0: + image_o = o[:, :, :-text_target_size, :] + if need_pad_text: + text_o = o[:, :, -text_target_size : -text_pad_size, :] + else: + text_o = o[:, :, -text_target_size:, :] + else: + image_o = o + + image_o = untile(image_o, canvas_thw, tile_thw) + + if need_pad: + # Remove padding from output + unpad_image_o = image_o.reshape(b, hd, t, h, w, d) + if pad_t > 0: + unpad_image_o = unpad_image_o[:, :, :-pad_t, :, :, :] + if pad_h > 0: + unpad_image_o = unpad_image_o[:, :, :, :-pad_h, :, :] + if pad_w > 0: + unpad_image_o = unpad_image_o[:, :, :, :, :-pad_w, :] + image_o = unpad_image_o.reshape(b, hd, -1, d) + + if text_len > 0: + o = torch.cat([image_o, text_o], dim=2) + else: + o = image_o + + return o, sparse_ratio \ No newline at end of file diff --git a/trainer/models/hyvideo/models/transformers/modules/token_refiner.py b/trainer/models/hyvideo/models/transformers/modules/token_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..69307315e6ec3f5f527f21ff7c633f0a2486ece3 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/token_refiner.py @@ -0,0 +1,283 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer +from .attention import attention + + + + +class IndividualTokenRefinerBlock(nn.Module): + """ + A single block for token refinement with self-attention and MLP. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Forward pass for IndividualTokenRefinerBlock. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + attn_mask: Optional attention mask of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + return x + + +class IndividualTokenRefiner(nn.Module): + """ + Stacks multiple IndividualTokenRefinerBlock modules. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList( + [ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + for _ in range(depth) + ] + ) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Forward pass for IndividualTokenRefiner. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + if mask is not None: + mask = mask.clone().bool() + mask[:, 0] = True # Prevent attention weights from becoming NaN + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """ + Single token refiner block for LLM text embedding refinement. + + Args: + in_channels: Input feature dimension. + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + in_channels: int, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs) + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + """ + Forward pass for SingleTokenRefiner. + + Args: + x: Input tensor of shape [B, L, in_channels]. + t: Timestep tensor of shape [B]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, hidden_size]. + """ + timestep_aware_representations = self.t_embedder(t) + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [B, L, 1] + context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder(context_aware_representations) + c = timestep_aware_representations + context_aware_representations + x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/trainer/models/hyvideo/models/transformers/modules/upsample.py b/trainer/models/hyvideo/models/transformers/modules/upsample.py new file mode 100644 index 0000000000000000000000000000000000000000..8cca495660d4a45b7a3b608a9fe937639c6caab2 --- /dev/null +++ b/trainer/models/hyvideo/models/transformers/modules/upsample.py @@ -0,0 +1,169 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from trainer.models.hyvideo.models.autoencoders.hunyuanvideo_15_vae_w_cache import ( + CausalConv3d, + ResnetBlock, + RMS_norm, + forward_with_checkpointing, + swish, +) + + +class UpsamplerType(Enum): + LEARNED = "learned" + FIXED = "fixed" + NONE = "none" + LEARNED_FIXED = "learned_fixed" + + +@dataclass +class UpsamplerConfig: + load_from: str + enable: bool = False + hidden_channels: int = 128 + num_blocks: int = 16 + model_type: UpsamplerType = UpsamplerType.NONE + version: str = "720p" + + +class SRResidualCausalBlock3D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.block = nn.Sequential( + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.block(x) + + +class SRTo720pUpsampler(ModelMixin, ConfigMixin): + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + hidden_channels: int | None = None, + num_blocks: int = 6, + global_residual: bool = False, + ): + super().__init__() + if hidden_channels is None: + hidden_channels = 64 + self.in_conv = CausalConv3d(in_channels, hidden_channels, kernel_size=3) + self.blocks = nn.ModuleList([SRResidualCausalBlock3D(hidden_channels) for _ in range(num_blocks)]) + self.out_conv = CausalConv3d(hidden_channels, out_channels, kernel_size=3) + self.global_residual = bool(global_residual) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + y = self.in_conv(x) + for blk in self.blocks: + y = blk(y) + y = self.out_conv(y) + if self.global_residual and (y.shape == residual.shape): + y = y + residual + return y + + +class SRTo1080pUpsampler(ModelMixin, ConfigMixin): + + @register_to_config + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int = 2, + is_residual: bool = False, + ): + super().__init__() + self.num_res_blocks = num_res_blocks + self.block_out_channels = block_out_channels + self.z_channels = z_channels + + block_in = block_out_channels[0] + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + self.up.append(up) + + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, out_channels, kernel_size=3) + + self.gradient_checkpointing = False + self.is_residual = is_residual + + def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor: + """ + Args: + z: (B, C, T, H, W) + target_shape: (H, W) + """ + use_checkpointing = bool(self.training and self.gradient_checkpointing) + if target_shape is not None and z.shape[-2:] != target_shape: + bsz = z.shape[0] + z = rearrange(z, "b c f h w -> (b f) c h w") + z = F.interpolate(z, size=target_shape, mode="bilinear", align_corners=False) + z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing( + self.up[i_level].block[i_block], + h, + use_checkpointing=use_checkpointing, + ) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h diff --git a/trainer/models/hyvideo/utils/communications.py b/trainer/models/hyvideo/utils/communications.py new file mode 100644 index 0000000000000000000000000000000000000000..9d63ccae0eaecaa42b3c2610d887fa41a8fc2ddf --- /dev/null +++ b/trainer/models/hyvideo/utils/communications.py @@ -0,0 +1,294 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Any, Tuple +import torch +import torch.distributed as dist + +from torch.nn import functional as F + + +def broadcast(input_: torch.Tensor, group: dist.ProcessGroup): + src = dist.get_global_rank(group, 0) + dist.broadcast(input_, src=src, group=group) + + +def _all_to_all_4D( + input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None +) -> torch.tensor: + """ + all-to-all for QKV + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert ( + input.dim() == 4 + ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + + seq_lens = [None] * seq_world_size + dist.all_gather_object(seq_lens, input.shape[1], group) + # uneven + if seq_lens[-1] != seq_lens[0] : + assert seq_lens[0] > seq_lens[-1] + gap = seq_lens[0] - seq_lens[-1] + if dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1: + assert input.shape[1] == seq_lens[-1] + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + else: + gap = 0 + + + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + assert hc % seq_world_size == 0, f'Invalid Head size: {hc}, which should be divisible by spsize {seq_world_size}' + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = ( + input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs) + .transpose(0, 2) + .contiguous() + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + if gap > 0: + output = output[:, :-gap] + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + + hc = shard_hc * seq_world_size + if seqlen % seq_world_size != 0: + new_seqlen = (seqlen // seq_world_size + 1) * seq_world_size + gap = new_seqlen - seqlen + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + bs, seqlen, shard_hc, hs = input.shape + else: + gap = 0 + + + assert seqlen % seq_world_size == 0 + + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = ( + input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs) + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + + if gap > 0 and dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1: + output = output[:, :-gap] + + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + +class SeqAllToAll4D(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: torch.Tensor, + scatter_idx: int, + gather_idx: int, + ) -> torch.Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply( + ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx + ), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int = 2, gather_dim: int = 1, +): + return SeqAllToAll4D.apply(group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [ + t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim) + ] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all( + input_, ctx.world_size, process_group, scatter_dim, gather_dim + ) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + + +def all_to_all( + input_: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int = 2, gather_dim: int = 1 +): + return _AllToAll.apply(input_, group, scatter_dim, gather_dim) + + +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim, group): + ctx.dim = dim + ctx.group = group + world_size = dist.get_world_size(group) + input_size = list(input_.size()) + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_.shape, group) + + ctx.input_size = input_size[dim] + + tensor_list = [torch.empty(sizes[i], dtype=input_.dtype, device=input_.device) for i in range(world_size)] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + group = ctx.group + world_size = dist.get_world_size(group) + global_rank = dist.get_rank() + rank = dist.get_group_rank(group, global_rank) + dim = ctx.dim + input_size = ctx.input_size + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_size, group=group) + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None, None + + +def all_gather(input_: torch.Tensor, dim: int = 1, group=None): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim, group) diff --git a/trainer/models/hyvideo/utils/data_utils.py b/trainer/models/hyvideo/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6677d592ccccd3e46cd1c0f7ef95cb7930c3e96c --- /dev/null +++ b/trainer/models/hyvideo/utils/data_utils.py @@ -0,0 +1,73 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import numpy as np +from PIL import Image + +def resize_and_center_crop(image, target_width, target_height): + if target_height == image.shape[0] and target_width == image.shape[1]: + return image + + pil_image = Image.fromarray(image) + original_width, original_height = pil_image.size + scale_factor = max(target_width / original_width, target_height / original_height) + resized_width = int(round(original_width * scale_factor)) + resized_height = int(round(original_height * scale_factor)) + resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS) + left = (resized_width - target_width) / 2 + top = (resized_height - target_height) / 2 + right = (resized_width + target_width) / 2 + bottom = (resized_height + target_height) / 2 + cropped_image = resized_image.crop((left, top, right, bottom)) + return np.array(cropped_image) + + +def get_closest_ratio(height: float, width: float, ratios: list, buckets: list): + """ + Get the closest ratio in the buckets. + + Args: + height (float): video height + width (float): video width + ratios (list): video aspect ratio + buckets (list): buckets generated by `generate_crop_size_list` + + Returns: + the closest size in the buckets and the corresponding ratio + """ + aspect_ratio = float(height) / float(width) + + ratios_array = np.array(ratios) + closest_ratio_id = np.abs(ratios_array - aspect_ratio).argmin() + closest_size = buckets[closest_ratio_id] + closest_ratio = ratios_array[closest_ratio_id] + + return closest_size, closest_ratio + + +def generate_crop_size_list(base_size=256, patch_size=16, max_ratio=4.0): + num_patches = round((base_size / patch_size) ** 2) + assert max_ratio >= 1.0 + crop_size_list = [] + wp, hp = num_patches, 1 + while wp > 0: + if max(wp, hp) / min(wp, hp) <= max_ratio: + crop_size_list.append((wp * patch_size, hp * patch_size)) + if (hp + 1) * wp <= num_patches: + hp += 1 + else: + wp -= 1 + return crop_size_list diff --git a/trainer/models/hyvideo/utils/flash_attn_no_pad.py b/trainer/models/hyvideo/utils/flash_attn_no_pad.py new file mode 100644 index 0000000000000000000000000000000000000000..06db6f1335018c4f5adbf9734fd3035e7cfa43a4 --- /dev/null +++ b/trainer/models/hyvideo/utils/flash_attn_no_pad.py @@ -0,0 +1,173 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from einops import rearrange +from torch.nn.attention.flex_attention import flex_attention, create_block_mask, BlockMask + +# compile first +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune") + +def prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 880, causal_mask=None + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + + total_length = num_frames * frame_seqlen + + def attention_mask(b, h, q_idx, kv_idx): + return causal_mask[q_idx, kv_idx] + +# add new function for computing causal attention with flex attention +def flex_attn_no_pad( + qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + batch_size = qkv.shape[0] # qkv shape: [B, total_length, 3, num_head, D] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( + x, key_padding_mask + ) + + # ------------------------- for chunk-wise causal attention mask ------------------------- + # get the video sequence + latent_seq_length = 1560 # set for hunyuanvideo 1.5, which is for 480 * 832 resolution + chunk_seq_length = 1560 * 4 + text_seq_length = max_s % chunk_seq_length # including txt, byt5, sigclip vision token + latent_num = (max_s - text_seq_length) // latent_seq_length + chunk_num = (max_s - text_seq_length) // chunk_seq_length + causal_mask = torch.zeros((max_s, max_s), device=qkv.device) + + causal_mask[:text_seq_length, :text_seq_length] = 1 # no attention for the rest + for i in range(chunk_num): + start_i = text_seq_length + i * chunk_seq_length + end_i = min(start_i + chunk_seq_length, max_s) + for j in range(i + 1): + start_j = text_seq_length + j * chunk_seq_length + end_j = min(start_j + chunk_seq_length, max_s) + # full attention within chunk i for j == i, causal for j < i + causal_mask[start_i:end_i, start_j:end_j] = 1 + causal_mask = causal_mask.to(torch.bool) # Force bool dtype + + block_mask = prepare_blockwise_causal_attn_mask( + device=qkv.device, + num_frames=latent_num, + frame_seqlen=latent_seq_length, + causal_mask=causal_mask + ) + + q,k,v = qkv.chunk(3, dim=2) + q = q.squeeze(2).transpose(1, 2) + k = k.squeeze(2).transpose(1, 2) + v = v.squeeze(2).transpose(1, 2) + output_unpad = flex_attention( + q, k, v, block_mask=block_mask + ) # output_unpad: [B, num_head, L, D] + + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output + + +# add new function for computing causal attention with flex attention +def flash_attn_no_pad( + qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( + x, key_padding_mask + ) + + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output + + +def flash_attn_no_pad_v3( + qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3 + + if flash_attn_varlen_func_v3 is None: + raise ImportError("FlashAttention V3 backend not available") + + batch_size, seqlen, _, nheads, head_dim = qkv.shape + query, key, value = qkv.unbind(dim=2) + + query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + rearrange(query, "b s h d -> b s (h d)"), key_padding_mask + ) + key_unpad, _, cu_seqlens_k, _, _ = unpad_input( + rearrange(key, "b s h d -> b s (h d)"), key_padding_mask + ) + value_unpad, _, _, _, _ = unpad_input( + rearrange(value, "b s h d -> b s (h d)"), key_padding_mask + ) + + query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads) + key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads) + value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads) + + output_unpad = flash_attn_varlen_func_v3( + query_unpad, key_unpad, value_unpad, + cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_q, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic + ) + + output = rearrange( + pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen), + "b s (h d) -> b s h d", h=nheads + ) + return output diff --git a/trainer/models/hyvideo/utils/infer_utils.py b/trainer/models/hyvideo/utils/infer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b1975e362d19110420297a5d0a549eaa3b060fbd --- /dev/null +++ b/trainer/models/hyvideo/utils/infer_utils.py @@ -0,0 +1,31 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch + +from trainer.models.hyvideo.commons.infer_state import get_infer_state + +def torch_compile_wrapper(): + """返回一个装饰器，延迟决定是否使用torch.compile""" + def decorator(func): + def wrapper(*args, **kwargs): + if get_infer_state() and get_infer_state().enable_torch_compile: + compiled_func = torch.compile(func) + return compiled_func(*args, **kwargs) + else: + return func(*args, **kwargs) + return wrapper + return decorator diff --git a/trainer/models/hyvideo/utils/multitask_utils.py b/trainer/models/hyvideo/utils/multitask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eeff5814c06d5f1106dbc306841430e8bada8b5a --- /dev/null +++ b/trainer/models/hyvideo/utils/multitask_utils.py @@ -0,0 +1,56 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from typing import List +import numpy as np +from PIL import Image + + +def numpy_to_pil(images: np.ndarray) -> List[Image.Image]: + """ + Convert a numpy image or a batch of images to a PIL image. + + Args: + images (np.ndarray): The image array to convert to PIL format. + + Returns: + List[Image.Image]: A list of PIL images. + """ + if images.ndim == 3: + images = images[None, ...] + images = (images * 255).round().astype("uint8") + if images.shape[-1] == 1: + # special case for grayscale (single channel) images + pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] + else: + pil_images = [Image.fromarray(image) for image in images] + + return pil_images + + +def merge_tensor_by_mask(tensor_1, tensor_2, mask, dim): + assert tensor_1.shape == tensor_2.shape + # Mask is a 0/1 vector. Choose tensor_2 when the value is 1; otherwise, tensor_1 + masked_indices = torch.nonzero(mask).squeeze(1) + tmp = tensor_1.clone() + if dim == 0: + tmp[masked_indices] = tensor_2[masked_indices] + elif dim == 1: + tmp[:, masked_indices] = tensor_2[:, masked_indices] + elif dim == 2: + tmp[:, :, masked_indices] = tensor_2[:, :, masked_indices] + return tmp diff --git a/trainer/models/hyvideo/utils/rewrite/clients.py b/trainer/models/hyvideo/utils/rewrite/clients.py new file mode 100644 index 0000000000000000000000000000000000000000..bbcd0f1964766ea96f49bb00136b04405d0776f2 --- /dev/null +++ b/trainer/models/hyvideo/utils/rewrite/clients.py @@ -0,0 +1,257 @@ +# -*- coding: utf-8 -*- +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json +import time +import ast +import os +import io +import base64 + +import openai +from loguru import logger + + +class NonStreamResponse(object): + def __init__(self): + self.response = "" + + def _deserialize(self, obj): + self.response = json.dumps(obj) + +# DeepSeekClient +class DeepSeekClient(object): + def __init__(self, key_id, key_secret): + from tencentcloud.common.common_client import CommonClient + from tencentcloud.common import credential + from tencentcloud.common.profile.client_profile import ClientProfile + from tencentcloud.common.profile.http_profile import HttpProfile + cred = credential.Credential(key_id, key_secret) + httpProfile = HttpProfile() + httpProfile.endpoint = "lkeap.tencentcloudapi.com" + httpProfile.reqTimeout = 40000 # The streaming interface may take a longer time. + clientProfile = ClientProfile() + clientProfile.httpProfile = httpProfile + self.common_client = CommonClient("lkeap", "2024-05-22", cred, "ap-guangzhou", profile=clientProfile) + + def run_single_recaption(self, system_prompt, input_prompt): + post_dict = { + "Model": "deepseek-v3.1", + "Messages": [ + { + "Role": "system", + "Content": system_prompt + }, + { + "Role": "user", + "Content": input_prompt + } + ], + "Stream": False, + "Thinking": {"Type": "enabled"}, + } + while True: + try: + resp = self.common_client._call_and_deserialize("ChatCompletions", post_dict, NonStreamResponse) + break + except Exception as e: + logger.error(e) + time.sleep(1) + resp = self.common_client._call_and_deserialize("ChatCompletions", post_dict, NonStreamResponse) + response = resp.response + response = ast.literal_eval(response) + content = response["Choices"][0]["Message"]["Content"] + reasoning_content = response["Choices"][0]["Message"]["ReasoningContent"] + print('Initial prompt: ', input_prompt) + print('Recaption prompt: ', content) + + return content, reasoning_content + +class QwenClient(object): + def __init__(self, base_url=None, model_name=None): + # Recommended model: Qwen3-235B-A22B-Thinking-2507 + self.base_url = base_url + self.model_name = model_name + + def qwen_api_call(self, system_prompt: str, user_input: str, temperature: float, max_tokens: int): + """ + Use Qwen Chat API to perform text rewriting, parse ... sections for reasoning content, and return (thinking, result). + """ + client = openai.OpenAI(base_url=self.base_url, api_key="None", timeout=600) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_input}, + ] + + last_err = None + + for i in range(10): + try: + stream = False + response = client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + stream=stream, + ) + content = response.choices[0].message.content + thinking = "" + result = content or "" + if content and "" in content: + head, tail = content.split("", 1) + thinking = head.replace("", "").strip() + result = tail.strip() + return thinking, result + except Exception as e: + last_err = e + if i < 9: + time.sleep(2 ** i) + continue + raise last_err + + + def run_single_recaption(self, system_prompt, input_prompt, temperature=0.1, max_tokens=4096): + thinking, result = self.qwen_api_call(system_prompt, input_prompt, temperature, max_tokens) + return result + + +class QwenVLClient(object): + + def __init__(self, base_url=None, model_name=None): + self.base_url = base_url + self.model_name = model_name + self.api_key = "None" + self.max_image_size = int(os.getenv("I2V_REWRITE_MAX_IMAGE_SIZE", "1024")) # 控制送入模型的最大图像分辨率 + + def _encode_image_to_base64(self, image_path: str, max_dimension: int) -> str: + """ + 参考 hyvideo/utils/rewrite/qwen_vllm.py 的实现： + 加载本地图片，将其按比例缩放到 max_dimension，然后编码为 Base64 data URL。 + """ + try: + from PIL import Image + except ImportError as e: + logger.error("Pillow (PIL) is required for QwenVLClient image encoding but is not installed.") + raise e + + try: + with Image.open(image_path) as img: + if img.width > max_dimension or img.height > max_dimension: + img.thumbnail((max_dimension, max_dimension)) + + buffer = io.BytesIO() + + # 去除透明通道，统一转为 JPEG + if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info): + img = img.convert("RGB") + + img.save(buffer, format="JPEG") + mime_type = "image/jpeg" + + encoded_string = base64.b64encode(buffer.getvalue()).decode("utf-8") + return f"data:{mime_type};base64,{encoded_string}" + except Exception as e: + logger.error(f"Failed to encode image {image_path} to base64: {e}") + raise + + def qwen_api_call( + self, + system_prompt: str, + user_input: str, + temperature: float, + max_tokens: int, + img_path: str = None, + ): + """ + Use Qwen3-VL to perform text rewriting. + """ + client = openai.OpenAI( + base_url=self.base_url, + api_key=self.api_key, + timeout=3600, + ) + + assert "{}" in system_prompt, "system_prompt must contain {{}}" + prompt_text = system_prompt.format(user_input) + + + assert img_path is not None, "img_path is required" + base64_image = self._encode_image_to_base64(img_path, max_dimension=self.max_image_size) + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": base64_image, + }, + }, + { + "type": "text", + "text": prompt_text, + }, + ], + } + ] + + last_err = None + + for i in range(10): + try: + + response = client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + stream=False, + ) + content = response.choices[0].message.content + thinking = "" + result = content or "" + # 和 QwenClient 一样，兼容 思维模式输出 + if content and "" in content: + head, tail = content.split("", 1) + thinking = head.replace("", "").strip() + result = tail.strip() + return thinking, result + except Exception as e: + last_err = e + logger.error(f"QwenVLClient request failed (attempt {i + 1}/10): {e}") + if i < 9: + time.sleep(2) + else: + raise last_err + + def run_single_recaption( + self, + system_prompt, + input_prompt, + temperature=0.1, + max_tokens=4096, + img_path: str = None, + ): + thinking, result = self.qwen_api_call( + system_prompt, + input_prompt, + temperature, + max_tokens, + img_path=img_path, + ) + return result \ No newline at end of file diff --git a/trainer/models/hyvideo/utils/rewrite/i2v_prompt.py b/trainer/models/hyvideo/utils/rewrite/i2v_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..e52016ea22c9f79705883232ebd5b30da06646a4 --- /dev/null +++ b/trainer/models/hyvideo/utils/rewrite/i2v_prompt.py @@ -0,0 +1,118 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + + +i2v_rewrite_system_prompt = """ +## 角色 +你是一位顶级的图生视频（Image-to-Video）Prompt工程师。你的任务不是生成视频，而是将用户输入的自由形式的自然语言，改写成具有丰富视觉细节、精确动态描述、并采用专业影视语言的中文Prompt。改写后的措辞、句式、表达方式等，必须严格遵循并尽可能接近本指令中定义的语言风格和表达习惯。 + +## 任务 +你的核心任务是进行"文本改写"。接收用户的简短或模糊想法，输出一段符合以下规则的、详细、客观、可执行的中文视频脚本式描述。无论用户输入是中文还是英文，你的输出都必须是中文。 + +## 核心改写规则 + +### 1. 镜头语言标准化 (Camera Language Standardization) +当用户指令中包含镜头运镜描述时，尽可能转换为标准表述。如果不能完全对应到标准表述，保留原意。 +* **运镜标准表述**: `镜头缓缓拉远/后拉`, `镜头向前推进`, `镜头上/下/左/右移动`, `镜头摇动/摇移`, `镜头跟随`, `镜头环绕`, `镜头静止不动`, `手持镜头`。 +* **示例1**: （能对应标准）用户输入"镜头慢慢前推，跟踪小鸟飞行"，应改写为"**镜头向前推进**，跟踪小鸟飞行"。 +* **示例2**: （不能完全对应）用户输入"镜头缓缓顺时针旋转并向前推"，应改写为"镜头缓缓顺时针旋转，并**向前推进**"。 +* 改写时绝对禁止补充用户指令中没有明确提出的镜头运镜描述，除非为了解释或标准化用户已明确提到的运镜所必需；不得擅自联想新增任何镜头运动方式或运镜效果。特别地，在用户未明确说明“镜头静止”或具有同等含义的表述时，严禁擅自添加“镜头静止不动”等描述。 + +### 2. 动态化与时序性 (Dynamic & Sequential) +将用户的静态描述分解成一个微小的时间序列。使用连接词来串联连续发生的、或同时发生的动作，构建出清晰的叙事流。 +* **结构**: 动作A发生，**随后/然后**，动作B发生，**同时**，动作C发生。 +* **常用连接词**: `随后`, `然后`, `接着`, `同时`, `之后`。 +* **示例1**: 用户输入"两个人见面"，应改写为"...画面左侧走进来一个男子，画面右侧走进来一个女子，他们微笑着在中间的心形前停下，**随后**两个人双手握在一起对视，**接着**男子和女子接吻..."。 +* **示例2**: 用户输入"女孩跳舞"，应改写为"**女孩身体开始左右轻轻摇摆,同时双手缓缓举过头顶**。。。"。 + +### 3. 遵循“主体-动作-细节”的客观描述模式 +使用客观、中立的语言，像摄影师一样记录画面中发生的一切。避免使用主观或情感化的词语（如“美丽的”、“悲伤的”），而是通过描述具体的行为来暗示情感（如用“嘴角露出微笑”代替“开心的”）。 +* **句式结构**: `[主体] + [方式状语(如缓缓地)] + [动作(如转动头部)]`。 +* **示例**: "**黄色头的蜥蜴** **转动着头** **向前探出身子**。" 这是一个完美的客观描述链。 + +### 4. 空间与方位的精确化 +明确物体和人物在画面中的位置及其移动方向,如果有新加入物体，需要描述已有物体加以区分。 +* **方位词**: `画面左/右/上/下侧`, `从...上方/下方伸入`, `背景中`, `前景处`, `向...方向移动`。 +* **示例1**: 用户输入"有只手伸进来"，应改写为"**一只手从画面右侧伸出**，摸了摸黑色衣服上的吊牌...**从下方消失在画面中**"。 +* **示例2**: 参考图像中间有一条美人鱼，用户输入"旁边游来两条粉色鱼尾的美人鱼"，应改写为"**画面中间有一条美人鱼，两条粉色鱼尾的美人鱼从画面右侧游入**..."。 + + +### 5. 指代关系清晰 + *模糊指代必须显式化（针对用户输入）**：当用户输入中出现模糊指代（如“他们”“它”“这个/那个”“其”等），在改写中必须替换为明确的实体称谓，并最小合理补充主体类别、数量与性别，但不引入与原意无关的新实体。 + * **示例**： + * 文本指令“他们跳舞。” → 你的输出中写为“一对男女在舞池中跳着华尔兹...” + * 文本指令“把它拿起来。” → 你的输出中写为“一只手从画面左侧伸入，拿起绿色的洗面奶管子...” + * 文本指令“把炸弹递给他。” → 你的输出中写为“黑猫把手中的炸弹递给灰猫。” + +### 6. 过度克制联想 +在不改变用户意图与叙事目标的前提下，仅进行必要的镜头化与轻度细节扩展。 +* 不得新增用户未要求的关键事件或动作（如坐下、拿起、递交、拥抱等）。 +* 可适度补充环境"微动态"，但不得擅自改变主体、方向、数量与时序。 +* **避免补充光影相关描述**：，除非用户指令明确要求,绝对不要添加任何光影相关内容，包括但不限于"光线"、"阴影"、"光影斑驳"、"阳光照射"、"灯光变化"等。 + 同时，绝对严禁输出"无任何光影变化"、"无额外光影渲染"、"整体画面无额外光影变化"等强调不存在光影的表述，否则将视为严重错误！！！ +* 输出长度不必与样例等长，应与用户输入的信息量匹配；避免冗长与过度联想。 + +### 7. 对齐用户意图 +* 结合参考图像和用户指令，正确理解用户的指定的范畴 + - 例1： 参考图像是：三个漂浮在水面上的甜甜圈 + 用户指令是：“美味的甜甜圈从水上慢慢沉下去消失不见，水波浮动” + 不应该改写成："画面中央的巧克力甜甜圈缓缓下沉...旁边的粉色和白色甜甜圈仍在水上轻轻浮动..." + 应该改写成："画面中的三个甜甜圈缓缓下沉..."（因为这里用户指代的是参考图像当中所有的甜甜圈） + - 例2: 参考图像是：狗人、兔子人、猫人三个人形动物头的角色在故宫的背景，面向镜头奔跑 + 用户指令是：兔子人加速向前跑，然后剩下的两个追上去，镜头右摇，拍摄它们跑远宫殿的背影 + 不应该改写成：画面中，一只狗、一只白兔和一只猫正并排在石桥上向前奔跑。突然，中间的白兔加速向前冲刺。随后，它左侧的狗和右侧的猫也加快速度追了上去。镜头向右摇移，跟随拍摄它们跑向远处宫殿的背影。 + 应该改写成：画面中，一只狗、一只白兔和一只猫正并排在石桥上向前奔跑。突然，中间的白兔加速向前冲刺。随后，它左侧的狗和右侧的猫也加快速度追了上去。镜头向右摇移，跟随拍摄它们奔跑远离宫殿的背影。 + + +## 文字处理规则 +如果用户输入包含引用文字（如书名、品牌名、标语等），必须在改写后的Prompt中保持原文和原语言，并一定要用中文双引号包围。不要在没有原文引用的情况下自行添加引用文字。例如，如果用户提到 写着"GROOVY MANGO"，你必须改写为 写着“GROOVY MANGO”。 +注意：绝对不要在指令中让完整的一段文字，逐个字词的出现，例如： + 用户指令：男人在纸上写"iPhone 15" + 不要改写成：...依次写上“i”、“P”、“h”.../依次形成字母“i”、“P”、“h”.../按顺序写上文字“i”、“P”、“h”... （等等任何把文字拆开的描述方式） + 应该改写成：...写上文字“iPhone 15” + +## 改写范例 +以下是一些改写样例，请严格模仿“改写后”的风格。 + +**范例 1** +* **参考图像**：傍晚的海边沙滩，一位长发女子站在近景位置，身体微侧，远处海浪与天际线清晰可见，海风将发丝轻轻吹起。 +* **文本指令**: 在海边的沙滩上站着的一位女人，面朝大海凝望。 +* **你的输出**: 女人的发丝随风轻轻飘动。女人眨了一下眼，向画面右侧转头凝视前方。背景中海浪拍打着沙滩。 + +**范例 2** +* **参考图像**：户外阴天光线下，一朵花的近景特写，花瓣边缘附着细小水珠，背景为柔和的绿色叶片虚化。 +* **文本指令**: 一朵在小雨中被细密雨丝打湿的花朵，花瓣上挂着水珠。 +* **你的输出**: 雨水滴落在花朵上，水珠在花瓣间滚动，花瓣轻微摇动。 + +**范例 3** +* **参考图像**：室内柔和日光下，一个圆形绒面猫窝里蜷着一只橘猫，窗边可见绿色树叶的模糊影子落在地面上。 +* **文本指令**: The cat stretches its body and wags its tail. Camera pulls back. The cat opens and blinks its eyes. +* **你的输出**: 橘猫在柔软的猫窝中舒展身体，把头和一只爪子伸出猫窝外，不停摇晃着尾巴。镜头微微向后拉远，露出窗外更多的树叶景色。猫咪睁开眼睛看了一下前方，随后又眯上了眼睛，尾巴也渐渐停止摇晃。阳光照下的影子也随之动着。 + +**范例 4** +* **参考图像**：灯光温暖的舞池中央，一对男女面对面站立，男士穿深色西装，女士穿深色礼服，周围有零散观众与舞台乐队背景虚化。 +* **文本指令**: A man and a woman are dancing together as the camera follows their movements. +* **你的输出**: 一对男女在舞池中跳着华尔兹，男士穿着黑色西装轻握穿着深色礼服女士的手，扶着她的腰，女子一只手搭在他的肩上。两人眼神交流，面带微笑。女士裙摆随舞步轻轻摆动。镜头跟随拍摄，捕捉他们的动作和表情。后面的人们拍着手，舞台上的乐队在演奏。 + + +## 输出格式要求 +请注意： +1. 无论用户输入是中文还是英文，你的最终输出都必须是中文。 +2. 在你的最终输出中，不要包含任何星号（**）符号。上述示例中的星号仅用于强调重要概念，你的实际输出应该是干净的、不带任何格式标记的纯文本。 + +现在给定参考图像和文本指令: {} +请按照上述规则进行改写，输出改写后的文本: +""" \ No newline at end of file diff --git a/trainer/models/hyvideo/utils/rewrite/rewrite_utils.py b/trainer/models/hyvideo/utils/rewrite/rewrite_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8b52df487cd5bca7b66404abb7184d186c0be92a --- /dev/null +++ b/trainer/models/hyvideo/utils/rewrite/rewrite_utils.py @@ -0,0 +1,69 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from hyvideo.utils.rewrite.clients import QwenClient, QwenVLClient +from hyvideo.utils.rewrite.t2v_prompt import t2v_rewrite_system_prompt +from hyvideo.utils.rewrite.i2v_prompt import i2v_rewrite_system_prompt + +def t2v_rewrite(user_prompt, rewrite_client=None): + base_url = os.getenv("T2V_REWRITE_BASE_URL") + model_name = os.getenv("T2V_REWRITE_MODEL_NAME") + if not base_url or not model_name: + raise EnvironmentError( + "T2V_REWRITE_BASE_URL and T2V_REWRITE_MODEL_NAME must be set in the environment variables " + "when prompt rewriting is enabled. Please configure the rewrite service correctly." + ) + if rewrite_client is None: + rewrite_client = QwenClient(base_url, model_name) + try: + rewritten_prompt = rewrite_client.run_single_recaption(t2v_rewrite_system_prompt, user_prompt) + except Exception as e: + raise ValueError( + f"Failed to rewrite prompt using {type(rewrite_client).__name__}: {e}" + ) + return rewritten_prompt + +def i2v_rewrite(user_input, img_path, rewrite_client=None): + """ + Use a rewrite client to generate a rewritten prompt for image-to-video. + """ + i2v_base_url = os.getenv("I2V_REWRITE_BASE_URL") + i2v_model_name = os.getenv("I2V_REWRITE_MODEL_NAME") + if not i2v_base_url or not i2v_model_name: + raise EnvironmentError( + "I2V_REWRITE_BASE_URL and I2V_REWRITE_MODEL_NAME must be set in the environment variables " + "for image-to-video prompt rewriting. Please set them before running, e.g.:\n" + 'export I2V_REWRITE_BASE_URL="YOUR_I2V_REWRITE_BASE_URL"\n' + 'export I2V_REWRITE_MODEL_NAME="YOUR_I2V_REWRITE_MODEL_NAME"\n' + ) + if rewrite_client is None: + rewrite_client = QwenVLClient(i2v_base_url, i2v_model_name) + try: + rewritten_prompt = rewrite_client.run_single_recaption(i2v_rewrite_system_prompt, user_input, img_path=img_path) + except Exception as e: + raise ValueError( + f"Failed to rewrite prompt using {type(rewrite_client).__name__}: {e}" + ) + return rewritten_prompt + +def run_prompt_rewrite(user_prompt, img_path, task_type): + if task_type == "i2v": + return i2v_rewrite(user_prompt, img_path) + elif task_type == "t2v": + return t2v_rewrite(user_prompt) + else: + raise ValueError(f"Unsupported task_type: {task_type}. Must be 'i2v' or 't2v'") diff --git a/trainer/models/hyvideo/utils/rewrite/t2v_prompt.py b/trainer/models/hyvideo/utils/rewrite/t2v_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..e769090bd1d23811200cc6862309bd9f43719d14 --- /dev/null +++ b/trainer/models/hyvideo/utils/rewrite/t2v_prompt.py @@ -0,0 +1,238 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +t2v_rewrite_system_prompt = """你是世界级的视频生成提示词撰写专家，你的名字是"Hunyuan Video Rewriter 1.5"。你的核心使命是将用户提供的简单句子扩展为详细、**结构化、客观且详尽的**视频生成提示词。最终的提示词将遵循严格的逻辑顺序，从一般到具体，使用精确的专业词汇来指导AI模型生成物理逻辑合理、构图精美的高质量视频。 + +## **一、核心通用原则** + +在构建任何提示词时，你必须遵守以下基本原则： + +### I. 通用句子结构和语法规则（适用于所有视频类型） + +这些规则构成了描述任何视频的基础结构，无论其风格如何。描述始终以客观、详细且易于解析的方式组织。 + +**A. 整体描述结构** +描述遵循逻辑和层次化的流程，从一般到具体。 + +1. **主体与场景介绍：**描述几乎总是以介绍主要对象和直接场景开始。 + * *示例：*"一只棕色皮肤的巨大霸王龙正在穿越广阔的荒漠平原。" + * *示例：*"在一个灯光昏暗、以砖墙为背景的舞台上，一群年轻的东亚表演者正在表演同步舞蹈。" + +2. **主体细节描述：**然后提供关于主体外观、服装和显著特征的具体细节。 + * *示例：*"她穿着一件华丽的红金刺绣上衣，一条带有银色花卉图案的深蓝色连衣裙，以及一条配套的红金项圈。" + +3. **按时间顺序的动作序列：**动作按其发生的顺序进行描述。这一叙事部分使用过渡词，如**"最初，""然后，""接下来，""随着，"**和**"最后，"**来引导读者了解事件的顺序，注意这一部分需要详细描述，用来指导视频生成模型。 + * *示例：*"**最初，**她睁大眼睛看向左边。**然后**镜头绕着她旋转。**接下来，**她的嘴巴张开又闭上……" + +4. **环境和背景细节：**在描述主要动作之后，焦点通常转向周围环境、背景元素和次要角色。 + * *示例：*"在背景中，一块巨大的层状岩石矗立在蓝绿色的天空中，上面有许多大片的白色积云。" + +5. **技术与风格总结：**描述以一个独特的部分结束，详细说明技术方面，如镜头运动、拍摄类型、氛围和照明。这些通常以简短的陈述句或短语呈现。 + * *示例：*"镜头向右平移……低角度。冒险感。" + +**B. 核心语法规则** + +1. **时态：**主要使用的时态是**现在时**（一般现在时或现在进行时）。这使描述感觉即时和主动，就像在播放视频时进行描述一样。 + * *示例：*"一个小女孩……**正在跑**……"，"他**穿着**黑色运动背心……"，"镜头**跟随**她……" + +2. **句子结构：** + * 句子主要是**陈述性的**，陈述关于场景的事实。 + * 结构通常以主语开始，然后是动词：`[主语] + [动词] + [细节]`。 + * **介词短语**被广泛用于添加关于位置（`在土路上`）、外观（`有着苍白的皮肤`）和关系（`在瓢虫后面`）的细节。 + * **分词短语**经常用于简洁地组合动作或描述。 + * *示例：*"一位年轻女性……**戴着棕色帽子**，从马上下来，**微笑着等待**。" + +3. **词汇和语气：** + * **形容词：**语言中充满描述性形容词，用于指定颜色、大小、纹理、情感和外观（例如，"广阔的"，"卷曲的"，"同步的"，"充满活力的"）。 + * **动词：**动作动词精确而动态（例如，"跳跃"，"投掷"，"摇摆"，"攀爬"）。 + * **语气：**语气是客观和事实性的。它描述视觉上呈现的内容，而不做过度主观的解释（除非陈述情绪，例如，"冒险的"，"欢乐的重逢"）。 + +--- + +### II. 不同视频类型的规则和特征 + +虽然上述通用规则适用于所有视频，但某些视频类型在其描述中具有独特的特征。 + +**A. 写实/真人视频** + +1. **关注人物细节：**描述优先考虑人物特征：年龄、种族、肤色、发型/颜色和具体服装项目。情绪状态通过面部表情和肢体语言来描述。 + * *示例：*"这位女性，她的脸上带着**担忧的表情**，抬起头与他对视……她的表情突然转变为**震惊和恐慌**……" + +2. **真实世界的动作和互动：**所描述的动作以现实为基础，关注人与人之间或与物理环境的互动。 + * *示例：*"当跑步者到达她身边时，他们立即**紧紧地拥抱在一起**。" + * *示例：*"他用右手撑起自己，**努力站起来**……" + +3. **电影术语：**描述通常包括暗示物理摄像机的特定电影制作术语。 + * *示例：*"**手持拍摄**向前移动，""**中近景**，""**平视角度**，""镜头**以中景稍微移动拍摄**。" + +4. **照明描述：**照明通常用自然光源（"来自阴天的柔和、漫射自然光"）或刻意的电影布光设置（"高调、漫射的背光"，"高对比度、戏剧性的舞台照明"）来描述。 + +**B. 动画/CGI视频** + +1. **强调风格化：**描述突出动画风格和非写实特征。主体通常是奇幻的（恐龙、蓝精灵）或拟人化的（拿着炸药的北极熊）。 + * *示例：*"四只**动画蓝精灵**骑在一只**毛茸茸的大黄兔子**上。" + * *示例：*"一只有橙色翅膀的瓢虫正在高速飞行……在瓢虫后面，有五只黑色苍蝇，它们有**显眼的大红眼睛**。" + +2. **违反物理和夸张的动作：**所描述的动作通常超越现实世界的限制，反映了动画的创作自由。 + * *示例：*"熊用后腿站立，转身，**将点燃的炸药高高地投入拱形门道**。" + * *示例：*"她突然**被甩出车外**……**安全落地**。" + +3. **明确的风格识别：**技术总结通常明确命名动画风格。 + * *示例：*"一个动态的**风格化昆虫的3D CGI动画**集成在照片级真实的自然环境中……" + * *示例：*"视觉风格是**高质量的3D电脑动画电影**……" + +### **II. 镜头控制指南** +以下是镜头控制系统。你应该使用镜头控制系统来描述提示词中的镜头运动。如果遇到下面的类型，就参考对应的描述。 +* **镜头360度旋转**: ["镜头旋转360度", "镜头进行完整旋转", "镜头绕一圈旋转"] +* **镜头第一人称视角FPV**: ["镜头显示第一人称视角", "场景从第一人称视角拍摄", "镜头采用FPV角度", "镜头处于第一人称视角"] +* **镜头向上移动**: ["镜头向上移动", "镜头上升", "镜头升起"] +* **镜头向下移动**: ["镜头向下移动", "镜头下降", "镜头落下"] +* **镜头低角度/仰拍**: ["镜头从低角度拍摄", "镜头从下方捕捉场景", "镜头位于低视点"] +* **镜头向上倾斜**: ["镜头向上倾斜", "镜头进行向上倾斜运动"] +* **镜头向下倾斜**: ["镜头向下倾斜", "镜头进行向下倾斜运动"] +* **镜头地面拍摄**: ["镜头在地面高度", "镜头从地面拍摄", "镜头从地面视角捕捉场景"] +* **镜头向前推进**: ["镜头向前推进", "镜头向前移动", "镜头向前移"] +* **镜头向右平移**: ["镜头向右移动", "镜头向右移"] +* **镜头向后拉**: ["镜头向后拉", "镜头后退", "镜头向后移"] +* **镜头向左平移**: ["镜头向左移动", "镜头向左移"] +* **镜头延时摄影**: ["镜头捕捉延时摄影", "使用延时拍摄", "场景以延时方式显示"] +* **镜头微距拍摄**: ["镜头进行微距拍摄", "使用微距视角"] +* **镜头慢动作**: ["镜头以慢动作记录", "显示慢动作镜头", "场景以慢动作捕捉"] +* **镜头拉远**: ["镜头拉远", "镜头向后拉", "镜头远离主体"] +* **镜头推近**: ["镜头推近", "镜头向前推", "镜头接近主体"] +* **镜头向右平移**: ["镜头向右平移", "镜头向右摆动", "镜头进行向右平移运动"] +* **镜头向左平移**: ["镜头向左平移", "镜头向左摆动", "镜头进行向左平移运动"] +* **镜头无人机视角**: ["镜头显示无人机视角", "场景从无人机视角拍摄", "镜头采用空中无人机角度"] +* **镜头环绕**: ["镜头环绕主体", "镜头围绕主体旋转"] +* **镜头跟随拍摄**: ["镜头跟随主体", "镜头跟踪运动", "镜头与主体一起移动"] +* **镜头过肩拍摄**: ["镜头使用过肩镜头", "镜头位于主体肩膀后方"] +* **镜头逆时针旋转**: ["镜头逆时针旋转", "镜头以逆时针方向旋转"] +* **镜头静止**: ["镜头保持静止", "镜头静止不动", "镜头保持静态"] +* **镜头顺时针旋转**: ["镜头顺时针旋转", "镜头以顺时针方向旋转"] +* **镜头高角度/俯拍**: ["镜头从高角度拍摄", "镜头从上方捕捉场景", "镜头位于高视点"] +* **镜头鱼眼镜头**: ["镜头使用鱼眼镜头", "场景以鱼眼效果显示", "应用鱼眼透视"] +* **镜头鸟瞰视角**: ["镜头显示鸟瞰视角", "场景从正上方拍摄"] + +** 相机运动也可以是动态的，如"最初，摄像机跟随...然后，焦点平滑地转移，向后拉..." + +## **三、标准生成流程** + +在生成最终提示词之前，你必须在思考和构建时遵循以下步骤： + +0. **语言规则**: + * 整体输出的提示词保持为英文。 + * 文本渲染内容应与用户输入的语言相同。例如，如果用户希望视频显示文本"Hello"，则扩展的提示词应该是英文，但渲染的文本应该是"Hello"。如果是中文，文本渲染内容总是包含在“”内，如果是英文则总是放置在""内。 + * **宝可梦IP的特殊规则**：如果用户输入包含宝可梦的IP角色，始终使用英文的IP名称。（例如，使用Jigglypuff而不是胖丁，因为胖丁来自宝可梦） + +1. **分析核心元素并评估风险**: + * **摘要**：提示词以视频故事的摘要开始。所有主要主体都必须在摘要中描述。摘要应该是一个简洁的句子，并放在提示词的开头。 + * **识别核心元素**：从用户输入中清晰地识别主体（人物、对象）、关键动作/事件、运动、环境和整体叙事弧线。 + * **确定实体数量**：如果用户输入包含多个实体，首先检查用户是否给出了精确的实体数量（例如，"六口之家"，"四个朋友"，"大约五名舞者"），那么你必须严格遵循用户的原始提示词并使用相同的数量描述。否则，如果用户使用模糊的词语（例如，"一群警官"）并且没有精确给出实体数量，那么你应该严格将人物或对象的数量限制在三个或更少，以保持场景清晰度。 + * **识别高风险概念**：特别注意复杂的物理互动（例如，动态体育序列）、随时间展开的抽象概念（例如，"能量"脉动或流动）以及动态排版，需要用简单、可渲染的描述来呈现。 + * **可视化概念**：将非视觉概念转换为视觉可用的序列。例如，"发令枪的声音"暗示比赛的开始，运动员从起跑线冲出并沿着跑道移动。你可以从用户提供的这种非视觉描述中推断事件序列。或出现鸟鸣声，那么场景中就需要出现鸟。 + +2. **确定摄影和构图**: + * **摄影和镜头运动**：使用以下规则决定摄影： + * 如果用户输入包含摄影描述（例如，广角镜头、航拍镜头、跟踪镜头、变焦镜头），你必须严格遵循用户输入，不要更改它。 + * 如果用户没有提供描述，基于场景和核心事件，你应该使用你的摄影知识来选择合适的镜头工作。优先选择能完全捕捉关键动作和所有核心元素的镜头和运动（例如，**缓慢平移穿过场景的略高角度镜头**）。 + * **构图和场景调度**：应用构图技术，如三分法或对称，以确保每一帧都经过专业构图。至关重要的是，考虑主体和元素随时间在画面中的移动方式（调度）——它们的进入、退出和运动路径。 + +3. **选择艺术风格**: + * **风格选择**：如果用户指定了风格（例如，油画动画、动漫、动态图形），严格遵守它。否则，首先根据用户输入推断合适的风格。如果没有明确的风格偏好，默认为**电影感写实风格**。 + +4. **确定镜头运动**: + * **镜头运动**：使用以下规则决定镜头运动： + * 如果用户输入包含镜头运动描述（例如，向前推进、推近、向上倾斜），你必须根据输入描述使用镜头运动关键字。（例如，如果用户输入是"镜头向上移动"，你必须推断它对应关键字"camera upward"，并在"camera upward"的值列表中选择一个词，即["The camera moves upward", "The camera rises", "The camera ascends"]） + * *示例：*输入："镜头向上移动"，使用词语"The camera rises"， + * *示例：*输入："高角度镜头"，使用词语"The camera captures the scene from above"。 + * 如果用户没有提供镜头运动描述。镜头拍摄必须是"The shot is at an eye-level angle with the main subject."。 + 并且镜头运动应该从镜头静止、镜头向前推进、镜头向后拉、镜头向右平移、镜头向左平移、镜头推近、镜头拉远中选择。你需要分析主体的运动和背景的运动，然后选择最合适的镜头运动。首选轻微的运动，以防止画面变化过快产生畸变。 + +4. **填充细节和审查**: + * **时长**：视频时长固定为5秒。填充的细节必须能在5秒内完成。 + * **细节填充**：遵循结构化顺序，描述主体的材料、纹理、动作、手势以及随时间变化的表情，以及环境中的任何次要元素及其变化方式。注意添加恰到好处的细节量；除非汗水是视频叙事的关键，否则不要过分强调汗水等元素。火花、雷电也是容易生成错误的内容，需要注意。**镜头或手机屏幕的正反面**：当场景中出现相机镜头、手机屏幕等物体时，必须明确描述其朝向（正面/背面/侧面），以及观众看到的是屏幕内容还是设备背面，避免产生歧义。例如："手机屏幕朝向镜头，显示......"或"相机的镜头正对着......"。 + * **运动填充**：视频中主体的运动必须在场景中逻辑合理且一致。并且运动必须清楚说明。主体的运动必须能在5秒内完成。一个坏例子是描述大量无法在5秒内完成的动作。主体的动作需要详细描述。 + * **逻辑审查**：检查描述是否完全符合物理定律、因果关系和场景逻辑贯穿整个视频时长。例如，动态场景中的动作流程是否合理？角色移动时的视线和互动是否一致？剪辑开始和结束之间是否有连续性错误？ + * **完整性审查**：确保所有预期的关键元素和事件（例如，裁判发出信号、背景爆炸）都被明确描述。检查镜头运动或主体运动是否导致不自然的裁剪或遮挡，特别是在动作期间的人体肢体或关键对象。 + * **生成畸形审查**：必须高度警惕任何可能引发动作或身体畸形的描述，严密审查并排除一切可能导致角色动作不自然、关节弯曲异常或肢体结构错误的因素，确保所有人体结构、动作轨迹、关节形态均自然、连贯，完全符合物理规律。 + +5. **最终验证**: + * **用户输入遵循检查**：将最终结果与用户输入进行比较，以确保完全描述了用户的核心内容，例如核心实体及其属性、对象或人物的数量、指定的摄影、时长、节奏和事件顺序。最终提示词不得添加用户输入未暗示的任何新前景对象或主要事件。 + * **检查物理和时间逻辑**：检查提示词中是否存在任何物理或时间逻辑错误。例如，对象大小在移动期间是否一致？运动物理是否可信（例如，加速度、重力）？反应的时间是否合理？事件之间的因果关系是否清晰？ + * **检查宝可梦IP**：检查提示词是否包含任何宝可梦IP；你应该使用宝可梦角色的英文名称。（例如，使用Jigglypuff而不是胖丁） + +6. **如果验证失败则重试**: + * **从头开始**：如果验证失败，你应该从头开始重新生成扩展的提示词。 + + + +### **四、风格特定创作指南** + +根据确定的艺术风格激活相应的专业知识库。 + +**1. 摄影和写实风格** +* **总体规则**：想象你是一位摄影大师，用户的输入将被你转换为专业摄影师拍摄的照片。假设你正在查看用户描述的图片，你将使用你的专业摄影知识将用户输入转换为具有专业构图和专业照明的视觉图片。 +* **摄影风格视频的专业照明**：你应该使用你的专业摄影照明技术来增强真实感。你应该根据用户的输入选择合适的照明技术，以下是一些可供选择的示例技术，你可以使用你的世界知识来获得更好的选择。 + * 使用戏剧性照明，强调光影之间的高对比度以创造深度感。使用伦勃朗照明从45度角照亮主体的面部，在脸颊上形成三角形高光，并在面部的另一侧投下深深的阴影。这应该以强烈的维度感突出面部特征。 + * 对于整体场景，应用黄金时段照明，具有温暖的色温，光线在过渡到阴影时逐渐柔化。柔和的侧光应该照亮主体的身体，在光影之间创造柔和的渐变，增强三维形式。背景应该具有渐进的阴影过渡，唤起平静、诱人的氛围。 + * 为了增加真实感，使用背光在主体周围创造剪影效果，强光位于主体后面以突出轮廓。这应该将背景投入深深的阴影中，边缘由附近表面的反射光柔化。确保主体的轮廓清晰，但周围的阴影增加了一种神秘感。 + * 从左侧加入冷光，与温暖的背景光形成对比，创造视觉张力。光强度应该变化，在关键特征上使用强烈的高光，在其他区域使用柔和的阴影。光影之间的过渡应该感觉自然和平衡，焦点区域由锐利的硬光照亮，阴影中的区域更柔和。 + * 为了增加额外的纹理和真实感，在水或玻璃等反射表面上包括反射，光线应该反弹以柔化阴影并在表面上创造光斑。直射光和反射光之间的这种相互作用将为整体构图增加复杂性和趣味性。 + * 确保光线在引导观众的眼睛穿过视频中起着关键作用，无论是通过阴影创造的引导线，还是通过突出场景关键元素的光线定向流动。 + * 调整以获得更多控制的关键参数： + * 照明风格：（例如，伦勃朗、柔和、硬、背光、侧光） + * 光线方向：（例如，45度角、自上而下、侧光） + * 光线质量：（例如，柔和、刺眼、漫射、聚光灯） + * 阴影细节：（例如，深阴影、柔和渐变、高对比度） + * 色温：（例如，温暖的黄金时段、凉爽的日光） + * 反射：（例如，水、玻璃或金属表面上的反射光） + * 剪影和轮廓：（例如，主体背光，创造戏剧性的轮廓） + * 可定制元素： + * 照明情绪：为场景定义整体照明情绪（例如，戏剧性、柔和、高对比度、微妙渐变）。 + * 背景照明：调整光线与背景的相互作用，例如柔和渐变或强烈阴影区域。 + * 柔和对比硬阴影：指定阴影应该有多刺眼或多漫射。 + * 高光细节：关注光线应该突出关键特征的区域（例如，面部、眼睛、纹理）。 + * 氛围：（例如，忧郁、宁静、戏剧性、和平） + * **重要规则**：如果光源不可见，只需描述光效，不要提及视频中未出现的任何光源，一个坏例子是"一个次要的、看不见的温暖光源，可能是画面外的台灯"，像这个例子的描述是被禁止的。 +* **镜头效果**：使用专业术语来描述镜头效果（例如，广角透视、长焦压缩、浅景深），你应该使用你的世界知识为用户的输入选择最佳镜头效果。 +* **构图**：使用专业术语为用户的输入选择最佳构图（例如，引导线、框架构图、三分法），**但不要直接使用构图技术名称，你应该安排人物/对象/环境来反映构图本身**。 +* **极致细节**：深入描述材料纹理（例如，木纹反射、织物纤维）、角色细节（例如，眼睛中的眼神光、皮肤毛孔）和环境氛围（例如，空气中的灰尘颗粒）。 +* **多实体场景**：如果用户的输入表明视频中有多个实体，例如一群人、一个团队或场景中的多个人，如果用户给出了精确的实体数量，则严格遵守用户的输入，不要更改数量。但如果用户没有给出具体的人物或对象数量，严格将人物或对象的数量限制在三个或更少。详细描述每个人/对象，并将它们放在中景到前景，确保他们的面部和肢体清晰、未变形，并且在关节处没有被裁剪。**这一条非常重要** +* **电影摄影写实**：如果用户的输入表明视频是写实风格，风格必须是"cinematic realistic style"（电影摄影写实风格）。 + +**2. 插画和绘画风格（卡通、油画、水彩等）** +* **定义类型**：精确定义风格（例如，"日本赛璐珞动画风格"，"厚涂油画"，"湿画法水彩"，"印象派点彩画法"）。 +* **媒介特定特征**：专注于描述风格的独特视觉语言，例如线条的粗细（"G笔线稿"）、笔触的纹理（"可见的、三维的笔触"）和颜料的特性（"水彩边缘的自然水渍"）。 +* **角色设计**：强调夸张的特征（例如，"Q版身体比例"，"占据脸部三分之一的大眼睛"）和富有表现力的姿势。 + +**3. 字体艺术** +* **主体优先**：描述必须以类似`The words "[Text Content]" rendered as...`的短语开始，以将文本确立为绝对核心主体。 +* **安全透视和构图**：强制使用安全的正面或俯视图（`Front view`，`Top-down`），并将文本主体放在画面中心，使用简单或广阔的背景进行对比，从根本上防止裁剪。 +* **形式而非形成**：描述文本构成的最终"形式"，而不是其"形成过程"。（错误："两个GPU相互倾斜形成一个A"；正确："由GPU制成的金字塔形状的字母A"）。 +* **完整性保险**：在提示词末尾添加强制性指令，例如`The entire [object/phrase] is fully visible`，作为防止裁剪的最后防线。 +* **禁止高风险词汇**：避免使用"巨大的"、"特写"、"复杂的"和"精致的"等词，因为它们可能诱导AI放大主体的一部分，导致裁剪。 + +### **五、最终输出要求** + +1. **仅输出最终提示词**：不要显示任何思考过程、Markdown格式或与文本到视频提示词无关的任何表达（如"摘要显示"）。 +2. **输出语言**：扩展的提示词应该是英文，同时根据语言规则保持文本渲染内容与用户输入的语言相同。 +3. **忠实于输入**：你必须保留用户输入句子中的核心概念、属性、数量和文本渲染内容。 +4. **风格强化**：在提示词中提及核心风格3-5次，并以风格声明句结束（例如，"整个视频是电影感写实风格"）。 +5. **字数控制**：描述主要主体的长度应该在140个单词左右，这里面主体的动作是重要部分。描述背景的长度应该是70个单词。描述其他属性（包括构图、光线、风格、氛围、拍摄角度、拍摄类型）的总长度应该在140个单词以内。 +6. **避免自我引用**：在开头直接描述视频内容，删除冗余短语，如"这个视频"或"这个视频显示了"。 + +接下来，我将提供输入句子，你将提供扩展的提示词。 +""" \ No newline at end of file diff --git a/trainer/models/loader/__init__.py b/trainer/models/loader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/models/loader/component_loader.py b/trainer/models/loader/component_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..d3cbc07c21723dfa9efb3753854920d604f85ccb --- /dev/null +++ b/trainer/models/loader/component_loader.py @@ -0,0 +1,556 @@ +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import glob +import json +import os +import time +from abc import ABC, abstractmethod +from collections.abc import Generator, Iterable +from copy import deepcopy +from typing import cast + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors.torch import load_file as safetensors_load_file +from torch.distributed import init_device_mesh +from transformers import AutoImageProcessor, AutoTokenizer +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from trainer.configs.models import EncoderConfig +from trainer.distributed import get_local_torch_device +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.models.hf_transformer_utils import get_diffusers_config +from trainer.models.loader.fsdp_load import maybe_load_fsdp_model, shard_model +from trainer.models.loader.utils import set_default_torch_dtype +from trainer.models.loader.weight_utils import ( + filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, + pt_weights_iterator, safetensors_weights_iterator) +from trainer.models.registry import ModelRegistry +from trainer.platforms import current_platform +from trainer.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class ComponentLoader(ABC): + """Base class for loading a specific type of model component.""" + + def __init__(self, device=None) -> None: + self.device = device + + @abstractmethod + def load(self, model_path: str, trainer_args: TrainerArgs): + """ + Load the component based on the model path, architecture, and inference args. + + Args: + model_path: Path to the component model + trainer_args: TrainerArgs + + Returns: + The loaded component + """ + raise NotImplementedError + + @classmethod + def for_module_type(cls, module_type: str, + transformers_or_diffusers: str) -> 'ComponentLoader': + """ + Factory method to create a component loader for a specific module type. + + Args: + module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler") + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + A component loader for the specified module type + """ + # Map of module types to their loader classes and expected library + module_loaders = { + "scheduler": (SchedulerLoader, "diffusers"), + "transformer": (TransformerLoader, "diffusers"), + "transformer_2": (TransformerLoader, "diffusers"), + "vae": (VAELoader, "diffusers"), + "text_encoder": (TextEncoderLoader, "transformers"), + "text_encoder_2": (TextEncoderLoader, "transformers"), + "tokenizer": (TokenizerLoader, "transformers"), + "tokenizer_2": (TokenizerLoader, "transformers"), + "image_processor": (ImageProcessorLoader, "transformers"), + "image_encoder": (ImageEncoderLoader, "transformers"), + } + + if module_type in module_loaders: + loader_cls, expected_library = module_loaders[module_type] + # Assert that the library matches what's expected for this module type + assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + return loader_cls() + + # For unknown module types, use a generic loader + logger.warning( + "No specific loader found for module type: %s. Using generic loader.", + module_type) + return GenericComponentLoader(transformers_or_diffusers) + + +class TextEncoderLoader(ComponentLoader): + """Loader for text encoders.""" + + @dataclasses.dataclass + class Source: + """A source for weights.""" + + model_or_path: str + """The model ID or path.""" + + prefix: str = "" + """A prefix to prepend to all weights.""" + + fall_back_to_pt: bool = True + """Whether .pt weights can be used.""" + + allow_patterns_overrides: list[str] | None = None + """If defined, weights will load exclusively using these patterns.""" + + counter_before_loading_weights: float = 0.0 + counter_after_loading_weights: float = 0.0 + + def _prepare_weights( + self, + model_name_or_path: str, + fall_back_to_pt: bool, + allow_patterns_overrides: list[str] | None, + ) -> tuple[str, list[str], bool]: + """Prepare weights for the model. + + If the model is not local, it will be downloaded.""" + # model_name_or_path = (self._maybe_download_from_modelscope( + # model_name_or_path, revision) or model_name_or_path) + + is_local = os.path.isdir(model_name_or_path) + assert is_local, "Model path must be a local directory" + + use_safetensors = False + index_file = SAFE_WEIGHTS_INDEX_NAME + allow_patterns = ["*.safetensors", "*.bin"] + + if fall_back_to_pt: + allow_patterns += ["*.pt"] + + if allow_patterns_overrides is not None: + allow_patterns = allow_patterns_overrides + + hf_folder = model_name_or_path + + hf_weights_files: list[str] = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if len(hf_weights_files) > 0: + if pattern == "*.safetensors": + use_safetensors = True + break + + if use_safetensors: + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file) + else: + hf_weights_files = filter_files_not_needed_for_inference( + hf_weights_files) + + if len(hf_weights_files) == 0: + raise RuntimeError( + f"Cannot find any model weights with `{model_name_or_path}`") + + return hf_folder, hf_weights_files, use_safetensors + + def _get_weights_iterator( + self, source: "Source", + to_cpu: bool) -> Generator[tuple[str, torch.Tensor], None, None]: + """Get an iterator for the model weights based on the load format.""" + hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + source.model_or_path, source.fall_back_to_pt, + source.allow_patterns_overrides) + if use_safetensors: + weights_iterator = safetensors_weights_iterator(hf_weights_files, + to_cpu=to_cpu) + else: + weights_iterator = pt_weights_iterator(hf_weights_files, + to_cpu=to_cpu) + + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # Apply the prefix. + return ((source.prefix + name, tensor) + for (name, tensor) in weights_iterator) + + def _get_all_weights( + self, + model: nn.Module, + model_path: str, + to_cpu: bool, + ) -> Generator[tuple[str, torch.Tensor], None, None]: + primary_weights = TextEncoderLoader.Source( + model_path, + prefix="", + fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), + allow_patterns_overrides=getattr(model, "allow_patterns_overrides", + None), + ) + yield from self._get_weights_iterator(primary_weights, to_cpu) + + secondary_weights = cast( + Iterable[TextEncoderLoader.Source], + getattr(model, "secondary_weights", ()), + ) + for source in secondary_weights: + yield from self._get_weights_iterator(source, to_cpu) + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the text encoders based on the model path, and inference args.""" + model_config = get_diffusers_config(model=model_path) + model_config.pop("_name_or_path", None) + model_config.pop("transformers_version", None) + model_config.pop("model_type", None) + model_config.pop("tokenizer_class", None) + model_config.pop("torch_dtype", None) + logger.info("HF Model config: %s", model_config) + + # @TODO(Wei): Better way to handle this? + try: + encoder_config = trainer_args.pipeline_config.text_encoder_configs[ + 0] + encoder_config.update_model_arch(model_config) + encoder_precision = trainer_args.pipeline_config.text_encoder_precisions[ + 0] + except Exception: + encoder_config = trainer_args.pipeline_config.text_encoder_configs[ + 1] + encoder_config.update_model_arch(model_config) + encoder_precision = trainer_args.pipeline_config.text_encoder_precisions[ + 1] + + target_device = get_local_torch_device() + # TODO(will): add support for other dtypes + return self.load_model(model_path, encoder_config, target_device, + trainer_args, encoder_precision) + + def load_model(self, + model_path: str, + model_config: EncoderConfig, + target_device: torch.device, + trainer_args: TrainerArgs, + dtype: str = "fp16"): + use_cpu_offload = trainer_args.text_encoder_cpu_offload and len( + getattr(model_config, "_fsdp_shard_conditions", [])) > 0 + + if trainer_args.text_encoder_cpu_offload: + target_device = torch.device( + "mps") if current_platform.is_mps() else torch.device("cpu") + + with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): + with target_device: + architectures = getattr(model_config, "architectures", []) + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + model = model_cls(model_config) + + weights_to_load = {name for name, _ in model.named_parameters()} + loaded_weights = model.load_weights( + self._get_all_weights(model, model_path, + to_cpu=use_cpu_offload)) + self.counter_after_loading_weights = time.perf_counter() + logger.info( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights - + self.counter_before_loading_weights) + + # Explicitly move model to target device after loading weights + model = model.to(target_device) + + if use_cpu_offload: + # Disable FSDP for MPS as it's not compatible + if current_platform.is_mps(): + logger.info( + "Disabling FSDP sharding for MPS platform as it's not compatible" + ) + else: + mesh = init_device_mesh( + "cuda", + mesh_shape=(1, dist.get_world_size()), + mesh_dim_names=("offload", "replicate"), + ) + shard_model( + model, + cpu_offload=True, + reshard_after_forward=True, + mesh=mesh["offload"], + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=trainer_args.pin_cpu_memory) + # We only enable strict check for non-quantized models + # that have loaded weights tracking currently. + # if loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + raise ValueError("Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}") + + return model.eval() + + +class ImageEncoderLoader(TextEncoderLoader): + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the text encoders based on the model path, and inference args.""" + with open(os.path.join(model_path, "config.json")) as f: + model_config = json.load(f) + model_config.pop("_name_or_path", None) + model_config.pop("transformers_version", None) + model_config.pop("torch_dtype", None) + model_config.pop("model_type", None) + logger.info("HF Model config: %s", model_config) + + encoder_config = trainer_args.pipeline_config.image_encoder_config + encoder_config.update_model_arch(model_config) + + if trainer_args.image_encoder_cpu_offload: + target_device = torch.device("mps") if current_platform.is_mps() else torch.device("cpu") + else: + target_device = get_local_torch_device() + # TODO(will): add support for other dtypes + return self.load_model( + model_path, encoder_config, target_device, trainer_args, + trainer_args.pipeline_config.image_encoder_precision) + + +class ImageProcessorLoader(ComponentLoader): + """Loader for image processor.""" + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the image processor based on the model path, and inference args.""" + logger.info("Loading image processor from %s", model_path) + + image_processor = AutoImageProcessor.from_pretrained(model_path, ) + logger.info("Loaded image processor: %s", + image_processor.__class__.__name__) + return image_processor + + +class TokenizerLoader(ComponentLoader): + """Loader for tokenizers.""" + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the tokenizer based on the model path, and inference args.""" + logger.info("Loading tokenizer from %s", model_path) + + tokenizer = AutoTokenizer.from_pretrained( + model_path, # "/tokenizer" + # in v0, this was same string as encoder_name "ClipTextModel" + # TODO(will): pass these tokenizer kwargs from inference args? Maybe + # other method of config? + padding_size='right', + ) + logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__) + return tokenizer + + +class VAELoader(ComponentLoader): + """Loader for VAE.""" + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the VAE based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." + trainer_args.model_paths["vae"] = model_path + + vae_config = trainer_args.pipeline_config.vae_config + vae_config.update_model_arch(config) + + if trainer_args.vae_cpu_offload: + target_device = torch.device("mps") if current_platform.is_mps() else torch.device("cpu") + else: + target_device = get_local_torch_device() + + with set_default_torch_dtype(PRECISION_TO_TYPE[ + trainer_args.pipeline_config.vae_precision]): + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = vae_cls(vae_config).to(target_device) + + # Find all safetensors files + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors")) + # TODO(PY) + assert len( + safetensors_list + ) == 1, f"Found {len(safetensors_list)} safetensors files in {model_path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict( + loaded, strict=False) # We might only load encoder or decoder + + return vae.eval() + + +class TransformerLoader(ComponentLoader): + """Loader for transformer.""" + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the transformer based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + hf_config = deepcopy(config) + cls_name = config.pop("_class_name") + if cls_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") + + trainer_args.model_paths["transformer"] = model_path + + if trainer_args.module_name == "transformer": + cls_name = trainer_args.cls_name + load_from_dir = trainer_args.load_from_dir + ar_action_load_from_dir = trainer_args.ar_action_load_from_dir + else: + cls_name = "WanTransformer3DBidirectActionModel" + load_from_dir = model_path + ar_action_load_from_dir = None + model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) + + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + + logger.info("Loading model from %s safetensors files in %s", + len(safetensors_list), model_path) + + default_dtype = PRECISION_TO_TYPE[ + trainer_args.pipeline_config.dit_precision] + + # Load the model using FSDP loader + logger.info("Loading model from %s, default_dtype: %s", cls_name, + default_dtype) + assert trainer_args.hsdp_shard_dim is not None + model = maybe_load_fsdp_model( + load_from_dir=load_from_dir, + ar_action_load_from_dir=ar_action_load_from_dir, + cls_name=cls_name, + model_cls=model_cls, + init_params={}, + weight_dir_list=safetensors_list, + device=get_local_torch_device(), + hsdp_replicate_dim=trainer_args.hsdp_replicate_dim, + hsdp_shard_dim=trainer_args.hsdp_shard_dim, + cpu_offload=trainer_args.dit_cpu_offload, + pin_cpu_memory=trainer_args.pin_cpu_memory, + fsdp_inference=trainer_args.use_fsdp_inference, + # TODO(will): make these configurable + # param_dtype=torch.bfloat16, + param_dtype=torch.float32, + reduce_dtype=torch.float32, + output_dtype=None, + training_mode=trainer_args.training_mode) + + + total_params = sum(p.numel() for p in model.parameters()) + logger.info("Loaded model with %.2fB parameters", total_params / 1e9) + + dtypes = set(param.dtype for param in model.parameters()) + if len(dtypes) > 1: + model = model.to(default_dtype) + model = model.eval() + return model + + +class SchedulerLoader(ComponentLoader): + """Loader for scheduler.""" + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load the scheduler based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." + + scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + scheduler = scheduler_cls(**config) + if trainer_args.pipeline_config.flow_shift is not None: + scheduler.set_shift(trainer_args.pipeline_config.flow_shift) + if trainer_args.pipeline_config.timesteps_scale is not None: + scheduler.set_timesteps_scale( + trainer_args.pipeline_config.timesteps_scale) + return scheduler + + +class GenericComponentLoader(ComponentLoader): + """Generic loader for components that don't have a specific loader.""" + + def __init__(self, library="transformers") -> None: + super().__init__() + self.library = library + + def load(self, model_path: str, trainer_args: TrainerArgs): + """Load a generic component based on the model path, and inference args.""" + logger.warning("Using generic loader for %s with library %s", + model_path, self.library) + + if self.library == "transformers": + from transformers import AutoModel + + model = AutoModel.from_pretrained( + model_path, + trust_remote_code=trainer_args.trust_remote_code, + revision=trainer_args.revision, + ) + logger.info("Loaded generic transformers model: %s", + model.__class__.__name__) + return model + elif self.library == "diffusers": + logger.warning( + "Generic loading for diffusers components is not fully implemented" + ) + + model_config = get_diffusers_config(model=model_path) + logger.info("Diffusers Model config: %s", model_config) + # This is a placeholder - in a real implementation, you'd need to handle this properly + return None + else: + raise ValueError(f"Unsupported library: {self.library}") + + +class PipelineComponentLoader: + """ + Utility class for loading pipeline components. + This replaces the chain of if-else statements in load_pipeline_module. + """ + + @staticmethod + def load_module(module_name: str, component_model_path: str, + transformers_or_diffusers: str, + trainer_args: TrainerArgs): + """ + Load a pipeline module. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + pipeline_args: Inference arguments + + Returns: + The loaded module + """ + logger.info( + "Loading %s using %s from %s", + module_name, + transformers_or_diffusers, + component_model_path, + ) + + # Get the appropriate loader for this module type + loader = ComponentLoader.for_module_type(module_name, + transformers_or_diffusers) + + # Load the module + return loader.load(component_model_path, trainer_args) diff --git a/trainer/models/loader/fsdp_load.py b/trainer/models/loader/fsdp_load.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f1ed4425d31b81fef4fb70b351942107cca23e --- /dev/null +++ b/trainer/models/loader/fsdp_load.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from torchtune +# Copyright 2024 The TorchTune Authors. +# Copyright 2025 The Trainer Authors. + +import contextlib +from collections.abc import Callable, Generator +from itertools import chain +from typing import Any +import os + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed import DeviceMesh, init_device_mesh +from torch.distributed._tensor import distribute_tensor +from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, + MixedPrecisionPolicy, fully_shard) +from torch.nn.modules.module import _IncompatibleKeys +from torch.nn.parallel import DistributedDataParallel as DDP + +from functools import partial +from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType +from torch.distributed.fsdp.api import CPUOffload +from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy + +from trainer.distributed import get_world_group +from trainer.logger import init_logger +from trainer.models.loader.utils import (get_param_names_mapping, + hf_to_custom_state_dict) +from trainer.models.loader.weight_utils import safetensors_weights_iterator +from trainer.utils import set_mixed_precision_policy + +logger = init_logger(__name__) + + +# TODO(PY): move this to utils elsewhere +@contextlib.contextmanager +def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: + """ + Context manager to set torch's default dtype. + + Args: + dtype (torch.dtype): The desired default dtype inside the context manager. + + Returns: + ContextManager: context manager for setting default dtype. + + Example: + >>> with set_default_dtype(torch.bfloat16): + >>> x = torch.tensor([1, 2, 3]) + >>> x.dtype + torch.bfloat16 + + + """ + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(old_dtype) + + +# TODO(PY): add compile option +def maybe_load_fsdp_model( + load_from_dir: str, + ar_action_load_from_dir: str, + cls_name: str, + model_cls: type[nn.Module], + init_params: dict[str, Any], + weight_dir_list: list[str], + device: torch.device, + hsdp_replicate_dim: int, + hsdp_shard_dim: int, + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + cpu_offload: bool = False, + fsdp_inference: bool = False, + output_dtype: torch.dtype | None = None, + training_mode: bool = True, + pin_cpu_memory: bool = True, +) -> torch.nn.Module: + """ + Load the model with FSDP if is training, else load the model without FSDP. + """ + # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are + # manually casting the inputs to the model + mp_policy = MixedPrecisionPolicy(param_dtype, + reduce_dtype, + output_dtype, + cast_forward_inputs=False) + + set_mixed_precision_policy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + + # with set_default_dtype(param_dtype), torch.device("meta"): + with set_default_dtype(param_dtype): + logger.info(f"loading from: {load_from_dir}") + logger.info(f"model:", model_cls) + # model = model_cls(**init_params) + model = model_cls.from_pretrained( + load_from_dir, local_attn_size=-1, sink_size=0) + + + if "Action" in cls_name: + model.add_discrete_action_parameters() + if ar_action_load_from_dir is not None: + from safetensors.torch import load_file + state_dict = load_file(ar_action_load_from_dir) + model.load_state_dict(state_dict, strict=True) + logger.info(f"loading from: {ar_action_load_from_dir}") + # model.add_channel_concat_parameters() + + model.to(torch.float32) + + # Check if we should use FSDP + use_fsdp = training_mode or fsdp_inference + + # Disable FSDP for MPS as it's not compatible + from trainer.platforms import current_platform + if current_platform.is_mps(): + use_fsdp = False + logger.info("Disabling FSDP for MPS platform as it's not compatible") + + if use_fsdp: + world_size = hsdp_replicate_dim * hsdp_shard_dim + if not training_mode and not fsdp_inference: + hsdp_replicate_dim = world_size + hsdp_shard_dim = 1 + + device_mesh = init_device_mesh( + "cuda", + # (Replicate(), Shard(dim=0)) + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + shard_model(model, + cpu_offload=cpu_offload, + reshard_after_forward=True, + mp_policy=mp_policy, + mesh=device_mesh, + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory) + + for n, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError( + f"Unexpected param or buffer {n} on meta device.") + # Avoid unintended computation graph accumulation during inference + if isinstance(p, torch.nn.Parameter): + p.requires_grad = False + return model + +def build_2d_groups(rank: int, world_size: int, fsdp_group_size: int): + assert world_size % fsdp_group_size == 0, "world_size 必须能整除 FSDP_GROUP_SIZE" + num_replicas = world_size // fsdp_group_size + + shard_group_id = rank // fsdp_group_size + shard_index = rank % fsdp_group_size + shard_ranks = list(range(shard_group_id * fsdp_group_size, + (shard_group_id + 1) * fsdp_group_size)) + fsdp_pg = dist.new_group(ranks=shard_ranks) + + replica_ranks = [shard_index + k * fsdp_group_size for k in range(num_replicas)] + ddp_pg = dist.new_group(ranks=replica_ranks) + + ddp_rank_in_group = replica_ranks.index(rank) + + info = { + "fsdp_pg": fsdp_pg, + "ddp_pg": ddp_pg, + "num_replicas": num_replicas, + "ddp_rank": ddp_rank_in_group, + "shard_group_id": shard_group_id, + "shard_index": shard_index, + "shard_ranks": shard_ranks, + "replica_ranks": replica_ranks, + } + return info + +def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", process_group=None, min_num_params=int(5e7), transformer_module=None, cpu_offload=False): + if mixed_precision: + mixed_precision_policy = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False + ) + else: + mixed_precision_policy = None + + if wrap_strategy == "transformer": + auto_wrap_policy = partial( + transformer_auto_wrap_policy, + transformer_layer_cls=transformer_module + ) + elif wrap_strategy == "size": + auto_wrap_policy = partial( + size_based_auto_wrap_policy, + min_num_params=min_num_params + ) + else: + raise ValueError(f"Invalid wrap strategy: {wrap_strategy}") + + os.environ["NCCL_CROSS_NIC"] = "1" + + sharding_strategy = { + "full": ShardingStrategy.FULL_SHARD, + "hybrid_full": ShardingStrategy.HYBRID_SHARD, + "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2, + "no_shard": ShardingStrategy.NO_SHARD, + }[sharding_strategy] + + module = FSDP( + module, + auto_wrap_policy=auto_wrap_policy, + sharding_strategy=sharding_strategy, + mixed_precision=mixed_precision_policy, + device_id=torch.cuda.current_device(), + limit_all_gathers=True, + use_orig_params=True, + cpu_offload=CPUOffload(offload_params=cpu_offload), + sync_module_states=False, # Load ckpt on rank 0 and sync to other ranks + # process_group=process_group, + ) + return module + +def shard_model( + model, + *, + cpu_offload: bool, + reshard_after_forward: bool = True, + mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(), # noqa + mesh: DeviceMesh | None = None, + fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [], # noqa + pin_cpu_memory: bool = True, +) -> None: + """ + Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API. + + This method will over the model's named modules from the bottom-up and apply shard modules + based on whether they meet any of the criteria from shard_conditions. + + Args: + model (TransformerDecoder): Model to shard with FSDP. + shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine + which modules to shard with FSDP. Each function should take module name (relative to root) + and the module itself, returning True if FSDP should shard the module and False otherwise. + If any of shard_conditions return True for a given module, it will be sharded by FSDP. + cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer + states to CPU. + reshard_after_forward (bool): Whether to reshard parameters and buffers after + the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy + from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy. + mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism. + Default to None. + fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine + which modules to shard with FSDP. + pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters. + + Raises: + ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered. + """ + if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0: + logger.warning( + "The FSDP shard condition list is empty or None. No modules will be sharded in %s", + type(model).__name__) + return + + fsdp_kwargs = { + "reshard_after_forward": reshard_after_forward, + "mesh": mesh, + "mp_policy": mp_policy, + } + if cpu_offload: + fsdp_kwargs["offload_policy"] = CPUOffloadPolicy( + pin_memory=pin_cpu_memory) + + # iterating in reverse to start with + # lowest-level modules first + num_layers_sharded = 0 + # TODO(will): don't reshard after forward for the last layer to save on the + # all-gather that will immediately happen Shard the model with FSDP, + for n, m in reversed(list(model.named_modules())): + if any([ + shard_condition(n, m) + for shard_condition in fsdp_shard_conditions + ]): + fully_shard(m, **fsdp_kwargs) + num_layers_sharded += 1 + + if num_layers_sharded == 0: + raise ValueError( + "No layer modules were sharded. Please check if shard conditions are working as expected." + ) + + # Finally shard the entire model to account for any stragglers + fully_shard(model, **fsdp_kwargs) + + +# TODO(PY): device mesh for cfg parallel +def load_model_from_full_model_state_dict( + model: FSDPModule | torch.nn.Module, + full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None], + device: torch.device, + param_dtype: torch.dtype, + strict: bool = False, + cpu_offload: bool = False, + param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, + training_mode: bool = True, +) -> _IncompatibleKeys: + """ + Converting full state dict into a sharded state dict + and loading it into FSDP model (if training) or normal huggingface model + Args: + model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict + full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs + device (torch.device): device used to move full state dict tensors + param_dtype (torch.dtype): dtype used to move full state dict tensors + strict (bool): flag to check if to load the model in strict mode + cpu_offload (bool): flag to check if FSDP offload is enabled + param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name + training_mode (bool): apply FSDP only for training + Returns: + ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: + * **missing_keys** is a list of str containing the missing keys + * **unexpected_keys** is a list of str containing the unexpected keys + + Raises: + NotImplementedError: If got FSDP with more than 1D. + """ + meta_sd = model.state_dict() + sharded_sd = {} + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( + full_sd_iterator, param_names_mapping) # type: ignore + for target_param_name, full_tensor in custom_param_sd.items(): + meta_sharded_param = meta_sd.get(target_param_name) + if meta_sharded_param is None: + raise ValueError( + f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." + ) + if not hasattr(meta_sharded_param, "device_mesh"): + full_tensor = full_tensor.to(device=device, dtype=param_dtype) + # In cases where parts of the model aren't sharded, some parameters will be plain tensors + sharded_tensor = full_tensor + else: + full_tensor = full_tensor.to(device=device, dtype=param_dtype) + sharded_tensor = distribute_tensor( + full_tensor, + meta_sharded_param.device_mesh, + meta_sharded_param.placements, + ) + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() + sharded_sd[target_param_name] = nn.Parameter(sharded_tensor) + + model.reverse_param_names_mapping = reverse_param_names_mapping + unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) + if unused_keys: + logger.warning("Found unloaded parameters in meta state dict: %s", + unused_keys) + + # List of allowed parameter name patterns + ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed + for new_param_name in unused_keys: + if not any(pattern in new_param_name + for pattern in ALLOWED_NEW_PARAM_PATTERNS): + logger.error("Unsupported new parameter: %s. Allowed patterns: %s", + new_param_name, ALLOWED_NEW_PARAM_PATTERNS) + raise ValueError( + f"New parameter '{new_param_name}' is not supported. " + f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed." + ) + meta_sharded_param = meta_sd.get(new_param_name) + if not hasattr(meta_sharded_param, "device_mesh"): + # Initialize with zeros + sharded_tensor = torch.zeros_like(meta_sharded_param, + device=device, + dtype=param_dtype) + else: + # Initialize with zeros and distribute + full_tensor = torch.zeros_like(meta_sharded_param, + device=device, + dtype=param_dtype) + sharded_tensor = distribute_tensor( + full_tensor, + meta_sharded_param.device_mesh, + meta_sharded_param.placements, + ) + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() + sharded_sd[new_param_name] = nn.Parameter(sharded_tensor) + + # choose `assign=True` since we cannot call `copy_` on meta tensor + return model.load_state_dict(sharded_sd, strict=strict, assign=True) diff --git a/trainer/models/loader/utils.py b/trainer/models/loader/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..17d3d92135797f17ffc918a67049d14a65c21a5d --- /dev/null +++ b/trainer/models/loader/utils.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for selecting and loading models.""" +import contextlib +import re +from collections import defaultdict +from collections.abc import Callable, Iterator +from typing import Any + +import torch + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +@contextlib.contextmanager +def set_default_torch_dtype(dtype: torch.dtype): + """Sets the default torch dtype to the given dtype.""" + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + yield + torch.set_default_dtype(old_dtype) + + +def get_param_names_mapping( + mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]: + """ + Creates a mapping function that transforms parameter names using regex patterns. + + Args: + mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns + param_name (str): The parameter name to be transformed + + Returns: + Callable[[str], str]: A function that maps parameter names from source to target format + """ + + def mapping_fn(name: str) -> tuple[str, Any, Any]: + # Try to match and transform the name using the regex patterns in mapping_dict + for pattern, replacement in mapping_dict.items(): + match = re.match(pattern, name) + if match: + merge_index = None + total_splitted_params = None + if isinstance(replacement, tuple): + merge_index = replacement[1] + total_splitted_params = replacement[2] + replacement = replacement[0] + name = re.sub(pattern, replacement, name) + return name, merge_index, total_splitted_params + + # If no pattern matches, return the original name + return name, None, None + + return mapping_fn + + +def hf_to_custom_state_dict( + hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], + param_names_mapping: Callable[[str], tuple[str, Any, Any]] +) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: + """ + Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary. + + Args: + hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary + param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format + + Returns: + custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict + reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf + """ + custom_param_sd = {} + to_merge_params = defaultdict(dict) # type: ignore + reverse_param_names_mapping = {} + if isinstance(hf_param_sd, dict): + hf_param_sd = hf_param_sd.items() # type: ignore + for source_param_name, full_tensor in hf_param_sd: # type: ignore + target_param_name, merge_index, num_params_to_merge = param_names_mapping( + source_param_name) + reverse_param_names_mapping[target_param_name] = (source_param_name, + merge_index, + num_params_to_merge) + if merge_index is not None: + to_merge_params[target_param_name][merge_index] = full_tensor + if len(to_merge_params[target_param_name]) == num_params_to_merge: + # cat at output dim according to the merge_index order + sorted_tensors = [ + to_merge_params[target_param_name][i] + for i in range(num_params_to_merge) + ] + full_tensor = torch.cat(sorted_tensors, dim=0) + del to_merge_params[target_param_name] + else: + continue + custom_param_sd[target_param_name] = full_tensor + return custom_param_sd, reverse_param_names_mapping diff --git a/trainer/models/loader/weight_utils.py b/trainer/models/loader/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c7b8b96827636e55260bb47e391f53276d927ff9 --- /dev/null +++ b/trainer/models/loader/weight_utils.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py +"""Utilities for downloading and initializing model weights.""" +import hashlib +import json +import os +import tempfile +from collections.abc import Generator +from pathlib import Path + +import filelock +import huggingface_hub.constants +import torch +from safetensors.torch import safe_open +from tqdm.auto import tqdm + +from trainer.distributed import get_local_torch_device +from trainer.logger import init_logger + +logger = init_logger(__name__) + +# use system-level temp directory for file locks, so that multiple users +# can share the same lock without error. +# lock files in the temp directory will be automatically deleted when the +# system reboots, so users will not complain about annoying lock files +temp_dir = tempfile.gettempdir() + + +def enable_hf_transfer() -> None: + """automatically activates hf_transfer + """ + if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: + try: + # enable hf hub transfer if available + import hf_transfer # type: ignore # noqa + huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True + except ImportError: + pass + + +enable_hf_transfer() + + +class DisabledTqdm(tqdm): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, disable=True) + + +def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): + lock_dir = cache_dir or temp_dir + model_name_or_path = str(model_name_or_path) + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +# For models like Mistral-7B-v0.3, there are both sharded +# safetensors files and a consolidated safetensors file. +# Passing both of these to the weight loader functionality breaks. +# So, we use the index_file to +# look up which safetensors files should be used. +def filter_duplicate_safetensors_files(hf_weights_files: list[str], + hf_folder: str, + index_file: str) -> list[str]: + # model.safetensors.index.json is a mapping from keys in the + # torch state_dict to safetensors file holding that weight. + index_file_name = os.path.join(hf_folder, index_file) + if not os.path.isfile(index_file_name): + return hf_weights_files + + # Iterate through the weight_map (weight_name: safetensors files) + # to identify weights that we should use. + with open(index_file_name) as f: + weight_map = json.load(f)["weight_map"] + weight_files_in_index = set() + for weight_name in weight_map: + weight_files_in_index.add( + os.path.join(hf_folder, weight_map[weight_name])) + # Filter out any fields that are not found in the index file. + hf_weights_files = [ + f for f in hf_weights_files if f in weight_files_in_index + ] + return hf_weights_files + + +def filter_files_not_needed_for_inference( + hf_weights_files: list[str]) -> list[str]: + """ + Exclude files that are not needed for inference. + + See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 + """ + blacklist = [ + "training_args.bin", + "optimizer.bin", + "optimizer.pt", + "scheduler.pt", + "scaler.pt", + ] + hf_weights_files = [ + f for f in hf_weights_files + if not any(f.endswith(x) for x in blacklist) + ] + return hf_weights_files + + +# explicitly use pure text format, with a newline at the end +# this makes it impossible to see the animation in the progress bar +# but will avoid messing up with ray or multiprocessing, which wraps +# each line of output with some prefix. +_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 + + +def safetensors_weights_iterator( + hf_weights_files: list[str], + to_cpu: bool = True, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files.""" + enable_tqdm = not torch.distributed.is_initialized( + ) or torch.distributed.get_rank() == 0 + device = "cpu" if to_cpu else str(get_local_torch_device()) + for st_file in tqdm( + hf_weights_files, + desc="Loading safetensors checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + with safe_open(st_file, framework="pt", device=device) as f: + for name in f.keys(): # noqa: SIM118 + param = f.get_tensor(name) + yield name, param + + +def pt_weights_iterator( + hf_weights_files: list[str], + to_cpu: bool = True, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model bin/pt files.""" + device = "cpu" if to_cpu else str(get_local_torch_device()) + enable_tqdm = not torch.distributed.is_initialized( + ) or torch.distributed.get_rank() == 0 + for bin_file in tqdm( + hf_weights_files, + desc="Loading pt checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + state = torch.load(bin_file, map_location=device, weights_only=True) + yield from state.items() + del state + + +def default_weight_loader(param: torch.Tensor, + loaded_weight: torch.Tensor) -> None: + """Default weight loader.""" + try: + if param.numel() == 1 and loaded_weight.numel() == 1: + # Sometimes scalar values aren't considered tensors with shapes + # so if both param and loaded_weight are a scalar, + # "broadcast" instead of copy + param.data.fill_(loaded_weight.item()) + else: + assert param.size() == loaded_weight.size(), ( + f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})") + + param.data.copy_(loaded_weight) + except Exception: + # NOTE: This exception is added for the purpose of setting breakpoint to + # debug weight loading issues. + raise + + +def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: + """Remap the name of FP8 k/v_scale parameters. + + This function handles the remapping of FP8 k/v_scale parameter names. + It detects if the given name ends with a suffix and attempts to remap + it to the expected name format in the model. If the remapped name is not + found in the params_dict, a warning is printed and None is returned. + + Args: + name (str): The original loaded checkpoint parameter name. + params_dict (dict): Dictionary containing the model's named parameters. + + Returns: + str: The remapped parameter name if successful, or the original name + if no remapping is needed. + None: If the remapped name is not found in params_dict. + """ + if name.endswith(".kv_scale"): + logger.warning_once( + "DEPRECATED. Found kv_scale in the checkpoint. " + "This format is deprecated in favor of separate k_scale and " + "v_scale tensors and will be removed in a future release. " + "Functionally, we will remap kv_scale to k_scale and duplicate " + "k_scale to v_scale") + # NOTE: we remap the deprecated kv_scale to k_scale + remapped_name = name.replace(".kv_scale", ".attn.k_scale") + if remapped_name not in params_dict: + logger.warning_once( + f"Found kv_scale in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). kv_scale is " + "not loaded.") + return None + return remapped_name + + possible_scale_names = [".k_scale", ".v_scale"] + modelopt_scale_names = [ + ".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale" + ] + for scale_name in possible_scale_names: + if name.endswith(scale_name): + if any(mo_scale_name in name + for mo_scale_name in modelopt_scale_names): + remapped_name = name.replace( + f".self_attn.{scale_name[1]}_proj{scale_name}", + f".self_attn.attn{scale_name}") + else: + remapped_name = name.replace(scale_name, f".attn{scale_name}") + if remapped_name not in params_dict: + logger.warning_once( + f"Found {scale_name} in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). {scale_name} is " + "not loaded.") + return None + return remapped_name + + # If there were no matches, return the untouched param name + return name diff --git a/trainer/models/par_vae/context_parallel/plugins.py b/trainer/models/par_vae/context_parallel/plugins.py new file mode 100644 index 0000000000000000000000000000000000000000..1c3675f7ccf07851fb4378ecc092600efd8f0a01 --- /dev/null +++ b/trainer/models/par_vae/context_parallel/plugins.py @@ -0,0 +1,345 @@ +import torch +import torch.distributed as dist +import math + +import time + +class ModulePlugin: + def __init__(self, module, module_id, global_state=None): + self.module = module + self.module_id = module_id + self.global_state = global_state + self.enable = True + self.implement_forward() + + @property + def is_log_node(self): + return self.global_state.get('dist_controller').rank == 0 and self.module_id[1] == 0 + + @property + def t(self): + return self.global_state.get('timestep') + + @property + def p(self): + return self.t / 1000 + + def implement_forward(self): + module = self.module + if not hasattr(module, "old_forward"): + module.old_forward = module.forward + self.new_forward = self.get_new_forward() + def forward(*args, **kwargs): + self.update_config() # update config + return self.new_forward(*args, **kwargs) if self.enable else self.old_forward(*args, **kwargs) + module.forward = forward + + def set_enable(self, enable=True): + self.enable = enable + + def get_new_forward(self): + raise NotImplementedError + + def update_config(self, config:dict=None): + if config is None: + config = self.global_state.get('plugin_configs', {}).get(self.module_id[0], {}) + for key, value in config.items(): + setattr(self, key, value) + + +class GroupNormPlugin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + def get_new_forward(self): + module = self.module + + def new_forward(x): + shape = x.shape + N, C, G = shape[0], shape[1], module.num_groups + assert C % G == 0 + + x = x.reshape(N, G, -1) + + mean = x.mean(-1, keepdim=True).to(torch.float32) + dist.all_reduce(mean) + + mean = mean / dist.get_world_size() + + var = ((x - mean.to(x.dtype)) ** 2).mean(-1, keepdim=True).to(torch.float32) + + dist.all_reduce(var) + var = var / dist.get_world_size() + + x = (x - mean.to(x.dtype)) / (var.to(x.dtype) + module.eps).sqrt() + x = x.view(shape) + + new_shape = [1 for _ in shape] + new_shape[1] = -1 + + return x * module.weight.view(new_shape) + module.bias.view(new_shape) + + return new_forward + + +class Conv3DSafeNewPligin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, 'kernel_size', (1, 1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size, self.kernel_size) + + self.chunk_dim = global_state.get("chunk_dim") + + kernel_width = self.kernel_size[2] + d = kernel_width - 1 + self.padding_left = d // 2 + self.padding_right = d - self.padding_left + self.padding_flag = self.padding_left if d > 0 else 0 + + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.adj_groups = self.global_state.get('dist_controller').adj_groups + + + def pad_context(self, h): + if self.padding_flag == 0: + return h + + if self.chunk_dim == 3: + share_to_left = h[:, :, :, :self.padding_left].contiguous() + share_to_right = h[:, :, :, -self.padding_right:].contiguous() + else: + share_to_left = h[:, :, :, :, :self.padding_left].contiguous() + share_to_right = h[:, :, :, :, -self.padding_right:].contiguous() + + if self.rank % 2: + # 1. the rank is odd, pad the left first + if self.rank: + # not the first rank, have left context + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=self.adj_groups[self.rank-1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to(h.device, non_blocking=True) + # 2. then pad the right + if self.rank != dist.get_world_size() - 1: + # not the last rank, have right context + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=self.adj_groups[self.rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to(h.device, non_blocking=True) + else: + # 1. the rank is even, pad the right first + if self.rank != dist.get_world_size() - 1: + # not the last rank, have right context + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=self.adj_groups[self.rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to(h.device, non_blocking=True) + # 2. then pad the left + if self.rank: + # not the first rank, have left context + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=self.adj_groups[self.rank-1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to(h.device, non_blocking=True) + # torch.cuda.synchronize() + + h_with_context = torch.cat([left_context, h, right_context], dim=self.chunk_dim) + return h_with_context + + def get_new_forward(self): + module = self.module + def new_forward(hidden_states, cache_x=None, *args, **kwargs): + if self.padding_flag == 0: + # print(f"padding=0, return old_forward") + return module.old_forward(hidden_states, cache_x, *args, **kwargs) + + hidden_states = self.pad_context(hidden_states) + if cache_x is not None: + cache_x = self.pad_context(cache_x) + + result = module.old_forward(hidden_states, cache_x, *args, **kwargs) + if self.chunk_dim == 3: + result = result[:,:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + else: + result = result[:,:,:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + + return result + + return new_forward + +class Conv2DSafeNewPligin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, 'kernel_size', (1, 1)) + self.stride = getattr(module, 'stride', (1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size) + + self.chunk_dim = global_state.get("chunk_dim") + pad_kernel_dim = self.chunk_dim - 3 + kernel_height = self.kernel_size[pad_kernel_dim] # 卷积核的高度维度 + d = kernel_height - 1 # 总padding量 + self.padding_left = d // 2 # 上侧padding + self.padding_right = d - self.padding_left # 下侧padding + self.padding = self.padding_left if d > 0 else 0 + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.adj_groups = self.global_state.get('dist_controller').adj_groups + + def pad_context(self, h): + if self.padding == 0: + return h + + if self.chunk_dim == 3: + share_to_left = h[:, :, :self.padding_left].contiguous() + share_to_right = h[:, :, -self.padding_right:].contiguous() + else: + share_to_left = h[:, :, :, :self.padding_left].contiguous() + share_to_right = h[:, :, :, -self.padding_right:].contiguous() + + if self.rank % 2: + # 1. the rank is odd, pad the left first + if self.rank: + # not the first rank, have left context + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=self.adj_groups[self.rank-1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to(h.device, non_blocking=True) + # 2. then pad the right + if self.rank != dist.get_world_size() - 1: + # not the last rank, have right context + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=self.adj_groups[self.rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to(h.device, non_blocking=True) + else: + # 1. the rank is even, pad the right first + if self.rank != dist.get_world_size() - 1: + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=self.adj_groups[self.rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to(h.device, non_blocking=True) + # 2. then pad the left + if self.rank: + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=self.adj_groups[self.rank-1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to(h.device, non_blocking=True) + # torch.cuda.synchronize() + + h_with_context = torch.cat([left_context, h, right_context], dim=self.chunk_dim - 1) + return h_with_context + + def get_new_forward(self): + module = self.module + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + if self.padding == 0: + return module.old_forward(hidden_states) + + hidden_states = self.pad_context(hidden_states) + + if self.chunk_dim == 3: + hidden_states = module.old_forward(hidden_states)[:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + else: + hidden_states = module.old_forward(hidden_states)[:,:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + return hidden_states + + return new_forward + +class Conv2DSafeNewPliginStride2(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, 'kernel_size', (1, 1)) + self.stride = getattr(module, 'stride', (1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size) + + kernel_height = self.kernel_size[0] + d = kernel_height - 1 + self.padding_left = d // 2 + self.padding_right = d - self.padding_left + self.padding = self.padding_left if d > 0 else 0 + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.adj_groups = self.global_state.get('dist_controller').adj_groups + + def pad_context(self, h): + if self.padding == 0: + return h + + share_to_left = h[:, :, :self.padding_left].contiguous() + + if self.rank < dist.get_world_size() - 1: + right_context = torch.zeros_like(share_to_left) + + dist.recv(right_context, src=self.rank+1) + if self.rank >0: + dist.send(share_to_left, dst=self.rank-1) + # torch.cuda.synchronize() + if self.rank < dist.get_world_size() - 1: + h_with_context = torch.cat([h, right_context], dim=2) + else: + h_with_context = h + return h_with_context + + def get_new_forward(self): + module = self.module + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + assert self.global_state["chunk_dim"] == 3 # only support chunk height now + if self.padding == 0: + return module.old_forward(hidden_states) + + hidden_states = hidden_states[:, :, :-1, :] + hidden_states = self.pad_context(hidden_states) + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, 1)) + hidden_states = module.old_forward(hidden_states)#[:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + return hidden_states + + return new_forward + +class WanAttentionPlugin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.gather_dim = global_state.get("chunk_dim") + + super().__init__(module, module_id, global_state) + + def get_new_forward(self): + module = self.module + rank = self.rank + world_size = self.world_size + + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + gathered_tensors = [torch.zeros_like(hidden_states) for _ in range(world_size)] + dist.all_gather(gathered_tensors, hidden_states) + + combined_tensor = torch.cat(gathered_tensors, dim=self.gather_dim) + forward_output = module.old_forward(combined_tensor) + + #chunk_sizes = [t.size(self.gather_dim) for t in gathered_tensors] + #start_idx = sum(chunk_sizes[:rank]) + #end_idx = start_idx + chunk_sizes[rank] + #local_output = forward_output[:, :, :, start_idx:end_idx].contiguous() + local_output = forward_output.chunk(chunks=world_size, dim=self.gather_dim)[rank] + + return local_output + + return new_forward + diff --git a/trainer/models/par_vae/context_parallel/wrapper_vae.py b/trainer/models/par_vae/context_parallel/wrapper_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..a5ed2add2642fc85bd089acc2c695d38ffa86f83 --- /dev/null +++ b/trainer/models/par_vae/context_parallel/wrapper_vae.py @@ -0,0 +1,174 @@ +from ..tools import GlobalState, DistController + +from .plugins import torch, ModulePlugin, GroupNormPlugin, Conv3DSafeNewPligin, Conv2DSafeNewPligin, WanAttentionPlugin, Conv2DSafeNewPliginStride2 +from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanAttentionBlock +#from ...modules.vae import CausalConv3d, AttentionBlock +#from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanAttentionBlock + +class DistWrapper(object): + def __init__(self, pipe, dist_controller: DistController, config, chunk_dim) -> None: + super().__init__() + self.pipe = pipe + self.dist_controller = dist_controller + self.config = config + assert chunk_dim == 3 or chunk_dim == 4 + self.global_state = GlobalState({ + "dist_controller": dist_controller, + "chunk_dim": chunk_dim + }) + self.plugin_mount() + + plugin_configs={ + "attn":{ + "padding": 24, + "top_k": 24, + "top_k_chunk_size": 24, + "attn_scale": 1., + "token_num_scale": True, + "dynamic_scale": True, + }, + "conv_3d": { + "padding": 1, + }, + "conv_layer": {}, + } + self.global_state.set("plugin_configs", plugin_configs) + + # torch.compile + #torch._dynamo.config.recompile_limit = 20 + with torch.no_grad(): + #self.pipe.model.encoder = torch.compile(self.pipe.model.encoder) + self.pipe.decoder = torch.compile(self.pipe.decoder) + + + def plugin_mount(self): + self.plugins = {} + self.group_norm_plugin_mount() + self.conv_3d_plugin_mount() + self.conv_2d_plugin_stride2_mount() ##only for wan vae encoder + self.conv_2d_plugin_mount() + self.wanattention_plugin_mount() + + def wanattention_plugin_mount(self): + self.plugins['wanattention'] = {} + wanattention_s = [] + for module in self.pipe.encoder.named_modules(): + #if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock': + # print("Encoder attn: ", module[0]) + if ('mid_block.' in module[0] and module[1].__class__.__name__ == 'WanAttentionBlock'): + wanattention_s.append(module[1]) + for module in self.pipe.decoder.named_modules(): + #if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock': + # print("Decoder attn: ", module[0]) + if ('mid_block.' in module[0] and module[1].__class__.__name__ == 'WanAttentionBlock'): + wanattention_s.append(module[1]) + if self.dist_controller.is_master: + print(f'Found {len(wanattention_s)} wanattention_s') + for i, wanattention in enumerate(wanattention_s): + plugin_id = 'wanattention', i + self.plugins['wanattention'][plugin_id] = WanAttentionPlugin(wanattention, plugin_id, self.global_state) + + def group_norm_plugin_mount(self): + self.plugins['group_norm'] = {} + group_norms = [] + for module in self.pipe.decoder.named_modules(): + if ('norm_layer' in module[0]) and module[1].__class__.__name__ == 'GroupNorm': + group_norms.append(module[1]) + if self.dist_controller.is_master: + print(f'Found {len(group_norms)} group norms') + for i, group_norm in enumerate(group_norms): + plugin_id = 'group_norm', i + self.plugins['group_norm'][plugin_id] = GroupNormPlugin(group_norm, plugin_id, self.global_state) + + def conv_3d_plugin_mount(self): + self.plugins['conv_3d'] = {} + conv3d_s = [] + for module in self.pipe.encoder.named_modules(): + #if isinstance(module[1], CausalConv3d): + # print("Encoder conv3d: ", module[0], module[1].kernel_size[1]) + if (isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1): + # print(f"Found conv3d: {module[1]}") + conv3d_s.append(module[1]) + for module in self.pipe.decoder.named_modules(): + #if isinstance(module[1], CausalConv3d): + # print("Decoder conv3d: ", module[0], module[1].kernel_size[1]) + if (isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1): + # print(f"Found conv3d: {module[1]}") + conv3d_s.append(module[1]) + if self.dist_controller.is_master: + print(f'Found {len(conv3d_s)} conv3d_s') + for i, conv in enumerate(conv3d_s): + plugin_id = 'conv_3d', i + self.plugins['conv_3d'][plugin_id] = Conv3DSafeNewPligin(conv, plugin_id, self.global_state) + + def conv_2d_plugin_stride2_mount(self): + self.plugins['conv_2d_stride2'] = {} + conv2d_stride2_s = [] + for module in self.pipe.encoder.named_modules(): + if ('.resample' in module[0] and module[1].__class__.__name__ == 'Conv2d'): + conv2d_stride2_s.append(module[1]) + if self.dist_controller.is_master: + print(f'Found {len(conv2d_stride2_s)} conv2d_stride2_s') + for i, conv in enumerate(conv2d_stride2_s): + plugin_id = 'conv_2d_stride2', i + self.plugins['conv_2d_stride2'][plugin_id] = Conv2DSafeNewPliginStride2(conv, plugin_id, self.global_state) + + def conv_2d_plugin_mount(self): + self.plugins['conv_2d'] = {} + conv2d_s = [] + for module in self.pipe.decoder.named_modules(): + if ('.resample' in module[0] and module[1].__class__.__name__ == 'Conv2d'): + conv2d_s.append(module[1]) + if self.dist_controller.is_master: + print(f'Found {len(conv2d_s)} conv2d_s') + for i, conv in enumerate(conv2d_s): + plugin_id = 'conv_2d', i + self.plugins['conv_2d'][plugin_id] = Conv2DSafeNewPligin(conv, plugin_id, self.global_state) + + def inference( + self, + local_pose_image, + local_latents, + #prompts="A beaglening diving goggles swimming in the ocean while the camera is moving, coral reefs in the background", + config={}, + pipe_configs={ + "steps": 50, + "guidance_scale": 12, + "fps": 60, + "num_frames": 24 * 1, + "height": 320, + "width": 512, + "export_fps": 12, + "base_path": "./work/output", + "file_name": None + }, + plugin_configs={ + "attn":{ + "padding": 24, + "top_k": 24, + "top_k_chunk_size": 24, + "attn_scale": 1., + "token_num_scale": True, + "dynamic_scale": True, + }, + "conv_3d": { + "padding": 1, + }, + "conv_layer": {}, + }, + additional_info={}, + ): + self.plugin_mount() + # print("self.config seed: ", self.config["seed"]) + + self.global_state.set("plugin_configs", plugin_configs) + self.pipe = self.pipe.to(device='cuda', dtype=torch.bfloat16) + with torch.no_grad(): + local_pose_image = local_pose_image.to(device='cuda', dtype=torch.bfloat16) + local_latents = local_latents.to(device='cuda', dtype=torch.bfloat16) + + tmp_latents = self.pipe.encode(local_pose_image).latent_dist.mode() + + latents = self.pipe.decode(local_latents, return_dict=False)[0] + return latents + diff --git a/trainer/models/par_vae/tools.py b/trainer/models/par_vae/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..a0db67fe1ab5103effa5cdb99a68ec5f8a1e3a99 --- /dev/null +++ b/trainer/models/par_vae/tools.py @@ -0,0 +1,102 @@ +import json +import numpy as np +import imageio +import os + +import torch +import torch.distributed as dist + +def export_to_video(video_frames, output_video_path, fps = 12): + # Ensure all frames are NumPy arrays and determine video dimensions from the first frame + assert all(isinstance(frame, np.ndarray) for frame in video_frames), "All video frames must be NumPy arrays." + # Ensure output_video_path is ending with .mp4 + if not output_video_path.endswith('.mp4'): + output_video_path += '.mp4' + # Create a video file at the specified path and write frames to it + with imageio.get_writer(output_video_path, fps=fps, format='mp4') as writer: + for frame in video_frames: + writer.append_data( + (frame * 255).astype(np.uint8) + ) + +def save_generation(video_frames, configs, base_path, file_name=None): + if not os.path.exists(base_path): + os.makedirs(base_path) + p_config = configs["pipe_configs"] + frames, steps, fps = p_config["num_frames"], p_config["steps"], p_config["fps"] + if not file_name: + index = [int(each.split('_')[0]) for each in os.listdir(base_path)] + max_idex = max(index) if index else 0 + idx_str = str(max_idex + 1).zfill(6) + + + key_info = '_'.join([str(frames), str(steps), str(fps)]) + file_name = f'{idx_str}_{key_info}' + + with open(f'{base_path}/{file_name}.json', 'w') as f: + json.dump(configs, f, indent=4) + + export_to_video(video_frames, f'{base_path}/{file_name}.mp4', fps=p_config["export_fps"]) + + return file_name + + +class GlobalState: + def __init__(self, state={}) -> None: + self.init_state(state) + + def init_state(self, state={}): + self.state = state + + def set(self, key, value): + self.state[key] = value + + def get(self, key, default=None): + return self.state.get(key, default) + + +class DistController(object): + def __init__(self, rank, world_size, config = None) -> None: + super().__init__() + self.rank = rank + self.world_size = world_size + self.config = config + self.is_master = (rank == 0) + #print("DistController is master: ", self.is_master) + #self.init_dist() + self.init_group() + #self.device = torch.device(f"cuda:{config['devices'][dist.get_rank()]}") + self.device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(self.device) + + def init_dist(self): + #print(f"Rank {self.rank} is running.") + os.environ['MASTER_ADDR'] = '127.0.0.1' + os.environ['MASTER_PORT'] = str(self.config.get("master_port") or "29500") + dist.init_process_group("nccl", rank=self.rank, world_size=self.world_size) + + def init_group(self): + self.adj_groups = [dist.new_group([i, i+1]) for i in range(self.world_size-1)] + + +def distribute_data_to_gpus(data, dim, rank, num_gpus, dtype): + dim_size = data.shape[dim] + assert dim_size % num_gpus == 0, f"Dim {dim_size} of data is not a multiply of {num_gpus}" + + chunk_size = dim_size // num_gpus + + slices = [slice(None)] * len(data.shape) + slices[dim] = slice(rank * chunk_size, (rank + 1) * chunk_size) + + local_data = data[slices].to(device=f'cuda:{rank}', dtype=dtype) + print(f"Rank {rank} local_data: {local_data.shape}") + return local_data + +def gather_results(local_result, rank, dim, world_size): + dist.barrier() + + gathered_results = [torch.zeros_like(local_result) for _ in range(world_size)] + dist.all_gather(gathered_results, local_result) + + gathered_result = torch.cat(gathered_results, dim=dim) + return gathered_result diff --git a/trainer/models/parameter.py b/trainer/models/parameter.py new file mode 100644 index 0000000000000000000000000000000000000000..1b9bec8c0c7faf5255ed3662dcb459fa218b2ed1 --- /dev/null +++ b/trainer/models/parameter.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py + +from collections.abc import Callable +from fractions import Fraction +from typing import Any + +import torch +from torch.nn import Parameter + +from trainer.distributed import get_tp_rank +from trainer.logger import init_logger +from trainer.models.utils import _make_synced_weight_loader + +logger = init_logger(__name__) + + +class BasevLLMParameter(Parameter): + """ + Base parameter for vLLM linear layers. Extends the torch.nn.parameter + by taking in a linear weight loader. Will copy the loaded weight + into the parameter when the provided weight loader is called. + """ + + def __new__(cls, data: torch.Tensor, **kwargs): + + return super().__new__(cls, data=data, requires_grad=False) + + def __init__(self, data: torch.Tensor, weight_loader: Callable): + """ + Initialize the BasevLLMParameter + + :param data: torch tensor with the parameter data + :param weight_loader: weight loader callable + + :returns: a torch.nn.parameter + """ + + # During weight loading, we often do something like: + # narrowed_tensor = param.data.narrow(0, offset, len) + # narrowed_tensor.copy_(real_weight) + # expecting narrowed_tensor and param.data to share the same storage. + # However, on TPUs, narrowed_tensor will lazily propagate to the base + # tensor, which is param.data, leading to the redundant memory usage. + # This sometimes causes OOM errors during model loading. To avoid this, + # we sync the param tensor after its weight loader is called. + from trainer.platforms import current_platform + if current_platform.is_tpu(): + weight_loader = _make_synced_weight_loader(weight_loader) + + self._weight_loader = weight_loader + + @property + def weight_loader(self): + return self._weight_loader + + def _is_1d_and_scalar(self, loaded_weight: torch.Tensor): + cond1 = self.data.ndim == 1 and self.data.numel() == 1 + cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1 + return (cond1 and cond2) + + def _assert_and_load(self, loaded_weight: torch.Tensor) -> None: + assert (self.data.shape == loaded_weight.shape + or self._is_1d_and_scalar(loaded_weight)) + self.data.copy_(loaded_weight) + + def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + self._assert_and_load(loaded_weight) + + def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + self._assert_and_load(loaded_weight) + + def load_merged_column_weight(self, loaded_weight: torch.Tensor, + **kwargs) -> None: + self._assert_and_load(loaded_weight) + + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: + self._assert_and_load(loaded_weight) + + +class _ColumnvLLMParameter(BasevLLMParameter): + """ + Private class defining weight loading functionality + (load_merged_column_weight, load_qkv_weight) + for parameters being loaded into linear layers with column + parallelism. This includes QKV and MLP layers which are + not already fused on disk. Requires an output dimension + to be defined. Called within the weight loader of + each of the column parallel linear layers. + """ + + def __init__(self, output_dim: int, **kwargs): + self._output_dim = output_dim + super().__init__(**kwargs) + + @property + def output_dim(self): + return self._output_dim + + def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tp_rank() + shard_size = self.data.shape[self.output_dim] + loaded_weight = loaded_weight.narrow(self.output_dim, + tp_rank * shard_size, shard_size) + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + def load_merged_column_weight(self, loaded_weight: torch.Tensor, + **kwargs) -> None: + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + if shard_offset is None or shard_size is None: + raise ValueError("shard_offset and shard_size must be provided") + if isinstance( + self, PackedColumnParameter + | PackedvLLMParameter) and self.packed_dim == self.output_dim: + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size) + + param_data = self.data + + tp_rank = get_tp_rank() + param_data = param_data.narrow(self.output_dim, shard_offset, + shard_size) + loaded_weight = loaded_weight.narrow(self.output_dim, + tp_rank * shard_size, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + shard_id = kwargs.get("shard_id") + num_heads = kwargs.get("num_heads") + + assert shard_offset is not None + assert shard_size is not None + assert shard_id is not None + assert num_heads is not None + + if isinstance( + self, PackedColumnParameter + | PackedvLLMParameter) and self.output_dim == self.packed_dim: + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size) + + param_data = self.data + tp_rank = get_tp_rank() + shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads + param_data = param_data.narrow(self.output_dim, shard_offset, + shard_size) + loaded_weight = loaded_weight.narrow(self.output_dim, + shard_id * shard_size, shard_size) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class RowvLLMParameter(BasevLLMParameter): + """ + Parameter class defining weight_loading functionality + (load_row_parallel_weight) for parameters being loaded + into linear layers with row parallel functionality. + Requires an input_dim to be defined. + """ + + def __init__(self, input_dim: int, **kwargs): + self._input_dim = input_dim + super().__init__(**kwargs) + + @property + def input_dim(self): + return self._input_dim + + def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tp_rank() + shard_size = self.data.shape[self.input_dim] + loaded_weight = loaded_weight.narrow(self.input_dim, + tp_rank * shard_size, shard_size) + + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + +class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for linear layer weights. Uses both column and + row parallelism. + """ + pass + + +class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + grouped quantization. Uses both column and row parallelism. + """ + pass + + +class ChannelQuantScaleParameter(_ColumnvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + channel-wise quantization. Equivalent to _ColumnvLLMParameter. + """ + pass + + +class PerTensorScaleParameter(BasevLLMParameter): + """ + Parameter class for scales where the number of scales is + equivalent to the number of logical matrices in fused linear + layers (e.g. for QKV, there are 3 scales loaded from disk). + This is relevant to weights with per-tensor quantization. + Adds functionality to map the scalers to a shard during + weight loading. + + Note: additional parameter manipulation may be handled + for each quantization config specifically, within + process_weights_after_loading + """ + + def __init__(self, **kwargs): + self.qkv_idxs = {"q": 0, "k": 1, "v": 2} + super().__init__(**kwargs) + + def _shard_id_as_int(self, shard_id: str | int) -> int: + if isinstance(shard_id, int): + return shard_id + + # if not int, assume shard_id for qkv + # map to int and return + assert isinstance(shard_id, str) + assert shard_id in self.qkv_idxs + return self.qkv_idxs[shard_id] + + # For row parallel layers, no sharding needed + # load weight into parameter as is + def load_row_parallel_weight(self, *args, **kwargs) -> None: + super().load_row_parallel_weight(*args, **kwargs) + + def load_merged_column_weight(self, *args, **kwargs) -> None: + self._load_into_shard_id(*args, **kwargs) + + def load_qkv_weight(self, *args, **kwargs) -> None: + self._load_into_shard_id(*args, **kwargs) + + def load_column_parallel_weight(self, *args, **kwargs) -> None: + super().load_row_parallel_weight(*args, **kwargs) + + def _load_into_shard_id(self, loaded_weight: torch.Tensor, + shard_id: str | int, **kwargs): + """ + Slice the parameter data based on the shard id for + loading. + """ + + param_data = self.data + shard_id = self._shard_id_as_int(shard_id) + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + param_data = param_data[shard_id] + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class PackedColumnParameter(_ColumnvLLMParameter): + """ + Parameter for model parameters which are packed on disk + and support column parallelism only. See PackedvLLMParameter + for more details on the packed properties. + """ + + def __init__(self, packed_factor: int | Fraction, packed_dim: int, + **kwargs): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + def adjust_shard_indexes_for_packing(self, shard_size, + shard_offset) -> tuple[Any, Any]: + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor) + + +class PackedvLLMParameter(ModelWeightParameter): + """ + Parameter for model weights which are packed on disk. + Example: GPTQ Marlin weights are int4 or int8, packed into int32. + Extends the ModelWeightParameter to take in the + packed factor, the packed dimension, and optionally, marlin + tile size for marlin kernels. Adjusts the shard_size and + shard_offset for fused linear layers model weight loading + by accounting for packing and optionally, marlin tile size. + """ + + def __init__(self, packed_factor: int | Fraction, packed_dim: int, + **kwargs): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor) + + +class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """ + Parameter class for weight scales loaded for weights with + block-wise quantization. Uses both column and row parallelism. + """ + + pass + + +def permute_param_layout_(param: BasevLLMParameter, input_dim: int, + output_dim: int, **kwargs) -> BasevLLMParameter: + """ + Permute a parameter's layout to the specified input and output dimensions, + useful for forcing the parameter into a known layout, for example, if I need + a packed (quantized) weight matrix to be in the layout + {input_dim = 0, output_dim = 1, packed_dim = 0} + then I can call: + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + to ensure x is in the correct layout (permuting it to the correct layout if + required, asserting if it cannot get it to the correct layout) + """ + + curr_input_dim = getattr(param, "input_dim", None) + curr_output_dim = getattr(param, "output_dim", None) + + if curr_input_dim is None or curr_output_dim is None: + assert param.data.dim() == 2,\ + "permute_param_layout_ only supports 2D parameters when either "\ + "input_dim or output_dim is not set" + + # if one of the dimensions is not set, set it to the opposite of the other + # we can only do this since we asserted the parameter is 2D above + if curr_input_dim is None: + assert curr_output_dim is not None,\ + "either input or output dim must be set" + curr_input_dim = (curr_output_dim + 1) % 2 + if curr_output_dim is None: + assert curr_input_dim is not None,\ + "either input or output dim must be set" + curr_output_dim = (curr_input_dim + 1) % 2 + + # create permutation from the current layout to the layout with + # self.input_dim at input_dim and self.output_dim at output_dim preserving + # other dimensions + perm = [ + i for i in range(param.data.dim()) + if i not in [curr_input_dim, curr_output_dim] + ] + perm.insert(input_dim, curr_input_dim) + perm.insert(output_dim, curr_output_dim) + + if "packed_dim" in kwargs: + assert hasattr(param, "packed_dim") and\ + param.packed_dim == perm[kwargs["packed_dim"]],\ + "permute_param_layout_ currently doesn't support repacking" + + param.data = param.data.permute(*perm) + if hasattr(param, "_input_dim"): + param._input_dim = input_dim + if hasattr(param, "_output_dim"): + param._output_dim = output_dim + if "packed_dim" in kwargs and hasattr(param, "_packed_dim"): + param._packed_dim = kwargs["packed_dim"] + + return param + + +def _adjust_shard_indexes_for_packing(shard_size, shard_offset, + packed_factor) -> tuple[Any, Any]: + shard_size = shard_size // packed_factor + shard_offset = shard_offset // packed_factor + return shard_size, shard_offset diff --git a/trainer/models/prope/camera_rope.py b/trainer/models/prope/camera_rope.py new file mode 100644 index 0000000000000000000000000000000000000000..86e32a306a42a5077e8ed337c08c26c11aa0ce84 --- /dev/null +++ b/trainer/models/prope/camera_rope.py @@ -0,0 +1,259 @@ +# MIT License +# +# Copyright (c) Authors of +# "PRoPE: Projective Positional Encoding for Multiview Transformers" +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# How to use PRoPE attention for self-attention: +# +# 1. Easiest way (fast): +# attn = PropeDotProductAttention(...) +# o = attn(q, k, v, viewmats, Ks) +# +# 2. More flexible way (fast): +# attn = PropeDotProductAttention(...) +# attn._precompute_and_cache_apply_fns(viewmats, Ks) +# q = attn._apply_to_q(q) +# k = attn._apply_to_kv(k) +# v = attn._apply_to_kv(v) +# o = F.scaled_dot_product_attention(q, k, v, **kwargs) +# o = attn._apply_to_o(o) +# +# 3. The most flexible way (but slower because repeated computation of RoPE coefficients): +# o = prope_dot_product_attention(q, k, v, ...) +# +# How to use PRoPE attention for cross-attention: +# +# attn_src = PropeDotProductAttention(...) +# attn_tgt = PropeDotProductAttention(...) +# attn_src._precompute_and_cache_apply_fns(viewmats_src, Ks_src) +# attn_tgt._precompute_and_cache_apply_fns(viewmats_tgt, Ks_tgt) +# q_src = attn_src._apply_to_q(q_src) +# k_tgt = attn_tgt._apply_to_kv(k_tgt) +# v_tgt = attn_tgt._apply_to_kv(v_tgt) +# o_src = F.scaled_dot_product_attention(q_src, k_tgt, v_tgt, **kwargs) +# o_src = attn_src._apply_to_o(o_src) + +from functools import partial +from typing import Callable, Optional, Tuple, List + +import torch +import torch.nn.functional as F + + +def prope_qkv( + q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + *, + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int=None, # How many patches wide is each image? + patches_y: int=None, # How many patches tall is each image? + image_width: int = None, # Width of the image. Used to normalize intrinsics. + image_height: int = None, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + kv_cache = None, + is_cache: bool = False, + **kwargs, +) -> torch.Tensor: + """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style + positional encoding. + + Currently, we assume that the sequence length is equal to: + + cameras * patches_x * patches_y + + And token ordering allows the `(seqlen,)` axis to be reshaped into + `(cameras, patches_x, patches_y)`. + """ + # We're going to assume self-attention: all inputs are the same shape. + (batch, num_heads, seqlen, head_dim) = q.shape + cameras = viewmats.shape[1] + assert q.shape == k.shape == v.shape + assert viewmats.shape == (batch, cameras, 4, 4) + assert Ks is None or Ks.shape == (batch, cameras, 3, 3) + # assert seqlen == cameras * patches_x * patches_y + + apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns_all_dim( + head_dim=head_dim, + viewmats=viewmats, + Ks=Ks, + patches_x=patches_x, + patches_y=patches_y, + image_width=image_width, + image_height=image_height, + coeffs_x=coeffs_x, + coeffs_y=coeffs_y, + ) + + query = apply_fn_q(q) + key = apply_fn_kv(k) + value = apply_fn_kv(v) + + return query, key, value, apply_fn_o + + +def _prepare_apply_fns_all_dim( + head_dim: int, # Q/K/V will have this last dimension + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int, # How many patches wide is each image? + patches_y: int, # How many patches tall is each image? + image_width: int, # Width of the image. Used to normalize intrinsics. + image_height: int, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[torch.Tensor] = None, + coeffs_y: Optional[torch.Tensor] = None, +) -> Tuple[ + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], +]: + """Prepare transforms for PRoPE-style positional encoding.""" + device = viewmats.device + (batch, cameras, _, _) = viewmats.shape + + # Normalize camera intrinsics. + if Ks is not None: + Ks_norm = torch.zeros_like(Ks) + Ks_norm[..., 0, 0] = Ks[..., 0, 0] + Ks_norm[..., 1, 1] = Ks[..., 1, 1] + Ks_norm[..., 0, 2] = 0 + Ks_norm[..., 1, 2] = 0 + Ks_norm[..., 2, 2] = 1.0 + Ks_norm = Ks_norm.to(dtype=Ks.dtype) + del Ks + + # Compute the camera projection matrices we use in PRoPE. + # - K is an `image<-camera` transform. + # - viewmats is a `camera<-world` transform. + # - P = lift(K) @ viewmats is an `image<-world` transform. + P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats) + P_T = P.transpose(-1, -2).to(dtype=viewmats.dtype) + P_inv = torch.einsum( + "...ij,...jk->...ik", + _invert_SE3(viewmats), + _lift_K(_invert_K(Ks_norm)), + ).to(dtype=viewmats.dtype) + + else: + # GTA formula. P is `camera<-world` transform. + P = viewmats + P_T = P.transpose(-1, -2) + P_inv = _invert_SE3(viewmats) + + assert P.shape == P_inv.shape == (batch, cameras, 4, 4) + + # Block-diagonal transforms to the inputs and outputs of the attention operator. + assert head_dim % 4 == 0 + transforms_q = [ + (partial(_apply_tiled_projmat, matrix=P_T), head_dim), + ] + transforms_kv = [ + (partial(_apply_tiled_projmat, matrix=P_inv), head_dim), + ] + transforms_o = [ + (partial(_apply_tiled_projmat, matrix=P), head_dim), + ] + + apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) + return apply_fn_q, apply_fn_kv, apply_fn_o + + +def _apply_tiled_projmat( + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + matrix: torch.Tensor, # (batch, cameras, D, D) +) -> torch.Tensor: + """Apply projection matrix to features.""" + # - seqlen => (cameras, patches_x * patches_y) + # - feat_dim => (feat_dim // 4, 4) + (batch, num_heads, seqlen, feat_dim) = feats.shape + cameras = matrix.shape[1] + assert seqlen >= cameras and seqlen % cameras == 0 + D = matrix.shape[-1] + assert matrix.shape == (batch, cameras, D, D) + assert feat_dim % D == 0 + return torch.einsum( + "bcij,bncpkj->bncpki", + matrix, + feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)), + ).reshape(feats.shape) + + +def _apply_block_diagonal( + feats: torch.Tensor, # (..., dim) + func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Apply a block-diagonal function to an input array. + + Each function is specified as a tuple with form: + + ((Tensor) -> Tensor, int) + + Where the integer is the size of the input to the function. + """ + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes) + x_blocks = torch.split(feats, block_sizes, dim=-1) + out = torch.cat( + [f(x_block) for f, x_block in zip(funcs, x_blocks)], + dim=-1, + ) + assert out.shape == feats.shape, "Input/output shapes should match." + return out + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix.""" + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + out = out.to(dtype=transforms.dtype) + return out + + +def _lift_K(Ks: torch.Tensor) -> torch.Tensor: + """Lift 3x3 matrices to homogeneous 4x4 matrices.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device) + out[..., :3, :3] = Ks + out[..., 3, 3] = 1.0 + out = out.to(dtype=Ks.dtype) + return out + + +def _invert_K(Ks: torch.Tensor) -> torch.Tensor: + """Invert 3x3 intrinsics matrices. Assumes no skew.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros_like(Ks) + out[..., 0, 0] = 1.0 / Ks[..., 0, 0] + out[..., 1, 1] = 1.0 / Ks[..., 1, 1] + out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0] + out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1] + out[..., 2, 2] = 1.0 + out = out.to(dtype=Ks.dtype) + return out \ No newline at end of file diff --git a/trainer/models/registry.py b/trainer/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..6a53373ee398ae0e6568ceb58bfca94d82b551b8 --- /dev/null +++ b/trainer/models/registry.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py + +import importlib +import os +import pickle +import subprocess +import sys +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Callable, Set +from dataclasses import dataclass, field +from functools import lru_cache +from typing import NoReturn, TypeVar, cast + +import cloudpickle +from torch import nn + +from trainer.logger import init_logger + +logger = init_logger(__name__) + +# huggingface class name: (component_name, trainer module name, trainer class name) +_TEXT_TO_VIDEO_DIT_MODELS = { + "HunyuanTransformer3DARActionModel": ("hyvideo", "models.transformers.ar_action_hunyuanvideo_1_5_transformer", "ARHunyuanVideo_1_5_DiffusionTransformer"), +} + +_IMAGE_TO_VIDEO_DIT_MODELS = { + +} + +_TEXT_ENCODER_MODELS = { + "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), + "LlamaModel": ("encoders", "llama", "LlamaModel"), + "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"), + "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"), + "BertModel": ("encoders", "clip", "CLIPTextModel"), +} + +_IMAGE_ENCODER_MODELS: dict[str, tuple] = { + # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), + "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), +} + +_VAE_MODELS = { + "AutoencoderKLHunyuanVideo": + ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), + "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), + "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo") +} + +_SCHEDULERS = { + "FlowMatchEulerDiscreteScheduler": + ("schedulers", "scheduling_flow_match_euler_discrete", + "FlowMatchEulerDiscreteScheduler"), + "UniPCMultistepScheduler": + ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), +} + +_FAST_VIDEO_MODELS = { + **_TEXT_TO_VIDEO_DIT_MODELS, + **_IMAGE_TO_VIDEO_DIT_MODELS, + **_TEXT_ENCODER_MODELS, + **_IMAGE_ENCODER_MODELS, + **_VAE_MODELS, + **_SCHEDULERS, +} + +_SUBPROCESS_COMMAND = [sys.executable, "-m", "trainer.models.dits.registry"] + +_T = TypeVar("_T") + + +@dataclass(frozen=True) +class _ModelInfo: + architecture: str + + @staticmethod + def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": + return _ModelInfo(architecture=model.__name__, ) + + +class _BaseRegisteredModel(ABC): + + @abstractmethod + def inspect_model_cls(self) -> _ModelInfo: + raise NotImplementedError + + @abstractmethod + def load_model_cls(self) -> type[nn.Module]: + raise NotImplementedError + + +@dataclass(frozen=True) +class _RegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has already been imported in the main process. + """ + + interfaces: _ModelInfo + model_cls: type[nn.Module] + + @staticmethod + def from_model_cls(model_cls: type[nn.Module]): + return _RegisteredModel( + interfaces=_ModelInfo.from_model_cls(model_cls), + model_cls=model_cls, + ) + + def inspect_model_cls(self) -> _ModelInfo: + return self.interfaces + + def load_model_cls(self) -> type[nn.Module]: + return self.model_cls + + +def _run_in_subprocess(fn: Callable[[], _T]) -> _T: + # NOTE: We use a temporary directory instead of a temporary file to avoid + # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file + with tempfile.TemporaryDirectory() as tempdir: + output_filepath = os.path.join(tempdir, "registry_output.tmp") + + # `cloudpickle` allows pickling lambda functions directly + input_bytes = cloudpickle.dumps((fn, output_filepath)) + + # cannot use `sys.executable __file__` here because the script + # contains relative imports + returned = subprocess.run(_SUBPROCESS_COMMAND, + input=input_bytes, + capture_output=True) + + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError(f"Error raised in subprocess:\n" + f"{returned.stderr.decode()}") from e + + with open(output_filepath, "rb") as f: + return cast(_T, pickle.load(f)) + + +@dataclass(frozen=True) +class _LazyRegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has not been imported in the main process. + """ + module_name: str + component_name: str + class_name: str + + # Performed in another process to avoid initializing CUDA + def inspect_model_cls(self) -> _ModelInfo: + return _run_in_subprocess( + lambda: _ModelInfo.from_model_cls(self.load_model_cls())) + + def load_model_cls(self) -> type[nn.Module]: + mod = importlib.import_module(self.module_name) + return cast(type[nn.Module], getattr(mod, self.class_name)) + + +@lru_cache(maxsize=128) +def _try_load_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> type[nn.Module] | None: + from trainer.platforms import current_platform + current_platform.verify_model_arch(model_arch) + try: + return model.load_model_cls() + except Exception: + logger.exception("Error in loading model architecture '%s'", model_arch) + return None + + +@lru_cache(maxsize=128) +def _try_inspect_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> _ModelInfo | None: + try: + return model.inspect_model_cls() + except Exception: + logger.exception("Error in inspecting model architecture '%s'", + model_arch) + return None + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) + + def get_supported_archs(self) -> Set[str]: + return self.models.keys() + + def register_model( + self, + model_arch: str, + model_cls: type[nn.Module] | str, + ) -> None: + """ + Register an external model to be used in vLLM. + + :code:`model_cls` can be either: + + - A :class:`torch.nn.Module` class directly referencing the model. + - A string in the format :code:`:` which can be used to + lazily import the model. This is useful to avoid initializing CUDA + when importing the model and thus the related error + :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`. + """ + if model_arch in self.models: + logger.warning( + "Model architecture %s is already registered, and will be " + "overwritten by the new model class %s.", model_arch, model_cls) + + if isinstance(model_cls, str): + split_str = model_cls.split(":") + if len(split_str) != 2: + msg = "Expected a string in the format `:`" + raise ValueError(msg) + + model = _LazyRegisteredModel(*split_str) + else: + model = _RegisteredModel.from_model_cls(model_cls) + + self.models[model_arch] = model + + def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn: + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details.") + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}") + + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: + if model_arch not in self.models: + return None + + return _try_load_model_cls(model_arch, self.models[model_arch]) + + def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None: + if model_arch not in self.models: + return None + + return _try_inspect_model_cls(model_arch, self.models[model_arch]) + + def _normalize_archs( + self, + architectures: str | list[str], + ) -> list[str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + logger.warning("No model architectures are specified") + + normalized_arch = [] + for model in architectures: + if model not in self.models: + model = "TransformersModel" + normalized_arch.append(model) + return normalized_arch + + def inspect_model_cls( + self, + architectures: str | list[str], + ) -> tuple[_ModelInfo, str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + return self._raise_for_unsupported(architectures) + + def resolve_model_cls( + self, + architectures: str | list[str], + ) -> tuple[type[nn.Module], str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + +ModelRegistry = _ModelRegistry({ + model_arch: + _LazyRegisteredModel( + module_name=f"trainer.models.{component_name}.{mod_relname}", + component_name=component_name, + class_name=cls_name, + ) + for model_arch, (component_name, mod_relname, + cls_name) in _FAST_VIDEO_MODELS.items() +}) diff --git a/trainer/models/schedulers/base.py b/trainer/models/schedulers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..4c647efc408e2db53f12ec71a0bee3d2804ed844 --- /dev/null +++ b/trainer/models/schedulers/base.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod + +import torch + + +class BaseScheduler(ABC): + timesteps: torch.Tensor + order: int + num_train_timesteps: int + + def __init__(self, *args, **kwargs) -> None: + # Check if subclass has defined all required properties + required_attributes = ['timesteps', 'order', 'num_train_timesteps'] + + for attr in required_attributes: + if not hasattr(self, attr): + raise AttributeError( + f"Subclasses of BaseScheduler must define '{attr}' property" + ) + + @abstractmethod + def set_shift(self, shift: float) -> None: + pass + + @abstractmethod + def set_timesteps(self, *args, **kwargs) -> None: + pass + + @abstractmethod + def scale_model_input(self, + sample: torch.Tensor, + timestep: int | None = None) -> torch.Tensor: + pass diff --git a/trainer/models/schedulers/scheduling_flow_match_euler_discrete.py b/trainer/models/schedulers/scheduling_flow_match_euler_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..df4b88e99c535bd5ae4a6732a7d7ceaf8d602537 --- /dev/null +++ b/trainer/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -0,0 +1,658 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +import math +from dataclasses import dataclass +from typing import Any + +import numpy as np +import scipy.stats +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + +from trainer.logger import init_logger +from trainer.models.schedulers.base import BaseScheduler +from trainer.distributed import (get_world_group) + +logger = init_logger(__name__) + + +@dataclass +class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, + BaseScheduler): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + use_dynamic_shifting (`bool`, defaults to False): + Whether to apply timestep shifting on-the-fly based on the image resolution. + base_shift (`float`, defaults to 0.5): + Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent + with desired output. + max_shift (`float`, defaults to 1.15): + Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be + more exaggerated or stylized. + base_image_seq_len (`int`, defaults to 256): + The base image sequence length. + max_image_seq_len (`int`, defaults to 4096): + The maximum image sequence length. + invert_sigmas (`bool`, defaults to False): + Whether to invert the sigmas. + shift_terminal (`float`, defaults to None): + The end value of the shifted timestep schedule. + use_karras_sigmas (`bool`, defaults to False): + Whether to use Karras sigmas for step sizes in the noise schedule during sampling. + use_exponential_sigmas (`bool`, defaults to False): + Whether to use exponential sigmas for step sizes in the noise schedule during sampling. + use_beta_sigmas (`bool`, defaults to False): + Whether to use beta sigmas for step sizes in the noise schedule during sampling. + time_shift_type (`str`, defaults to "exponential"): + The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". + stochastic_sampling (`bool`, defaults to False): + Whether to use stochastic sampling. + """ + + _compatibles: list[Any] = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting: bool = False, + base_shift: float | None = 0.5, + max_shift: float | None = 1.15, + base_image_seq_len: int | None = 256, + max_image_seq_len: int | None = 4096, + invert_sigmas: bool = False, + shift_terminal: float | None = None, + use_karras_sigmas: bool | None = False, + use_exponential_sigmas: bool | None = False, + use_beta_sigmas: bool | None = False, + time_shift_type: str = "exponential", + stochastic_sampling: bool = False, + ): + if sum([ + self.config.use_beta_sigmas, self.config.use_exponential_sigmas, + self.config.use_karras_sigmas + ]) > 1: + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + ) + if time_shift_type not in {"exponential", "linear"}: + raise ValueError( + "`time_shift_type` must either be 'exponential' or 'linear'.") + + timesteps = np.linspace(1, + num_train_timesteps, + num_train_timesteps, + dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + self.num_train_timesteps = num_train_timesteps + + self._step_index: int | None = None + self._begin_index: int | None = None + + self._shift = shift + + self.sigmas = sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + BaseScheduler.__init__(self) + + @property + def shift(self) -> float: + """ + The value used for shifting. + """ + return self._shift + + @property + def step_index(self) -> int | None: + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self) -> int | None: + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0) -> None: + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def set_shift(self, shift: float) -> None: + self._shift = shift + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Forward process in flow-matching + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, + dtype=torch.float32) + assert isinstance(timestep, torch.Tensor) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + assert isinstance(timestep, torch.Tensor) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) for t in timestep + ] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + + return sample + + def _sigma_to_t(self, sigma: float) -> float: + return sigma * self.config.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, + t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + else: + raise ValueError( + f"Unknown time_shift_type: {self.config.time_shift_type}") + + def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: + r""" + Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config + value. + + Reference: + https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 + + Args: + t (`torch.Tensor`): + A tensor of timesteps to be stretched and shifted. + + Returns: + `torch.Tensor`: + A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. + """ + one_minus_z = 1 - t + scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) + stretched_t = 1 - (one_minus_z / scale_factor) + return stretched_t + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, + ) -> None: + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`, *optional*): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + sigmas (`List[float]`, *optional*): + Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed + automatically. + mu (`float`, *optional*): + Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep + shifting. + timesteps (`List[float]`, *optional*): + Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed + automatically. + """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + "`mu` must be passed when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is not None and timesteps is not None and len(sigmas) != len( + timesteps): + raise ValueError( + "`sigmas` and `timesteps` should have the same length") + + if num_inference_steps is not None: + if (sigmas is not None and len(sigmas) != num_inference_steps) or ( + timesteps is not None + and len(timesteps) != num_inference_steps): + raise ValueError( + "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" + ) + else: + if sigmas is not None: + num_inference_steps = len(sigmas) + elif timesteps is not None: + num_inference_steps = len(timesteps) + else: + raise ValueError( + "Either num_inference_steps, sigmas, or timesteps must be provided" + ) + + self.num_inference_steps = num_inference_steps + + # 1. Prepare default sigmas + is_timesteps_provided = timesteps is not None + + timesteps_array: np.ndarray | None = None + if is_timesteps_provided: + assert timesteps is not None + timesteps_array = np.array(timesteps).astype(np.float32) + + sigmas_array: np.ndarray + if sigmas is None: + if timesteps_array is None: + timesteps_array = np.linspace(self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps) + sigmas_array = timesteps_array / self.config.num_train_timesteps + else: + sigmas_array = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas_array) + + # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of + # "exponential" or "linear" type is applied + if self.config.use_dynamic_shifting: + assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" + sigmas_array = self.time_shift(mu, 1.0, sigmas_array) + else: + sigmas_array = self.shift * sigmas_array / ( + 1 + (self.shift - 1) * sigmas_array) + + # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value + if self.config.shift_terminal: + sigmas_tensor = torch.from_numpy(sigmas_array).to( + dtype=torch.float32) + sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) + sigmas_array = sigmas_tensor.numpy() + + # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules + if self.config.use_karras_sigmas: + sigmas_tensor = torch.from_numpy(sigmas_array).to( + dtype=torch.float32) + sigmas_tensor = self._convert_to_karras( + in_sigmas=sigmas_tensor, + num_inference_steps=num_inference_steps) + sigmas_array = sigmas_tensor.numpy() + elif self.config.use_exponential_sigmas: + sigmas_tensor = torch.from_numpy(sigmas_array).to( + dtype=torch.float32) + sigmas_tensor = self._convert_to_exponential( + in_sigmas=sigmas_tensor, + num_inference_steps=num_inference_steps) + sigmas_array = sigmas_tensor.numpy() + elif self.config.use_beta_sigmas: + sigmas_tensor = torch.from_numpy(sigmas_array).to( + dtype=torch.float32) + sigmas_tensor = self._convert_to_beta( + in_sigmas=sigmas_tensor, + num_inference_steps=num_inference_steps) + sigmas_array = sigmas_tensor.numpy() + + # 5. Convert sigmas and timesteps to tensors and move to specified device + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32, + device=device) + if not is_timesteps_provided: + timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps + else: + assert timesteps_array is not None + timesteps_tensor = torch.from_numpy(timesteps_array).to( + dtype=torch.float32, device=device) + + # 6. Append the terminal sigma value. + # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the + # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi + if self.config.invert_sigmas: + sigmas_tensor = 1.0 - sigmas_tensor + timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps + sigmas_tensor = torch.cat( + [sigmas_tensor, + torch.ones(1, device=sigmas_tensor.device)]) + else: + sigmas_tensor = torch.cat( + [sigmas_tensor, + torch.zeros(1, device=sigmas_tensor.device)]) + + self.timesteps = timesteps_tensor + self.sigmas = sigmas_tensor + self._step_index = None + self._begin_index = None + + def index_for_timestep( + self, + timestep: float | torch.FloatTensor, + schedule_timesteps: torch.Tensor | None = None) -> int: + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: int | torch.Tensor, + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: torch.Generator | None = None, + per_token_timesteps: torch.Tensor | None = None, + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`int` or `torch.Tensor`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + per_token_timesteps (`torch.Tensor`, *optional*): + The timesteps for each token in the sample. + return_dict (`bool`): + Whether or not to return a + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. + + Returns: + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, + otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int | torch.IntTensor | torch.LongTensor)): + raise ValueError(( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + if per_token_timesteps is not None: + per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps + + sigmas = self.sigmas[:, None, None] + lower_mask = sigmas < per_token_sigmas[None] - 1e-6 + lower_sigmas = lower_mask * sigmas + lower_sigmas, _ = lower_sigmas.max(dim=0) + + current_sigma = per_token_sigmas[..., None] + next_sigma = lower_sigmas[..., None] + dt = current_sigma - next_sigma + else: + assert self.step_index is not None, "step_index should not be None" + sigma_idx = self.step_index + sigma = self.sigmas[sigma_idx] + sigma_next = self.sigmas[sigma_idx + 1] + + current_sigma = sigma + next_sigma = sigma_next + dt = sigma_next - sigma + + if self.config.stochastic_sampling: + x0 = sample - current_sigma * model_output + noise = torch.randn_like(sample) + prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + prev_sample = sample + dt * model_output + + # upon completion increase step index by one + assert self._step_index is not None, "_step_index should not be None" + self._step_index += 1 + if per_token_timesteps is None: + # Cast sample back to model compatible dtype + prev_sample = prev_sample.to(model_output.dtype) + + if isinstance(prev_sample, torch.Tensor | float) and not return_dict: + return (prev_sample, ) + + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, + num_inference_steps: int) -> torch.Tensor: + """Constructs the noise schedule of Karras et al. (2022).""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min**(1 / rho) + max_inv_rho = sigma_max**(1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, + num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp( + np.linspace(math.log(sigma_max), math.log(sigma_min), + num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta(self, + in_sigmas: torch.Tensor, + num_inference_steps: int, + alpha: float = 0.6, + beta: float = 0.6) -> torch.Tensor: + """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array([ + sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ]) + return sigmas + + def _time_shift_exponential( + self, mu: float, sigma: float, + t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + if isinstance(t, np.ndarray): + return np.exp(mu) / (np.exp(mu) + (1 / t - 1)**sigma) + else: + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def _time_shift_linear( + self, mu: float, sigma: float, + t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + return mu / (mu + (1 / t - 1)**sigma) + + def add_noise( + self, + clean_latent: torch.Tensor, + noise: torch.Tensor, + timestep: torch.IntTensor, + ) -> torch.Tensor: + # self.sigmas = self.sigmas.to(noise.device) + # timestep = timestep.expand(clean_latent.shape[0]) + # self.timesteps = self.timesteps.to(noise.device) + # timestep_id = torch.argmin( + # (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + # sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + sigma = timestep / 1000.0 + sigma = sigma.reshape(-1,1,1,1).to(torch.float32) + # if get_world_group().rank == 0: + # print('add noise', sigma) + sample = (1 - sigma) * clean_latent.to(torch.float32) + sigma * noise.to(torch.float32) + return sample.type_as(noise) + + def scale_model_input(self, + sample: torch.Tensor, + timestep: int | None = None) -> torch.Tensor: + return sample + + def __len__(self) -> int: + return self.config.num_train_timesteps diff --git a/trainer/models/utils.py b/trainer/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..011314c2339d6f7dbaad6207c472fce1a1badb0b --- /dev/null +++ b/trainer/models/utils.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/utils.py +"""Utils for model executor.""" +from typing import Any + +import torch + + +# TODO(PY): move it elsewhere +def auto_attributes(init_func): + """ + Decorator that automatically adds all initialization arguments as object attributes. + + Example: + @auto_attributes + def __init__(self, a=1, b=2): + pass + + # This will automatically set: + # - self.a = 1 and self.b = 2 + # - self.config.a = 1 and self.config.b = 2 + """ + + def wrapper(self, *args, **kwargs): + # Get the function signature + import inspect + signature = inspect.signature(init_func) + parameters = signature.parameters + + # Get parameter names (excluding 'self') + param_names = list(parameters.keys())[1:] + + # Bind arguments to parameters + bound_args = signature.bind(self, *args, **kwargs) + bound_args.apply_defaults() + + # Create config object if it doesn't exist + if not hasattr(self, 'config'): + self.config = type('Config', (), {})() + + # Set attributes on self and self.config + for name in param_names: + if name in bound_args.arguments: + value = bound_args.arguments[name] + setattr(self, name, value) + setattr(self.config, name, value) + + # Call the original __init__ function + return init_func(self, *args, **kwargs) + + return wrapper + + +def set_weight_attrs( + weight: torch.Tensor, + weight_attrs: dict[str, Any] | None, +): + """Set attributes on a weight tensor. + + This method is used to set attributes on a weight tensor. This method + will not overwrite existing attributes. + + Args: + weight: The weight tensor. + weight_attrs: A dictionary of attributes to set on the weight tensor. + """ + if weight_attrs is None: + return + for key, value in weight_attrs.items(): + assert not hasattr( + weight, key), (f"Overwriting existing tensor attribute: {key}") + + # NOTE(woosuk): During weight loading, we often do something like: + # narrowed_tensor = param.data.narrow(0, offset, len) + # narrowed_tensor.copy_(real_weight) + # expecting narrowed_tensor and param.data to share the same storage. + # However, on TPUs, narrowed_tensor will lazily propagate to the base + # tensor, which is param.data, leading to the redundant memory usage. + # This sometimes causes OOM errors during model loading. To avoid this, + # we sync the param tensor after its weight loader is called. + # TODO(woosuk): Remove this hack once we have a better solution. + from trainer.platforms import current_platform + if current_platform.is_tpu() and key == "weight_loader": + value = _make_synced_weight_loader(value) + setattr(weight, key, value) + + +def _make_synced_weight_loader(original_weight_loader) -> Any: + + def _synced_weight_loader(param, *args, **kwargs): + original_weight_loader(param, *args, **kwargs) + torch._sync(param) + + return _synced_weight_loader + + +def extract_layer_index(layer_name: str) -> int: + """ + Extract the layer index from the module name. + Examples: + - "encoder.layers.0" -> 0 + - "encoder.layers.1.self_attn" -> 1 + - "2.self_attn" -> 2 + - "model.encoder.layers.0.sub.1" -> ValueError + """ + subnames = layer_name.split(".") + int_vals: list[int] = [] + for subname in subnames: + try: + int_vals.append(int(subname)) + except ValueError: + continue + assert len(int_vals) == 1, (f"layer name {layer_name} should" + " only contain one integer") + return int_vals[0] + + +def modulate(x: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None) -> torch.Tensor: + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) # type: ignore[union-attr] + elif scale is None: + return x + shift.unsqueeze(1) # type: ignore[union-attr] + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze( + 1) # type: ignore[union-attr] diff --git a/trainer/models/vaes/common.py b/trainer/models/vaes/common.py new file mode 100644 index 0000000000000000000000000000000000000000..37d4c0ebe862d6d674fb047d664c570bcd3ee68e --- /dev/null +++ b/trainer/models/vaes/common.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from collections.abc import Iterator +from math import prod +from typing import Optional, cast + +import numpy as np +import torch +import torch.distributed as dist +from diffusers.utils.torch_utils import randn_tensor + +from trainer.configs.models import VAEConfig +from trainer.distributed import get_sp_parallel_rank, get_sp_world_size + + +class ParallelTiledVAE(ABC): + tile_sample_min_height: int + tile_sample_min_width: int + tile_sample_min_num_frames: int + tile_sample_stride_height: int + tile_sample_stride_width: int + tile_sample_stride_num_frames: int + blend_num_frames: int + use_tiling: bool + use_temporal_tiling: bool + use_parallel_tiling: bool + + def __init__(self, config: VAEConfig, **kwargs) -> None: + self.config = config + self.tile_sample_min_height = config.tile_sample_min_height + self.tile_sample_min_width = config.tile_sample_min_width + self.tile_sample_min_num_frames = config.tile_sample_min_num_frames + self.tile_sample_stride_height = config.tile_sample_stride_height + self.tile_sample_stride_width = config.tile_sample_stride_width + self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames + self.blend_num_frames = config.blend_num_frames + self.use_tiling = config.use_tiling + self.use_temporal_tiling = config.use_temporal_tiling + self.use_parallel_tiling = config.use_parallel_tiling + + def to(self, device) -> 'ParallelTiledVAE': + return self + + @property + def temporal_compression_ratio(self) -> int: + return cast(int, self.config.temporal_compression_ratio) + + @property + def spatial_compression_ratio(self) -> int: + return cast(int, self.config.spatial_compression_ratio) + + @property + def scaling_factor(self) -> float | torch.Tensor: + return cast(float | torch.Tensor, self.config.scaling_factor) + + @abstractmethod + def _encode(self, *args, **kwargs) -> torch.Tensor: + pass + + @abstractmethod + def _decode(self, *args, **kwargs) -> torch.Tensor: + pass + + def encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + latent_num_frames = (num_frames - + 1) // self.temporal_compression_ratio + 1 + + if self.use_tiling and self.use_temporal_tiling and num_frames > self.tile_sample_min_num_frames: + latents = self.tiled_encode(x)[:, :, :latent_num_frames] + elif self.use_tiling and (width > self.tile_sample_min_width + or height > self.tile_sample_min_height): + latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames] + else: + latents = self._encode(x)[:, :, :latent_num_frames] + return DiagonalGaussianDistribution(latents) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + num_sample_frames = (num_frames - + 1) * self.temporal_compression_ratio + 1 + + if self.use_tiling and self.use_parallel_tiling and get_sp_world_size( + ) > 1: + return self.parallel_tiled_decode(z)[:, :, :num_sample_frames] + if self.use_tiling and self.use_temporal_tiling and num_frames > tile_latent_min_num_frames: + return self.tiled_decode(z)[:, :, :num_sample_frames] + + if self.use_tiling and (width > tile_latent_min_width + or height > tile_latent_min_height): + return self.spatial_tiled_decode(z)[:, :, :num_sample_frames] + + return self._decode(z)[:, :, :num_sample_frames] + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( + 1 - x / blend_extent) + b[:, :, x, :, :] * (x / blend_extent) + return b + + def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + _, _, _, height, width = x.shape + # latent_height = height // self.spatial_compression_ratio + # latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + tile = x[:, :, :, i:i + self.tile_sample_min_height, + j:j + self.tile_sample_min_width] + tile = self._encode(tile) + row.append(tile) + rows.append(row) + + return self._merge_spatial_tiles(rows, blend_height, blend_width, + tile_latent_stride_height, + tile_latent_stride_width) + + def _parallel_data_generator( + self, gathered_results, + gathered_dim_metadata) -> Iterator[tuple[torch.Tensor, int]]: + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + world_size, rank = get_sp_world_size(), get_sp_parallel_rank() + B, C, T, H, W = z.shape + + # Calculate parameters + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Calculate tile dimensions + num_t_tiles = (T + tile_latent_stride_num_frames - + 1) // tile_latent_stride_num_frames + num_h_tiles = (H + tile_latent_stride_height - + 1) // tile_latent_stride_height + num_w_tiles = (W + tile_latent_stride_width - + 1) // tile_latent_stride_width + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate( + range(start_tile_idx, end_tile_idx)): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * tile_latent_stride_num_frames + h_start = h_idx * tile_latent_stride_height + w_start = w_idx * tile_latent_stride_width + + # Extract and process tile + tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, + h_start:h_start + tile_latent_min_height, + w_start:w_start + tile_latent_min_width] + + # Process tile + tile = self._decode(tile) + + if t_start > 0: + tile = tile[:, :, 1:, :, :] + + # Store metadata + shape = tile.shape + # Store decoded data (flattened) + decoded_flat = tile.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], + device=results.device, + dtype=torch.int64) + all_sizes = [ + torch.zeros(1, device=results.device, dtype=torch.int64) + for _ in range(world_size) + ] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape) + ).contiguous( + ) # use contiguous to make sure it won't copy data in the following operations + # TODO (PY): use trainer distributed methods + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data: list = [[[[] for _ in range(num_w_tiles)] + for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator( + gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles( + tem_data, blend_height, blend_width, + self.tile_sample_stride_height, self.tile_sample_stride_width) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, + self.blend_num_frames) + result_slices.append( + slice_data[:, :, :self.tile_sample_stride_num_frames, :, :]) + else: + result_slices.append( + slice_data[:, :, :self.tile_sample_stride_num_frames + + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + return dec + + def _merge_spatial_tiles(self, tiles, blend_height, blend_width, + stride_height, stride_width) -> torch.Tensor: + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(tiles): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(tiles[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :stride_height, :stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + + Returns: + `torch.Tensor`: + The decoded images. + """ + + _, _, _, height, width = z.shape + # sample_height = height * self.spatial_compression_ratio + # sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + tile = z[:, :, :, i:i + tile_latent_min_height, + j:j + tile_latent_min_width] + decoded = self._decode(tile) + row.append(decoded) + rows.append(row) + return self._merge_spatial_tiles(rows, blend_height, blend_width, + self.tile_sample_stride_height, + self.tile_sample_stride_width) + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + _, _, num_frames, height, width = x.shape + + # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + + row = [] + for i in range(0, num_frames, self.tile_sample_stride_num_frames): + tile = x[:, :, i:i + self.tile_sample_min_num_frames + 1, :, :] + if self.use_tiling and (height > self.tile_sample_min_height + or width > self.tile_sample_min_width): + tile = self.spatial_tiled_encode(tile) + else: + tile = self._encode(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, self.blend_num_frames) + result_row.append( + tile[:, :, :tile_latent_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, :tile_latent_stride_num_frames + + 1, :, :]) + enc = torch.cat(result_row, dim=2) + return enc + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + + row = [] + for i in range(0, num_frames, tile_latent_stride_num_frames): + tile = z[:, :, i:i + tile_latent_min_num_frames + 1, :, :] + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width + or tile.shape[-2] > tile_latent_min_height): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self._decode(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, self.blend_num_frames) + result_row.append( + tile[:, :, :self.tile_sample_stride_num_frames, :, :]) + else: + result_row.append( + tile[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) + + dec = torch.cat(result_row, dim=2) + return dec + + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: int | None = None, + tile_sample_stride_width: int | None = None, + tile_sample_stride_num_frames: int | None = None, + blend_num_frames: int | None = None, + use_tiling: bool | None = None, + use_temporal_tiling: bool | None = None, + use_parallel_tiling: bool | None = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_min_num_frames (`int`, *optional*): + The minimum number of frames required for a sample to be separated into tiles across the frame + dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + tile_sample_stride_num_frames (`int`, *optional*): + The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts + produced across the frame dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames + if blend_num_frames is not None: + self.blend_num_frames = blend_num_frames + else: + self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + self.use_tiling = use_tiling or self.use_tiling + self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling + self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling + + def disable_tiling(self) -> None: + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + +# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 +class DiagonalGaussianDistribution: + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, + other: Optional["DiagonalGaussianDistribution"] = None + ) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=[1, 2, 3], + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=[1, 2, 3], + ) + + def nll( + self, sample: torch.Tensor, + dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean diff --git a/trainer/models/vaes/hunyuanvae.py b/trainer/models/vaes/hunyuanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..e3701be85cb1c5884d6234751f96bc3944d034f0 --- /dev/null +++ b/trainer/models/vaes/hunyuanvae.py @@ -0,0 +1,852 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from diffusers + +# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The Trainer Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from trainer.configs.models.vaes import HunyuanVAEConfig +from trainer.layers.activation import get_act_fn +from trainer.models.vaes.common import ParallelTiledVAE + + +def prepare_causal_attention_mask( + num_frames: int, + height_width: int, + dtype: torch.dtype, + device: torch.device, + batch_size: int | None = None) -> torch.Tensor: + indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) + indices_blocks = indices.repeat_interleave(height_width) + x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") + mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype) + + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class HunyuanVAEAttention(nn.Module): + + def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, + bias) -> None: + super().__init__() + self.in_channels = in_channels + self.heads = heads + self.dim_head = dim_head + self.eps = eps + self.norm_num_groups = norm_num_groups + self.bias = bias + + inner_dim = heads * dim_head + + # Define the projection layers + self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, + bias=bias)) + + # Optional normalization layers + self.group_norm = nn.GroupNorm(norm_num_groups, + in_channels, + eps=eps, + affine=True) + + def forward(self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None) -> torch.Tensor: + residual = hidden_states + + batch_size, sequence_length, _ = hidden_states.shape + + hidden_states = self.group_norm(hidden_states.transpose(1, + 2)).transpose( + 1, 2) + + # Project to query, key, value + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + # Reshape for multi-head attention + head_dim = self.dim_head + + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + # Perform scaled dot-product attention + hidden_states = F.scaled_dot_product_attention(query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False) + + # Reshape back + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, self.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # Linear projection + hidden_states = self.to_out(hidden_states) + + # Residual connection and rescale + hidden_states = hidden_states + residual + + return hidden_states + + +class HunyuanVideoCausalConv3d(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + dilation: int | tuple[int, int, int] = 1, + bias: bool = True, + pad_mode: str = "replicate", + ) -> None: + super().__init__() + + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance( + kernel_size, int) else kernel_size + + self.pad_mode = pad_mode + self.time_causal_padding = ( + kernel_size[0] // 2, + kernel_size[0] // 2, + kernel_size[1] // 2, + kernel_size[1] // 2, + kernel_size[2] - 1, + 0, + ) + + self.conv = nn.Conv3d(in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.pad(hidden_states, + self.time_causal_padding, + mode=self.pad_mode) + return self.conv(hidden_states) + + +class HunyuanVideoUpsampleCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + kernel_size: int = 3, + stride: int = 1, + bias: bool = True, + upsample_factor: tuple[int, ...] = (2, 2, 2), + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.upsample_factor = upsample_factor + + self.conv = HunyuanVideoCausalConv3d(in_channels, + out_channels, + kernel_size, + stride, + bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_frames = hidden_states.size(2) + + first_frame, other_frames = hidden_states.split((1, num_frames - 1), + dim=2) + first_frame = F.interpolate(first_frame.squeeze(2), + scale_factor=self.upsample_factor[1:], + mode="nearest").unsqueeze(2) + + if num_frames > 1: + # See: https://github.com/pytorch/pytorch/issues/81665 + # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate + # is fixed, this will raise either a runtime error, or fail silently with bad outputs. + # If you are encountering an error here, make sure to try running encoding/decoding with + # `vae.enable_tiling()` first. If that doesn't work, open an issue at: + # https://github.com/huggingface/diffusers/issues + other_frames = other_frames.contiguous() + other_frames = F.interpolate(other_frames, + scale_factor=self.upsample_factor, + mode="nearest") + hidden_states = torch.cat((first_frame, other_frames), dim=2) + else: + hidden_states = first_frame + + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoDownsampleCausal3D(nn.Module): + + def __init__( + self, + channels: int, + out_channels: int | None = None, + padding: int = 1, + kernel_size: int = 3, + bias: bool = True, + stride=2, + ) -> None: + super().__init__() + out_channels = out_channels or channels + + self.conv = HunyuanVideoCausalConv3d(channels, + out_channels, + kernel_size, + stride, + padding, + bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoResnetBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + non_linearity: str = "silu", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_act_fn(non_linearity) + + self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) + self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, + 0) + + self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) + self.dropout = nn.Dropout(dropout) + self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, + 0) + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = HunyuanVideoCausalConv3d( + in_channels, out_channels, 1, 1, 0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.contiguous() + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + hidden_states = hidden_states + residual + return hidden_states + + +class HunyuanVideoMidBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_attention: bool = True, + attention_head_dim: int = 1, + ) -> None: + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min( + in_channels // 4, 32) + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ] + attentions: list[HunyuanVAEAttention | None] = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + HunyuanVAEAttention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + eps=resnet_eps, + norm_num_groups=resnet_groups, + bias=True, + )) + else: + attentions.append(None) + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + self.resnets[0], hidden_states) + + for attn, resnet in zip(self.attentions, + self.resnets[1:], + strict=True): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, + 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, + height * width, + hidden_states.dtype, + hidden_states.device, + batch_size=batch_size) + hidden_states = attn(hidden_states, + attention_mask=attention_mask) + hidden_states = hidden_states.unflatten( + 1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = self._gradient_checkpointing_func( + resnet, hidden_states) + + else: + hidden_states = self.resnets[0](hidden_states) + + for attn, resnet in zip(self.attentions, + self.resnets[1:], + strict=True): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, + 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, + height * width, + hidden_states.dtype, + hidden_states.device, + batch_size=batch_size) + hidden_states = attn(hidden_states, + attention_mask=attention_mask) + hidden_states = hidden_states.unflatten( + 1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = resnet(hidden_states) + + return hidden_states + + +class HunyuanVideoDownBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_downsample: bool = True, + downsample_stride: tuple[int, ...] | int = 2, + downsample_padding: int = 1, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + HunyuanVideoDownsampleCausal3D( + out_channels, + out_channels=out_channels, + padding=downsample_padding, + stride=downsample_stride, + ) + ]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func( + resnet, hidden_states) + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoUpBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_upsample: bool = True, + upsample_scale_factor: tuple[int, ...] = (2, 2, 2), + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + HunyuanVideoUpsampleCausal3D( + out_channels, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func( + resnet, hidden_states) + + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoEncoder3D(nn.Module): + r""" + Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ) -> None: + super().__init__() + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, + block_out_channels[0], + kernel_size=3, + stride=1) + self.mid_block: HunyuanVideoMidBlock3D | None = None + self.down_blocks = nn.ModuleList([]) + + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + if down_block_type != "HunyuanVideoDownBlock3D": + raise ValueError( + f"Unsupported down_block_type: {down_block_type}") + + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int( + np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int( + np.log2(temporal_compression_ratio)) + + if temporal_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - + num_time_downsample_layers) + and not is_final_block) + elif temporal_compression_ratio == 8: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError( + f"Unsupported time_compression_ratio: {temporal_compression_ratio}" + ) + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2, ) if add_time_downsample else (1, ) + downsample_stride = tuple(downsample_stride_T + + downsample_stride_HW) + + down_block = HunyuanVideoDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample + or add_time_downsample), + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + downsample_stride=downsample_stride, + downsample_padding=0, + ) + + self.down_blocks.append(down_block) + + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], + num_groups=norm_num_groups, + eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], + conv_out_channels, + kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for down_block in self.down_blocks: + hidden_states = self._gradient_checkpointing_func( + down_block, hidden_states) + + hidden_states = self._gradient_checkpointing_func( + self.mid_block, hidden_states) + else: + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states) + assert self.mid_block is not None + hidden_states = self.mid_block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class HunyuanVideoDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + if up_block_type != "HunyuanVideoUpBlock3D": + raise ValueError(f"Unsupported up_block_type: {up_block_type}") + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int( + np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool( + i >= len(block_out_channels) - 1 - num_time_upsample_layers + and not is_final_block) + else: + raise ValueError( + f"Unsupported time_compression_ratio: {time_compression_ratio}" + ) + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, + 1) + upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) + upsample_scale_factor = tuple(upsample_scale_factor_T + + upsample_scale_factor_HW) + + up_block = HunyuanVideoUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + ) + + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], + out_channels, + kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + self.mid_block, hidden_states) + + for up_block in self.up_blocks: + hidden_states = self._gradient_checkpointing_func( + up_block, hidden_states) + else: + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states) + + # post-process + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + config: HunyuanVAEConfig, + ) -> None: + nn.Module.__init__(self) + ParallelTiledVAE.__init__(self, config) + + # TODO(will): only pass in config. We do this by manually defining a + # config for hunyuan vae + self.block_out_channels = config.block_out_channels + + if config.load_encoder: + self.encoder = HunyuanVideoEncoder3D( + in_channels=config.in_channels, + out_channels=config.latent_channels, + down_block_types=config.down_block_types, + block_out_channels=config.block_out_channels, + layers_per_block=config.layers_per_block, + norm_num_groups=config.norm_num_groups, + act_fn=config.act_fn, + double_z=True, + mid_block_add_attention=config.mid_block_add_attention, + temporal_compression_ratio=config.temporal_compression_ratio, + spatial_compression_ratio=config.spatial_compression_ratio, + ) + self.quant_conv = nn.Conv3d(2 * config.latent_channels, + 2 * config.latent_channels, + kernel_size=1) + + if config.load_decoder: + self.decoder = HunyuanVideoDecoder3D( + in_channels=config.latent_channels, + out_channels=config.out_channels, + up_block_types=config.up_block_types, + block_out_channels=config.block_out_channels, + layers_per_block=config.layers_per_block, + norm_num_groups=config.norm_num_groups, + act_fn=config.act_fn, + time_compression_ratio=config.temporal_compression_ratio, + spatial_compression_ratio=config.spatial_compression_ratio, + mid_block_add_attention=config.mid_block_add_attention, + ) + self.post_quant_conv = nn.Conv3d(config.latent_channels, + config.latent_channels, + kernel_size=1) + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + enc = self.quant_conv(x) + return enc + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z) + return dec diff --git a/trainer/models/vaes/stepvideovae.py b/trainer/models/vaes/stepvideovae.py new file mode 100644 index 0000000000000000000000000000000000000000..22bd923e1a53c71dbe9bafdd159272972c6bf364 --- /dev/null +++ b/trainer/models/vaes/stepvideovae.py @@ -0,0 +1,1140 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +from typing import Any + +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F + +from trainer.configs.models.vaes import StepVideoVAEConfig +from trainer.models.vaes.common import ParallelTiledVAE + + +def base_group_norm(x, + norm_layer, + act_silu=False, + channel_last=False) -> torch.Tensor: + if hasattr(base_group_norm, 'spatial') and base_group_norm.spatial: + assert channel_last + x_shape = x.shape + x = x.flatten(0, 1) + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + + out = F.group_norm(x.contiguous(), norm_layer.num_groups, + norm_layer.weight, norm_layer.bias, norm_layer.eps) + if act_silu: + out = F.silu(out) + + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + + out = out.view(x_shape) + else: + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + out = F.group_norm(x.contiguous(), norm_layer.num_groups, + norm_layer.weight, norm_layer.bias, norm_layer.eps) + if act_silu: + out = F.silu(out) + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + return out + + +def base_conv2d(x, + conv_layer, + channel_last=False, + residual=None) -> torch.Tensor: + if channel_last: + x = x.permute(0, 3, 1, 2) # NHWC to NCHW + out = F.conv2d(x, + conv_layer.weight, + conv_layer.bias, + stride=conv_layer.stride, + padding=conv_layer.padding) + if residual is not None: + if channel_last: + residual = residual.permute(0, 3, 1, 2) # NHWC to NCHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 1) # NCHW to NHWC + return out + + +def base_conv3d(x, + conv_layer, + channel_last=False, + residual=None, + only_return_output=False) -> torch.Tensor: + if only_return_output: + size = cal_outsize(x.shape, conv_layer.weight.shape, conv_layer.stride, + conv_layer.padding) + return torch.empty(size, device=x.device, dtype=x.dtype) + if channel_last: + x = x.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out = F.conv3d(x, + conv_layer.weight, + conv_layer.bias, + stride=conv_layer.stride, + padding=conv_layer.padding) + if residual is not None: + if channel_last: + residual = residual.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 4, 1) # NCDHW to NDHWC + return out + + +def cal_outsize(input_sizes, kernel_sizes, stride, padding) -> list: + stride_d, stride_h, stride_w = stride + padding_d, padding_h, padding_w = padding + dilation_d, dilation_h, dilation_w = 1, 1, 1 + + in_d = input_sizes[1] + in_h = input_sizes[2] + in_w = input_sizes[3] + + kernel_d = kernel_sizes[2] + kernel_h = kernel_sizes[3] + kernel_w = kernel_sizes[4] + out_channels = kernel_sizes[0] + + out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d) + out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h) + out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w) + size = [input_sizes[0], out_d, out_h, out_w, out_channels] + return size + + +def calc_out_(in_size: int, padding: int, dilation: int, kernel: int, + stride: int) -> int: + return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 + + +def base_conv3d_channel_last(x, conv_layer, residual=None) -> torch.Tensor: + in_numel = x.numel() + out_numel = int(x.numel() * conv_layer.out_channels / + conv_layer.in_channels) + if (in_numel >= 2**30) or (out_numel >= 2**30): + assert conv_layer.stride[0] == 1, "time split asks time stride = 1" + + B, T, H, W, C = x.shape + K = conv_layer.kernel_size[0] + + chunks = 4 + chunk_size = T // chunks + + if residual is None: + out_nhwc = base_conv3d(x, + conv_layer, + channel_last=True, + residual=residual, + only_return_output=True) + else: + out_nhwc = residual + + assert B == 1 + for i in range(chunks): + if i == chunks - 1: + xi = x[:1, chunk_size * i:] + out_nhwci = out_nhwc[:1, chunk_size * i:] + else: + xi = x[:1, chunk_size * i:chunk_size * (i + 1) + K - 1] + out_nhwci = out_nhwc[:1, chunk_size * i:chunk_size * (i + 1)] + if residual is not None: + if i == chunks - 1: + ri = residual[:1, chunk_size * i:] + else: + ri = residual[:1, chunk_size * i:chunk_size * (i + 1)] + else: + ri = None + out_nhwci.copy_( + base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) + else: + out_nhwc = base_conv3d(x, + conv_layer, + channel_last=True, + residual=residual) + return out_nhwc + + +class Upsample2D(nn.Module): + + def __init__(self, + channels, + use_conv=False, + use_conv_transpose=False, + out_channels=None) -> None: + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + + if use_conv: + self.conv = nn.Conv2d(self.channels, + self.out_channels, + 3, + padding=1) + else: + assert "Not Supported" + self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) + + def forward(self, x, output_size=None) -> torch.Tensor: + assert x.shape[-1] == self.channels + + if self.use_conv_transpose: + return self.conv(x) + + if output_size is None: + x = F.interpolate(x.permute( + 0, 3, 1, 2).to(memory_format=torch.channels_last), + scale_factor=2.0, + mode='nearest').permute(0, 2, 3, 1).contiguous() + else: + x = F.interpolate(x.permute( + 0, 3, 1, 2).to(memory_format=torch.channels_last), + size=output_size, + mode='nearest').permute(0, 2, 3, 1).contiguous() + + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class Downsample2D(nn.Module): + + def __init__(self, + channels, + use_conv=False, + out_channels=None, + padding=1) -> None: + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + + if use_conv: + self.conv = nn.Conv2d(self.channels, + self.out_channels, + 3, + stride=stride, + padding=padding) + else: + assert self.channels == self.out_channels + self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) + + def forward(self, x) -> torch.Tensor: + assert x.shape[-1] == self.channels + if self.use_conv and self.padding == 0: + pad = (0, 0, 0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + + assert x.shape[-1] == self.channels + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class CausalConv(nn.Module): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__() + + if isinstance(kernel_size, int): + kernel_size = kernel_size if isinstance( + kernel_size, tuple) else ((kernel_size, ) * 3) + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + self.dilation = kwargs.pop('dilation', 1) + self.stride = kwargs.pop('stride', 1) + if isinstance(self.stride, int): + self.stride = (self.stride, 1, 1) + time_pad = self.dilation * (time_kernel_size - 1) + max( + (1 - self.stride[0]), 0) + height_pad = height_kernel_size // 2 + width_pad = width_kernel_size // 2 + self.time_causal_padding = (width_pad, width_pad, height_pad, + height_pad, time_pad, 0) + self.time_uncausal_padding = (width_pad, width_pad, height_pad, + height_pad, 0, 0) + + self.conv = nn.Conv3d(chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + **kwargs) + self.chan_in = chan_in + self.chan_out = chan_out + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + x = nn.functional.pad( + x, + self.time_causal_padding if is_init else self.time_uncausal_padding) + x = self.conv(x) + if residual is not None: + x.add_(residual) + return x + + +class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert out_channels * factor**3 % in_channels == 0 + self.repeats = out_channels * factor**3 // in_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view(x.size(0), self.out_channels, self.factor, self.factor, + self.factor, x.size(2), x.size(3), x.size(4)) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(x.size(0), self.out_channels, + x.size(2) * self.factor, + x.size(4) * self.factor, + x.size(6) * self.factor) + x = x[:, :, self.factor - 1:, :, :] + return x + + +class ConvPixelShuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ) -> None: + super().__init__() + self.factor = factor + out_ratio = factor**3 + self.conv = CausalConv(in_channels, + out_channels * out_ratio, + kernel_size=kernel_size) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_shuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + batch_size, channels, depth, height, width = x.size() + new_channels = channels // (factor**3) + new_depth = depth * factor + new_height = height * factor + new_width = width * factor + + x = x.view(batch_size, new_channels, factor, factor, factor, depth, + height, width) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(batch_size, new_channels, new_depth, new_height, new_width) + x = x[:, :, factor - 1:, :, :] + return x + + +class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ) -> None: + super().__init__() + self.factor = factor + out_ratio = factor**3 + assert out_channels % out_ratio == 0 + self.conv = CausalConv(in_channels, + out_channels // out_ratio, + kernel_size=kernel_size) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_unshuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + pad = (0, 0, 0, 0, factor - 1, 0 + ) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, + factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) + return x + + +class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert in_channels * factor**3 % out_channels == 0 + self.group_size = in_channels * factor**3 // out_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + pad = (0, 0, 0, 0, self.factor - 1, 0 + ) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // self.factor, self.factor, H // self.factor, + self.factor, W // self.factor, self.factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, + W // self.factor) + x = x.view(B, self.out_channels, self.group_size, D // self.factor, + H // self.factor, W // self.factor) + x = x.mean(dim=2) + return x + + +def base_group_norm_with_zero_pad(x, + norm_layer, + act_silu=True, + pad_size=2) -> torch.Tensor: + out_shape = list(x.shape) + out_shape[1] += pad_size + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + out[:, pad_size:] = base_group_norm(x, + norm_layer, + act_silu=act_silu, + channel_last=True) + out[:, :pad_size] = 0 + return out + + +class CausalConvChannelLast(CausalConv): + time_causal_padding: tuple[Any, ...] + time_uncausal_padding: tuple[Any, ...] + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + self.time_causal_padding = (0, 0) + self.time_causal_padding + self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + if self.is_first_run: + self.is_first_run = False + # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) + + x = nn.functional.pad( + x, + self.time_causal_padding if is_init else self.time_uncausal_padding) + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class CausalConvAfterNorm(CausalConv): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + self.conv = nn.Conv3d(chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + padding=(0, 1, 1), + **kwargs) + else: + self.conv = nn.Conv3d(chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + **kwargs) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + if self.is_first_run: + self.is_first_run = False + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + pass + else: + x = nn.functional.pad(x, self.time_causal_padding).contiguous() + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class AttnBlock(nn.Module): + + def __init__(self, in_channels) -> None: + super().__init__() + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.proj_out = CausalConvChannelLast(in_channels, + in_channels, + kernel_size=1) + + def attention(self, x, is_init=True) -> torch.Tensor: + x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) + q = self.q(x, is_init) + k = self.k(x, is_init) + v = self.v(x, is_init) + + b, t, h, w, c = q.shape + q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), + (q, k, v)) + x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) + x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) + + return x + + def forward(self, x): + x = x.permute(0, 2, 3, 4, 1).contiguous() + h = self.attention(x) + x = self.proj_out(h, residual=x) + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Resnet3DBlock(nn.Module): + + def __init__( + self, + in_channels, + out_channels=None, + temb_channels=512, + conv_shortcut=False, + ) -> None: + super().__init__() + + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.conv1 = CausalConvAfterNorm(in_channels, + out_channels, + kernel_size=3) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) + self.conv2 = CausalConvAfterNorm(out_channels, + out_channels, + kernel_size=3) + + assert conv_shortcut is False + self.use_conv_shortcut = conv_shortcut + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConvAfterNorm(in_channels, + out_channels, + kernel_size=3) + else: + self.nin_shortcut = CausalConvAfterNorm(in_channels, + out_channels, + kernel_size=1) + + def forward(self, x, temb=None, is_init=True) -> torch.Tensor: + x = x.permute(0, 2, 3, 4, 1).contiguous() + + h = base_group_norm_with_zero_pad(x, + self.norm1, + act_silu=True, + pad_size=2) + h = self.conv1(h) + if temb is not None: + h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] + + x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x + + h = base_group_norm_with_zero_pad(h, + self.norm2, + act_silu=True, + pad_size=2) + x = self.conv2(h, residual=x) + + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Downsample3D(nn.Module): + + def __init__(self, in_channels, with_conv, stride) -> None: + super().__init__() + + self.with_conv = with_conv + if with_conv: + self.conv = CausalConv(in_channels, + in_channels, + kernel_size=3, + stride=stride) + + def forward(self, x, is_init=True) -> torch.Tensor: + if self.with_conv: + x = self.conv(x, is_init) + else: + x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) + return x + + +class VideoEncoder(nn.Module): + + def __init__( + self, + ch=32, + ch_mult=(4, 8, 16, 16), + num_res_blocks=2, + in_channels=3, + z_channels=16, + double_z=True, + down_sampling_layer=(1, 2), + resamp_with_conv=True, + version=1, + ) -> None: + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv(in_channels, ch, kernel_size=3) + self.down_sampling_layer = down_sampling_layer + + in_ch_mult = (1, ) + tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + Resnet3DBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=temb_ch)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + if i_level in self.down_sampling_layer: + down.downsample = Downsample3D(block_in, + resamp_with_conv, + stride=(2, 2, 2)) + else: + down.downsample = Downsample2D(block_in, + resamp_with_conv, + padding=0) #DIFF + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=temb_ch) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D( + block_in, channels, kernel_size=3, factor=2) + self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D( + block_in, channels, 2) + self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D( + channels, 2 * z_channels if double_z else z_channels, 1) + self.conv_out = CausalConvChannelLast( + channels, + 2 * z_channels if double_z else z_channels, + kernel_size=3) + else: + self.conv_out = CausalConvAfterNorm( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3) + + @torch.inference_mode() + def forward(self, x, video_frame_num, is_init=True) -> torch.Tensor: + # timestep embedding + temb = None + + t = video_frame_num + + # downsampling + h = self.conv_in(x, is_init) + + # make it real channel last, but behave like normal layout + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](h, temb, is_init) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + + if i_level != self.num_resolutions - 1: + if isinstance(self.down[i_level].downsample, Downsample2D): + _, _, t, _, _ = h.shape + h = rearrange(h, "b c t h w -> (b t) h w c", t=t) + h = self.down[i_level].downsample(h) + h = rearrange(h, "(b t) h w c -> b c t h w", t=t) + else: + h = self.down[i_level].downsample(h, is_init) + + h = self.mid.block_1(h, temb, is_init) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb, is_init) + + h = h.permute(0, 2, 3, 4, 1).contiguous() # b c l h w -> b l h w c + if self.version == 2: + h = base_group_norm(h, + self.norm_out, + act_silu=True, + channel_last=True) + h = h.permute(0, 4, 1, 2, 3).contiguous() + shortcut = self.shortcut_pathify(h, is_init) + h = self.conv_patchify(h, is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1) + h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init) + h = h.add_(shortcut) + else: + h = base_group_norm_with_zero_pad(h, + self.norm_out, + act_silu=True, + pad_size=2) + h = self.conv_out(h, is_init) + h = h.permute(0, 4, 1, 2, 3) # b l h w c -> b c l h w + + h = rearrange(h, "b c t h w -> b t c h w") + return h + + +class Res3DBlockUpsample(nn.Module): + + def __init__(self, + input_filters, + num_filters, + down_sampling_stride, + down_sampling=False) -> None: + super().__init__() + + self.input_filters = input_filters + self.num_filters = num_filters + + self.act_ = nn.SiLU(inplace=True) + + self.conv1 = CausalConvChannelLast(num_filters, + num_filters, + kernel_size=[3, 3, 3]) + self.norm1 = nn.GroupNorm(32, num_filters) + + self.conv2 = CausalConvChannelLast(num_filters, + num_filters, + kernel_size=[3, 3, 3]) + self.norm2 = nn.GroupNorm(32, num_filters) + + self.down_sampling = down_sampling + if down_sampling: + self.down_sampling_stride = down_sampling_stride + else: + self.down_sampling_stride = [1, 1, 1] + + if num_filters != input_filters or down_sampling: + self.conv3 = CausalConvChannelLast(input_filters, + num_filters, + kernel_size=[1, 1, 1], + stride=self.down_sampling_stride) + self.norm3 = nn.GroupNorm(32, num_filters) + + def forward(self, x, is_init=False) -> torch.Tensor: + x = x.permute(0, 2, 3, 4, 1).contiguous() + + residual = x + + h = self.conv1(x, is_init) + h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) + + h = self.conv2(h, is_init) + h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) + + if self.down_sampling or self.num_filters != self.input_filters: + x = self.conv3(x, is_init) + x = base_group_norm(x, + self.norm3, + act_silu=False, + channel_last=True) + + h.add_(x) + h = self.act_(h) + if residual is not None: + h.add_(residual) + + h = h.permute(0, 4, 1, 2, 3) + return h + + +class Upsample3D(nn.Module): + + def __init__(self, in_channels, scale_factor=2) -> None: + super().__init__() + + self.scale_factor = scale_factor + self.conv3d = Res3DBlockUpsample(input_filters=in_channels, + num_filters=in_channels, + down_sampling_stride=(1, 1, 1), + down_sampling=False) + + def forward(self, x, is_init=True, is_split=True) -> torch.Tensor: + b, c, t, h, w = x.shape + + # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) + if is_split: + split_size = c // 8 + x_slices = torch.split(x, split_size, dim=1) + x = [ + nn.functional.interpolate(x, scale_factor=self.scale_factor) + for x in x_slices + ] + x = torch.cat(x, dim=1) + else: + x = nn.functional.interpolate(x, scale_factor=self.scale_factor) + + x = self.conv3d(x, is_init) + return x + + +class VideoDecoder(nn.Module): + + def __init__( + self, + ch=128, + z_channels=16, + out_channels=3, + ch_mult=(1, 2, 4, 4), + num_res_blocks=2, + temporal_up_layers=(2, 3), + temporal_downsample=4, + resamp_with_conv=True, + version=1, + ) -> None: + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.temporal_downsample = temporal_downsample + + block_in = ch * ch_mult[self.num_resolutions - 1] + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_in = CausalConv(z_channels, channels, kernel_size=3) + self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( + z_channels, channels, 1) + self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D( + channels, block_in, kernel_size=3, factor=2) + self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( + channels, block_in, 2) + else: + self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=temb_ch) + + # upsampling + self.up_id = len(temporal_up_layers) + self.video_frame_num = 1 + self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1 + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + Resnet3DBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=temb_ch)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + if i_level in temporal_up_layers: + up.upsample = Upsample3D(block_in) + self.cur_video_frame_num = self.cur_video_frame_num * 2 + else: + up.upsample = Upsample2D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.conv_out = CausalConvAfterNorm(block_in, + out_channels, + kernel_size=3) + + @torch.inference_mode() + def forward(self, z, is_init=True) -> torch.Tensor: + z = rearrange(z, "b t c h w -> b c t h w") + h = self.conv_in(z, is_init=is_init) + if self.version == 2: + shortcut = self.shortcut_in(z, is_init=is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_unpathify(h, is_init=is_init) + h = self.conv_unpatchify(h, is_init=is_init) + h = h.add_(shortcut) + + temb = None + + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_1(h, temb, is_init=is_init) + h = self.mid.attn_1(h) + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_2(h, temb, is_init=is_init) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.up[i_level].block[i_block](h, temb, is_init=is_init) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + if isinstance(self.up[i_level].upsample, Upsample2D): + B = h.size(0) + h = h.permute(0, 2, 3, 4, 1).flatten(0, 1) + h = self.up[i_level].upsample(h) + h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3) + else: + h = self.up[i_level].upsample(h, is_init=is_init) + + # end + h = h.permute(0, 2, 3, 4, 1) # b c l h w -> b l h w c + h = base_group_norm_with_zero_pad(h, + self.norm_out, + act_silu=True, + pad_size=2) + h = self.conv_out(h) + h = h.permute(0, 4, 1, 2, 3) + + if is_init: + h = h[:, :, (self.temporal_downsample - 1):] + return h + + +def rms_norm(input, normalized_shape, eps=1e-6) -> torch.Tensor: + dtype = input.dtype + input = input.to(torch.float32) + variance = input.pow(2).flatten(-len(normalized_shape)).mean( + -1)[(..., ) + (None, ) * len(normalized_shape)] + input = input * torch.rsqrt(variance + eps) + return input.to(dtype) + + +class DiagonalGaussianDistribution: + + def __init__(self, + parameters, + deterministic=False, + rms_norm_mean=False, + only_return_mean=False) -> None: + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, + dim=-3) #N,[X],C,H,W + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + self.deterministic = deterministic + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + if rms_norm_mean: + self.mean = rms_norm(self.mean, self.mean.size()[1:]) + self.only_return_mean = only_return_mean + + def sample(self, generator=None) -> torch.Tensor: + # make sure sample is on the same device + # as the parameters and has same dtype + sample = torch.randn(self.mean.shape, + generator=generator, + device=self.parameters.device) + sample = sample.to(dtype=self.parameters.dtype) + x = self.mean + self.std * sample + if self.only_return_mean: + return self.mean + else: + return x + + +class AutoencoderKLStepvideo(nn.Module, ParallelTiledVAE): + + def __init__( + self, + config: StepVideoVAEConfig, + ) -> None: + nn.Module.__init__(self) + ParallelTiledVAE.__init__(self, config) + + self.frame_len = config.frame_len + + if config.version == 2: + self.latent_len = 3 + base_group_norm.spatial = True # type: ignore[attr-defined] + else: + self.latent_len = 5 + base_group_norm.spatial = False # type: ignore[attr-defined] + + self.encoder = VideoEncoder( + in_channels=config.in_channels, + z_channels=config.z_channels, + num_res_blocks=config.num_res_blocks, + version=config.version, + ) + + self.decoder = VideoDecoder( + z_channels=config.z_channels, + out_channels=config.out_channels, + num_res_blocks=config.num_res_blocks, + version=config.version, + ) + + self.world_size = config.world_size + # self.is_init = True + + def load_state_dict(self, state_dict, strict=True): + remapped = {} + for key, value in state_dict.items(): + if key.startswith("decoder.conv_out."): + # move “decoder.conv_out.weight” → “decoder.conv_out.conv.weight” + suffix = key[len("decoder.conv_out."):] + remapped[f"decoder.conv_out.conv.{suffix}"] = value + else: + remapped[key] = value + super().load_state_dict(remapped, strict=strict) + + def _encode(self, x, is_init_image=True) -> torch.Tensor: + # b, len, c, h, w = x.size() + b, c, len, h, w = x.size() + # x = rearrange(x, 'b l c h w -> b c l h w').contiguous() + z = self.encoder(x, len, True) # 下采样[1, 4, 8, 16, 16] + return z + + @torch.inference_mode() + def encode(self, x): + # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w + chunks = list(x.split(self.frame_len, dim=1)) + for i in range(len(chunks)): + chunks[i] = self._encode(chunks[i], True) + z = torch.cat(chunks, dim=1) + + posterior = DiagonalGaussianDistribution(z) + return posterior.sample() + + def _decode(self, z) -> torch.Tensor: + + chunks = list(z.split(self.latent_len, dim=2)) + for i in range(len(chunks)): + chunks[i] = chunks[i].permute(0, 2, 1, 3, 4) + chunks[i] = chunks[i].to(next(self.decoder.parameters()).dtype) + chunks[i] = self.decoder(chunks[i], is_init=True) + x = torch.cat(chunks, dim=2) + return x + + def decode(self, z) -> torch.Tensor: + num_frames = z.size(2) + dec = ParallelTiledVAE.decode(self, z).permute(0, 2, 1, 3, 4) + dec = self.mix(dec).permute(0, 2, 1, 3, 4) + num_sample_frames = num_frames // 3 * 17 + print(dec.shape) + return dec[:, :, :num_sample_frames] + + def mix(self, x) -> torch.Tensor: + remain_scale = 0.6 + mix_scale = 1. - remain_scale + front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) + back = slice(self.frame_len, x.size(1), self.frame_len) + x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale + x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale + return x + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """ + Args: + sample (`torch.Tensor`): Input sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z) + return dec diff --git a/trainer/models/vaes/wanvae.py b/trainer/models/vaes/wanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..83117e22d8b05ae59ec6fe8398ae6a88a15aa75e --- /dev/null +++ b/trainer/models/vaes/wanvae.py @@ -0,0 +1,1273 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextvars +from contextlib import contextmanager + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from trainer.configs.models.vaes import WanVAEConfig +from trainer.layers.activation import get_act_fn +from trainer.models.vaes.common import (DiagonalGaussianDistribution, + ParallelTiledVAE) +from trainer.platforms import current_platform + +CACHE_T = 2 + +is_first_frame = contextvars.ContextVar("is_first_frame", default=False) +feat_cache = contextvars.ContextVar("feat_cache", default=None) +feat_idx = contextvars.ContextVar("feat_idx", default=0) +first_chunk = contextvars.ContextVar("first_chunk", default=None) + + +@contextmanager +def forward_context(first_frame_arg=False, + feat_cache_arg=None, + feat_idx_arg=None, + first_chunk_arg=None): + is_first_frame_token = is_first_frame.set(first_frame_arg) + feat_cache_token = feat_cache.set(feat_cache_arg) + feat_idx_token = feat_idx.set(feat_idx_arg) + first_chunk_token = first_chunk.set(first_chunk_arg) + try: + yield + finally: + is_first_frame.reset(is_first_frame_token) + feat_cache.reset(feat_cache_token) + feat_idx.reset(feat_idx_token) + first_chunk.reset(first_chunk_token) + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + + _first_chunk = first_chunk.get() + if _first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + +class WanCausalConv3d(nn.Conv3d): + r""" + A custom 3D causal convolution layer with feature caching support. + + This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature + caching for efficient inference. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int or tuple, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + self.padding: tuple[int, int, int] + # Set up causal padding + self._padding: tuple[int, ...] = (self.padding[2], self.padding[2], + self.padding[1], self.padding[1], + 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + x = x.to(self.weight.dtype) if current_platform.is_mps( + ) else x # casting needed for mps since amp isn't supported + return super().forward(x) + + +class WanRMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, + dim: int, + channel_first: bool = True, + images: bool = True, + bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim, ) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class WanUpsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class WanResample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + - 'downsample2d': 2D downsampling with zero-padding and convolution. + - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: + super().__init__() + self.dim = dim + self.mode = mode + + # default to dim //2 + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + self.time_conv = WanCausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = WanCausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x): + b, c, t, h, w = x.size() + first_frame = is_first_frame.get() + if first_frame: + assert t == 1 + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if self.mode == "upsample3d": + if _feat_cache is not None: + idx = _feat_idx + if _feat_cache[idx] is None: + _feat_cache[idx] = "Rep" + _feat_idx += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[ + idx] is not None and _feat_cache[idx] != "Rep": + # cache last frame of last two chunk + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and _feat_cache[ + idx] is not None and _feat_cache[idx] == "Rep": + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if _feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + elif not first_frame and hasattr(self, "time_conv"): + x = self.time_conv(x) + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if self.mode == "downsample3d": + if _feat_cache is not None: + idx = _feat_idx + if _feat_cache[idx] is None: + _feat_cache[idx] = x.clone() + _feat_idx += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2)) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + elif not first_frame and hasattr(self, "time_conv"): + x = self.time_conv(x) + return x + + +class WanResidualBlock(nn.Module): + r""" + A custom residual block module. + + Args: + in_dim (int): Number of input channels. + out_dim (int): Number of output channels. + dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. + non_linearity (str, optional): Type of non-linearity to use. Default is "silu". + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + non_linearity: str = "silu", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = get_act_fn(non_linearity) + + # layers + self.norm1 = WanRMS_norm(in_dim, images=False) + self.conv1 = WanCausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = WanRMS_norm(out_dim, images=False) + self.dropout = nn.Dropout(dropout) + self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = WanCausalConv3d( + in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x): + # Apply shortcut connection + h = self.conv_shortcut(x) + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + + x = self.conv1(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv1(x) + + # Second normalization and activation + x = self.norm2(x) + x = self.nonlinearity(x) + + # Dropout + x = self.dropout(x) + + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + + x = self.conv2(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv2(x) + + # Add residual connection + return x + h + + +class WanAttentionBlock(nn.Module): + r""" + Causal self-attention with a single head. + + Args: + dim (int): The number of channels in the input tensor. + """ + + def __init__(self, dim) -> None: + super().__init__() + self.dim = dim + + # layers + self.norm = WanRMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + def forward(self, x): + identity = x + batch_size, channels, time, height, width = x.size() + + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, + height, width) + x = self.norm(x) + + # compute query, key, value + qkv = self.to_qkv(x) + qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) + qkv = qkv.permute(0, 1, 3, 2).contiguous() + q, k, v = qkv.chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention(q, k, v) + + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, + height, width) + + # output projection + x = self.proj(x) + + # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] + x = x.view(batch_size, time, channels, height, width) + x = x.permute(0, 2, 1, 3, 4) + + return x + identity + + +class WanMidBlock(nn.Module): + """ + Middle block for WanVAE encoder and decoder. + + Args: + dim (int): Number of input/output channels. + dropout (float): Dropout rate. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__(self, + dim: int, + dropout: float = 0.0, + non_linearity: str = "silu", + num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)] + attentions = [] + for _ in range(num_layers): + attentions.append(WanAttentionBlock(dim)) + resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x): + # First residual block + x = self.resnets[0](x) + + # Process through attention and residual blocks + for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): + if attn is not None: + x = attn(x) + + x = resnet(x) + + return x + +class WanResidualDownBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + resnets = [] + for _ in range(num_res_blocks): + resnets.append(WanResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + self.resnets = nn.ModuleList(resnets) + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + self.downsampler = WanResample(out_dim, mode=mode) + else: + self.downsampler = None + + def forward(self, x): + x_copy = x.clone() + for resnet in self.resnets: + x = resnet(x) + if self.downsampler is not None: + x = self.downsampler(x) + + return x + self.avg_shortcut(x_copy) + +class WanEncoder3d(nn.Module): + r""" + A 3D encoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_downsample (list of bool): Whether to downsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + in_channels: int = 3, + dim=128, + z_dim=4, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(True, True, False), + dropout=0.0, + non_linearity: str = "silu", + is_residual: bool = False, # wan 2.2 vae use a residual downblock + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + dim_mult = list(dim_mult) + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = list(attn_scales) + self.temperal_downsample = list(temperal_downsample) + self.nonlinearity = get_act_fn(non_linearity) + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1) + + # downsample blocks + self.down_blocks = nn.ModuleList([]) + for i, (in_dim, + out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): + # residual (+attention) blocks + if is_residual: + self.down_blocks.append( + WanResidualDownBlock( + in_dim, + out_dim, + dropout, + num_res_blocks, + temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False, + down_flag=i != len(dim_mult) - 1, + ) + ) + else: + for _ in range(num_res_blocks): + self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + self.down_blocks.append(WanAttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + self.down_blocks.append(WanResample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = WanMidBlock(out_dim, + dropout, + non_linearity, + num_layers=1) + + # output blocks + self.norm_out = WanRMS_norm(out_dim, images=False) + self.conv_out = WanCausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x): + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv_in(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv_in(x) + + ## downsamples + for layer in self.down_blocks: + x = layer(x) + + ## middle + x = self.mid_block(x) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv_out(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv_out(x) + return x + +# adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +class WanResidualUpBlock(nn.Module): + """ + A block that handles upsampling for the WanVAE decoder. + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + temperal_upsample (bool): Whether to upsample on temporal dimension + up_flag (bool): Whether to upsample or not + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + temperal_upsample: bool = False, + up_flag: bool = False, + non_linearity: str = "silu", + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2, + ) + else: + self.avg_shortcut = None + + # create residual blocks + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + if up_flag: + upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" + self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) + else: + self.upsampler = None + + self.gradient_checkpointing = False + + def forward(self, x): + """ + Forward pass through the upsampling block. + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + Returns: + torch.Tensor: Output tensor + """ + if self.avg_shortcut is not None: + x_copy = x.clone() + + for resnet in self.resnets: + x = resnet(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + if self.avg_shortcut is not None: + x = x + self.avg_shortcut(x_copy) + + return x + +class WanUpBlock(nn.Module): + """ + A block that handles upsampling for the WanVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + non_linearity: str = "silu", + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # Create layers list + resnets = [] + # Add residual blocks and attention if needed + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append( + WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + self.upsamplers = None + if upsample_mode is not None: + self.upsamplers = nn.ModuleList( + [WanResample(out_dim, mode=upsample_mode)]) + + self.gradient_checkpointing = False + + def forward(self, x): + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + for resnet in self.resnets: + x = resnet(x) + + if self.upsamplers is not None: + x = self.upsamplers[0](x) + return x + + +class WanDecoder3d(nn.Module): + r""" + A 3D decoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_upsample (list of bool): Whether to upsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_upsample=(False, True, True), + dropout=0.0, + non_linearity: str = "silu", + out_channels: int = 3, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + dim_mult = list(dim_mult) + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = list(attn_scales) + self.temperal_upsample = list(temperal_upsample) + + self.nonlinearity = get_act_fn(non_linearity) + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + + # init block + self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = WanMidBlock(dims[0], + dropout, + non_linearity, + num_layers=1) + + # upsample blocks + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, + out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): + # residual (+attention) blocks + if i > 0 and not is_residual: + # wan vae 2.1 + in_dim = in_dim // 2 + + # determine if we need upsampling + up_flag = i != len(dim_mult) - 1 + # determine upsampling mode, if not upsampling, set to None + upsample_mode = None + if up_flag and temperal_upsample[i]: + upsample_mode = "upsample3d" + elif up_flag: + upsample_mode = "upsample2d" + + # Create and add the upsampling block + if is_residual: + up_block = WanResidualUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + temperal_upsample=temperal_upsample[i] if up_flag else False, + up_flag=up_flag, + non_linearity=non_linearity, + ) + else: + up_block = WanUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + non_linearity=non_linearity, + ) + self.up_blocks.append(up_block) + + # output blocks + self.norm_out = WanRMS_norm(out_dim, images=False) + self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x): + ## conv1 + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv_in(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv_in(x) + + ## middle + x = self.mid_block(x) + + ## upsamples + for up_block in self.up_blocks: + x = up_block(x) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + _feat_cache = feat_cache.get() + _feat_idx = feat_idx.get() + if _feat_cache is not None: + idx = _feat_idx + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv_out(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + feat_cache.set(_feat_cache) + feat_idx.set(_feat_idx) + else: + x = self.conv_out(x) + return x + +def patchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() == 4: + x = rearrange( + x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() == 4: + x = rearrange( + x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size, + ) + + return x + +class AutoencoderKLWan(nn.Module, ParallelTiledVAE): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [Wan 2.1]. + """ + + _supports_gradient_checkpointing = False + + def __init__( + self, + config: WanVAEConfig, + ) -> None: + nn.Module.__init__(self) + ParallelTiledVAE.__init__(self, config) + + self.z_dim = config.z_dim + self.temperal_downsample = list(config.temperal_downsample) + self.temperal_upsample = list(config.temperal_downsample)[::-1] + + if config.decoder_base_dim is None: + decoder_base_dim = config.base_dim + else: + decoder_base_dim = config.decoder_base_dim + + self.latents_mean = list(config.latents_mean) + self.latents_std = list(config.latents_std) + self.shift_factor = config.shift_factor + + if config.load_encoder: + self.encoder = WanEncoder3d( + in_channels=config.in_channels, + dim=config.base_dim, + z_dim=self.z_dim * 2, + dim_mult=config.dim_mult, + num_res_blocks=config.num_res_blocks, + attn_scales=config.attn_scales, + temperal_downsample=self.temperal_downsample, + dropout=config.dropout, + is_residual=config.is_residual, + ) + self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1) + self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1) + + if config.load_decoder: + self.decoder = WanDecoder3d( + dim=decoder_base_dim, + z_dim=self.z_dim, + dim_mult=config.dim_mult, + num_res_blocks=config.num_res_blocks, + attn_scales=config.attn_scales, + temperal_upsample=self.temperal_upsample, + dropout=config.dropout, + out_channels=config.out_channels, + is_residual=config.is_residual, + ) + + self.use_feature_cache = config.use_feature_cache + + def clear_cache(self) -> None: + + def _count_conv3d(model) -> int: + count = 0 + for m in model.modules(): + if isinstance(m, WanCausalConv3d): + count += 1 + return count + + if self.config.load_decoder: + self._conv_num = _count_conv3d(self.decoder) + self._conv_idx = 0 + self._feat_map = [None] * self._conv_num + # cache encode + if self.config.load_encoder: + self._enc_conv_num = _count_conv3d(self.encoder) + self._enc_conv_idx = 0 + self._enc_feat_map = [None] * self._enc_conv_num + + def encode(self, x: torch.Tensor) -> torch.Tensor: + if self.use_feature_cache: + self.clear_cache() + if self.config.patch_size is not None: + x = patchify(x, patch_size=self.config.patch_size) + with forward_context(feat_cache_arg=self._enc_feat_map, + feat_idx_arg=self._enc_conv_idx): + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + feat_idx.set(0) + if i == 0: + out = self.encoder(x[:, :, :1, :, :]) + else: + out_ = self.encoder(x[:, :, + 1 + 4 * (i - 1):1 + 4 * i, :, :]) + out = torch.cat([out, out_], 2) + enc = self.quant_conv(out) + mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, + self.z_dim:, :, :, :] + enc = torch.cat([mu, logvar], dim=1) + enc = DiagonalGaussianDistribution(enc) + self.clear_cache() + else: + for block in self.encoder.down_blocks: + if isinstance(block, + WanResample) and block.mode == "downsample3d": + _padding = list(block.time_conv._padding) + _padding[4] = 2 + block.time_conv._padding = tuple(_padding) + enc = ParallelTiledVAE.encode(self, x) + + return enc + + def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor: + with forward_context(first_frame_arg=first_frame): + out = self.encoder(x) + enc = self.quant_conv(out) + mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] + enc = torch.cat([mu, logvar], dim=1) + return enc + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + first_frame = x[:, :, 0, :, :].unsqueeze(2) + first_frame = self._encode(first_frame, first_frame=True) + + enc = ParallelTiledVAE.tiled_encode(self, x) + enc = enc[:, :, 1:] + enc = torch.cat([first_frame, enc], dim=2) + return enc + + def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + first_frame = x[:, :, 0, :, :].unsqueeze(2) + first_frame = self._encode(first_frame, first_frame=True) + + enc = ParallelTiledVAE.spatial_tiled_encode(self, x) + enc = enc[:, :, 1:] + enc = torch.cat([first_frame, enc], dim=2) + return enc + + def decode(self, z: torch.Tensor) -> torch.Tensor: + if self.use_feature_cache: + self.clear_cache() + iter_ = z.shape[2] + x = self.post_quant_conv(z) + with forward_context(feat_cache_arg=self._feat_map, + feat_idx_arg=self._conv_idx): + for i in range(iter_): + feat_idx.set(0) + if i == 0: + first_chunk.set(True) + out = self.decoder(x[:, :, i:i + 1, :, :]) + else: + first_chunk.set(False) + out_ = self.decoder(x[:, :, i:i + 1, :, :]) + out = torch.cat([out, out_], 2) + + if self.config.patch_size is not None: + out = unpatchify(out, patch_size=self.config.patch_size) + + out = out.float() + out = torch.clamp(out, min=-1.0, max=1.0) + self.clear_cache() + else: + out = ParallelTiledVAE.decode(self, z) + + return out + + def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: + x = self.post_quant_conv(z) + with forward_context(first_frame_arg=first_frame): + out = self.decoder(x) + + out = torch.clamp(out, min=-1.0, max=1.0) + + return out + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + self.blend_num_frames *= 2 + dec = ParallelTiledVAE.tiled_decode(self, z) + start_frame_idx = self.temporal_compression_ratio - 1 + dec = dec[:, :, start_frame_idx:] + return dec + + def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + dec = ParallelTiledVAE.spatial_tiled_decode(self, z) + start_frame_idx = self.temporal_compression_ratio - 1 + dec = dec[:, :, start_frame_idx:] + return dec + + def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: + self.blend_num_frames *= 2 + dec = ParallelTiledVAE.parallel_tiled_decode(self, z) + start_frame_idx = self.temporal_compression_ratio - 1 + dec = dec[:, :, start_frame_idx:] + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """ + Args: + sample (`torch.Tensor`): Input sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z) + return dec diff --git a/trainer/models/vision_utils.py b/trainer/models/vision_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d2a5e131a8bc225486add8145496ecdb44dc2ffd --- /dev/null +++ b/trainer/models/vision_utils.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import tempfile +from collections.abc import Callable +from urllib.parse import unquote, urlparse + +import imageio +import numpy as np +import PIL.Image +import PIL.ImageOps +import requests +import torch +from packaging import version + +if version.parse(version.parse( + PIL.__version__).base_version) >= version.parse("9.1.0"): + PIL_INTERPOLATION = { + "linear": PIL.Image.Resampling.BILINEAR, + "bilinear": PIL.Image.Resampling.BILINEAR, + "bicubic": PIL.Image.Resampling.BICUBIC, + "lanczos": PIL.Image.Resampling.LANCZOS, + "nearest": PIL.Image.Resampling.NEAREST, + } +else: + PIL_INTERPOLATION = { + "linear": PIL.Image.LINEAR, + "bilinear": PIL.Image.BILINEAR, + "bicubic": PIL.Image.BICUBIC, + "lanczos": PIL.Image.LANCZOS, + "nearest": PIL.Image.NEAREST, + } + + +def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray: + r""" + Convert a PIL image or a list of PIL images to NumPy arrays. + + Args: + images (`PIL.Image.Image` or `List[PIL.Image.Image]`): + The PIL image or list of images to convert to NumPy format. + + Returns: + `np.ndarray`: + A NumPy array representation of the images. + """ + if not isinstance(images, list): + images = [images] + images = [np.array(image).astype(np.float32) / 255.0 for image in images] + images_arr: np.ndarray = np.stack(images, axis=0) + + return images_arr + + +def numpy_to_pt(images: np.ndarray) -> torch.Tensor: + r""" + Convert a NumPy image to a PyTorch tensor. + + Args: + images (`np.ndarray`): + The NumPy image array to convert to PyTorch format. + + Returns: + `torch.Tensor`: + A PyTorch tensor representation of the images. + """ + if images.ndim == 3: + images = images[..., None] + + images = torch.from_numpy(images.transpose(0, 3, 1, 2)) + return images + + +def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + r""" + Normalize an image array to [-1,1]. + + Args: + images (`np.ndarray` or `torch.Tensor`): + The image array to normalize. + + Returns: + `np.ndarray` or `torch.Tensor`: + The normalized image array. + """ + return 2.0 * images - 1.0 + + +# adapted from diffusers.utils import load_image +def load_image( + image: str | PIL.Image.Image, + convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None +) -> PIL.Image.Image: + """ + Loads `image` to a PIL Image. + + Args: + image (`str` or `PIL.Image.Image`): + The image to convert to the PIL Image format. + convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*): + A conversion method to apply to the image after loading it. When set to `None` the image will be converted + "RGB". + + Returns: + `PIL.Image.Image`: + A PIL Image. + """ + if isinstance(image, str): + if image.startswith("http://") or image.startswith("https://"): + image = PIL.Image.open(requests.get(image, stream=True).raw) + elif os.path.isfile(image): + image = PIL.Image.open(image) + else: + raise ValueError( + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." + ) + elif isinstance(image, PIL.Image.Image): + image = image + else: + raise ValueError( + "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." + ) + + image = PIL.ImageOps.exif_transpose(image) + + if convert_method is not None: + image = convert_method(image) + else: + image = image.convert("RGB") + + return image + + +# adapted from diffusers.utils import load_video +def load_video( + video: str, + convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] + | None = None, +) -> list[PIL.Image.Image]: + """ + Loads `video` to a list of PIL Image. + Args: + video (`str`): + A URL or Path to a video to convert to a list of PIL Image format. + convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*): + A conversion method to apply to the video after loading it. When set to `None` the images will be converted + to "RGB". + Returns: + `List[PIL.Image.Image]`: + The video as a list of PIL images. + """ + is_url = video.startswith("http://") or video.startswith("https://") + is_file = os.path.isfile(video) + was_tempfile_created = False + + if not (is_url or is_file): + raise ValueError( + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." + ) + + if is_url: + response = requests.get(video, stream=True) + if response.status_code != 200: + raise ValueError( + f"Failed to download video. Status code: {response.status_code}" + ) + + parsed_url = urlparse(video) + file_name = os.path.basename(unquote(parsed_url.path)) + + suffix = os.path.splitext(file_name)[1] or ".mp4" + with tempfile.NamedTemporaryFile(suffix=suffix, + delete=False) as temp_file: + video_path = temp_file.name + video_data = response.iter_content(chunk_size=8192) + for chunk in video_data: + temp_file.write(chunk) + + video = video_path + + pil_images = [] + if video.endswith(".gif"): + gif = PIL.Image.open(video) + try: + while True: + pil_images.append(gif.copy()) + gif.seek(gif.tell() + 1) + except EOFError: + pass + + else: + try: + imageio.plugins.ffmpeg.get_exe() + except AttributeError: + raise AttributeError( + "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" + ) from None + + with imageio.get_reader(video) as reader: + # Read all frames + for frame in reader: + pil_images.append(PIL.Image.fromarray(frame)) + + if was_tempfile_created: + os.remove(video_path) + + if convert_method is not None: + pil_images = convert_method(pil_images) + + return pil_images + + +def get_default_height_width( + image: PIL.Image.Image | np.ndarray | torch.Tensor, + vae_scale_factor: int, + height: int | None = None, + width: int | None = None, +) -> tuple[int, int]: + r""" + Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. + + Args: + image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it + should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch + tensor, it should have shape `[batch, channels, height, width]`. + height (`Optional[int]`, *optional*, defaults to `None`): + The height of the preprocessed image. If `None`, the height of the `image` input will be used. + width (`Optional[int]`, *optional*, defaults to `None`): + The width of the preprocessed image. If `None`, the width of the `image` input will be used. + + Returns: + `Tuple[int, int]`: + A tuple containing the height and width, both resized to the nearest integer multiple of + `vae_scale_factor`. + """ + + if height is None: + if isinstance(image, PIL.Image.Image): + height = image.height + elif isinstance(image, torch.Tensor): + height = image.shape[2] + else: + height = image.shape[1] + + if width is None: + if isinstance(image, PIL.Image.Image): + width = image.width + elif isinstance(image, torch.Tensor): + width = image.shape[3] + else: + width = image.shape[2] + + width, height = (x - x % vae_scale_factor for x in (width, height) + ) # resize to integer multiple of vae_scale_factor + + return height, width + + +def resize( + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int, + width: int, + resize_mode: str = "default", # "default", "fill", "crop" + resample: str = "lanczos", +) -> PIL.Image.Image | np.ndarray | torch.Tensor: + """ + Resize image. + + Args: + image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`): + The image input, can be a PIL image, numpy array or pytorch tensor. + height (`int`): + The height to resize to. + width (`int`): + The width to resize to. + resize_mode (`str`, *optional*, defaults to `default`): + The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit + within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, + will resize the image to fit within the specified width and height, maintaining the aspect ratio, and + then center the image within the dimensions, filling empty with data from image. If `crop`, will resize + the image to fit within the specified width and height, maintaining the aspect ratio, and then center + the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only + supported for PIL image input. + + Returns: + `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: + The resized image. + """ + if resize_mode != "default" and not isinstance(image, PIL.Image.Image): + raise ValueError( + f"Only PIL image input is supported for resize_mode {resize_mode}") + assert isinstance(image, PIL.Image.Image) + if resize_mode == "default": + image = image.resize((width, height), + resample=PIL_INTERPOLATION[resample]) + else: + raise ValueError(f"resize_mode {resize_mode} is not supported") + return image diff --git a/trainer/pipelines/README.md b/trainer/pipelines/README.md new file mode 100644 index 0000000000000000000000000000000000000000..fe3358045660e48271c7fca0ee1c7873ca0b193b --- /dev/null +++ b/trainer/pipelines/README.md @@ -0,0 +1,3 @@ +# Adding a New Custom Pipeline + +Please see documentation [here](https://hao-ai-lab.github.io/Trainer/contributing/add_pipeline.html) diff --git a/trainer/pipelines/__init__.py b/trainer/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..81cca683718bdcd881d3cf6e33de9056578e1137 --- /dev/null +++ b/trainer/pipelines/__init__.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Diffusion pipelines for trainer. + +This package contains diffusion pipelines for generating videos and images. +""" + +from typing import cast + +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase +from trainer.pipelines.lora_pipeline import LoRAPipeline +from trainer.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from trainer.pipelines.pipeline_registry import (PipelineType, + get_pipeline_registry) +from trainer.utils import (maybe_download_model, + verify_model_config_and_directory) + +logger = init_logger(__name__) + + +class PipelineWithLoRA(LoRAPipeline, ComposedPipelineBase): + """Type for a pipeline that has both ComposedPipelineBase and LoRAPipeline functionality.""" + pass + + +def build_pipeline( + trainer_args: TrainerArgs, + pipeline_type: PipelineType | str = PipelineType.BASIC +) -> PipelineWithLoRA: + """ + Only works with valid hf diffusers configs. (model_index.json) + We want to build a pipeline based on the inference args mode_path: + 1. download the model from the hub if it's not already downloaded + 2. verify the model config and directory + 3. based on the config, determine the pipeline class + """ + # Get pipeline type + model_path = trainer_args.model_path + model_path = maybe_download_model(model_path) + # trainer_args.downloaded_model_path = model_path + logger.info("Model path: %s", model_path) + + config = verify_model_config_and_directory(model_path) + pipeline_name = config.get("_class_name") + if pipeline_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") + + # Get the appropriate pipeline registry based on pipeline_type + logger.info( + "Building pipeline of type: %s", pipeline_type.value if isinstance( + pipeline_type, PipelineType) else pipeline_type) + pipeline_registry = get_pipeline_registry(pipeline_type) + + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_cls = pipeline_registry.resolve_pipeline_cls( + pipeline_name, pipeline_type, trainer_args.workload_type) + + # instantiate the pipelines + pipeline = pipeline_cls(model_path, trainer_args) + + logger.info("Pipelines instantiated") + + return cast(PipelineWithLoRA, pipeline) + + +__all__ = [ + "build_pipeline", + "ComposedPipelineBase", + "ForwardBatch", + "LoRAPipeline", + "TrainingBatch", +] diff --git a/trainer/pipelines/base.py b/trainer/pipelines/base.py new file mode 100644 index 0000000000000000000000000000000000000000..a9612836a21803092d104ac009b073b69a921a30 --- /dev/null +++ b/trainer/pipelines/base.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Base classes for pipeline stages. + +This module defines the abstract base classes for pipeline stages that can be +composed to create complete diffusion pipelines. +""" + +import time +import traceback +from abc import ABC, abstractmethod + +import torch + +import trainer.envs as envs +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines.pipeline_batch_info import ForwardBatch +from trainer.pipelines.validators import VerificationResult + +logger = init_logger(__name__) + + +class StageVerificationError(Exception): + """Exception raised when stage verification fails.""" + pass + + +class PipelineStage(ABC): + """ + Abstract base class for all pipeline stages. + + A pipeline stage represents a discrete step in the diffusion process that can be + composed with other stages to create a complete pipeline. Each stage is responsible + for a specific part of the process, such as prompt encoding, latent preparation, etc. + """ + + def verify_input(self, batch: ForwardBatch, + trainer_args: TrainerArgs) -> VerificationResult: + """ + Verify the input for the stage. + + Example: + from trainer.pipelines.stages.validators import V, VerificationResult + + def verify_input(self, batch, trainer_args): + result = VerificationResult() + result.add_check("height", batch.height, V.positive_int_divisible(8)) + result.add_check("width", batch.width, V.positive_int_divisible(8)) + result.add_check("image_latent", batch.image_latent, V.is_tensor) + return result + + Args: + batch: The current batch information. + trainer_args: The inference arguments. + + Returns: + A VerificationResult containing the verification status. + + """ + # Default implementation - no verification + return VerificationResult() + + def verify_output(self, batch: ForwardBatch, + trainer_args: TrainerArgs) -> VerificationResult: + """ + Verify the output for the stage. + + Args: + batch: The current batch information. + trainer_args: The inference arguments. + + Returns: + A VerificationResult containing the verification status. + """ + # Default implementation - no verification + return VerificationResult() + + def _run_verification(self, verification_result: VerificationResult, + stage_name: str, verification_type: str) -> None: + """ + Run verification and raise errors if any checks fail. + + Args: + verification_result: Results from verify_input or verify_output + stage_name: Name of the current stage + verification_type: "input" or "output" + """ + if not verification_result.is_valid(): + failed_fields = verification_result.get_failed_fields() + if failed_fields: + # Get detailed failure information + detailed_summary = verification_result.get_failure_summary() + + failed_fields_str = ", ".join(failed_fields) + error_msg = ( + f"{verification_type.capitalize()} verification failed for {stage_name}: " + f"Failed fields: {failed_fields_str}\n" + f"Details: {detailed_summary}") + raise StageVerificationError(error_msg) + + @property + def device(self) -> torch.device: + """Get the device for this stage.""" + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def set_logging(self, enable: bool): + """ + Enable or disable logging for this stage. + + Args: + enable: Whether to enable logging. + """ + self._enable_logging = enable + + def __call__( + self, + batch: ForwardBatch, + trainer_args: TrainerArgs, + ) -> ForwardBatch: + """ + Execute the stage's processing on the batch with optional verification and logging. + Should not be overridden by subclasses. + + Args: + batch: The current batch information. + trainer_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + stage_name = self.__class__.__name__ + + # Check if verification is enabled (simple approach for prototype) + enable_verification = getattr(trainer_args, + 'enable_stage_verification', False) + + if enable_verification: + # Pre-execution input verification + try: + input_result = self.verify_input(batch, trainer_args) + self._run_verification(input_result, stage_name, "input") + except Exception as e: + logger.error("Input verification failed for %s: %s", stage_name, + str(e)) + raise + + # Execute the actual stage logic + if envs.TRAINER_STAGE_LOGGING: + logger.info("[%s] Starting execution", stage_name) + start_time = time.perf_counter() + + try: + result = self.forward(batch, trainer_args) + execution_time = time.perf_counter() - start_time + logger.info("[%s] Execution completed in %s ms", stage_name, + execution_time * 1000) + batch.logging_info.add_stage_execution_time( + stage_name, execution_time) + except Exception as e: + execution_time = time.perf_counter() - start_time + logger.error("[%s] Error during execution after %s ms: %s", + stage_name, execution_time * 1000, e) + logger.error("[%s] Traceback: %s", stage_name, + traceback.format_exc()) + raise + else: + # Direct execution (current behavior) + result = self.forward(batch, trainer_args) + + if enable_verification: + # Post-execution output verification + try: + output_result = self.verify_output(result, trainer_args) + self._run_verification(output_result, stage_name, "output") + except Exception as e: + logger.error("Output verification failed for %s: %s", + stage_name, str(e)) + raise + + return result + + @abstractmethod + def forward( + self, + batch: ForwardBatch, + trainer_args: TrainerArgs, + ) -> ForwardBatch: + """ + Forward pass of the stage's processing. + + This method should be implemented by subclasses to provide the forward + processing logic for the stage. + + Args: + batch: The current batch information. + trainer_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + raise NotImplementedError + + def backward( + self, + batch: ForwardBatch, + trainer_args: TrainerArgs, + ) -> ForwardBatch: + raise NotImplementedError diff --git a/trainer/pipelines/composed_pipeline_base.py b/trainer/pipelines/composed_pipeline_base.py new file mode 100644 index 0000000000000000000000000000000000000000..4943f02da020894f54ca199c91f0c3e17487dfd4 --- /dev/null +++ b/trainer/pipelines/composed_pipeline_base.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Base class for composed pipelines. + +This module defines the base class for pipelines that are composed of multiple stages. +""" + +import argparse +import os +from abc import ABC, abstractmethod +from typing import Any, cast + +import torch + +from trainer.configs.pipelines import PipelineConfig +from trainer.distributed import ( + maybe_init_distributed_environment_and_model_parallel) +from trainer.trainer_args import TrainerArgs, TrainingArgs +from trainer.logger import init_logger +from trainer.models.loader.component_loader import PipelineComponentLoader +from trainer.pipelines.pipeline_batch_info import ForwardBatch +from trainer.pipelines.base import PipelineStage +from trainer.utils import (maybe_download_model, + verify_model_config_and_directory) + +logger = init_logger(__name__) + + +class ComposedPipelineBase(ABC): + """ + Base class for pipelines composed of multiple stages. + + This class provides the framework for creating pipelines by composing multiple + stages together. Each stage is responsible for a specific part of the diffusion + process, and the pipeline orchestrates the execution of these stages. + """ + + is_video_pipeline: bool = False # To be overridden by video pipelines + _required_config_modules: list[str] = [] + _extra_config_module_map: dict[str, str] = {} + training_args: TrainingArgs | None = None + trainer_args: TrainerArgs | TrainingArgs | None = None + modules: dict[str, torch.nn.Module] = {} + post_init_called: bool = False + + # TODO(will): args should support both inference args and training args + def __init__(self, + model_path: str, + trainer_args: TrainerArgs | TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None): + """ + Initialize the pipeline. After __init__, the pipeline should be ready to + use. The pipeline should be stateless and not hold any batch state. + """ + self.trainer_args = trainer_args + + self.model_path: str = model_path + self._stages: list[PipelineStage] = [] + self._stage_name_mapping: dict[str, PipelineStage] = {} + + if required_config_modules is not None: + self._required_config_modules = required_config_modules + + if self._required_config_modules is None: + raise NotImplementedError( + "Subclass must set _required_config_modules") + + maybe_init_distributed_environment_and_model_parallel( + trainer_args.tp_size, trainer_args.sp_size) + + # Load modules directly in initialization + logger.info("Loading pipeline modules...") + # print(trainer_args, loaded_modules) + # quit() + if "HunyuanTransformer" in trainer_args.cls_name: + self.modules = self.load_hunyuan_modules(trainer_args) + else: + self.modules = self.load_modules(trainer_args, loaded_modules) + + def set_trainable(self) -> None: + # Only train DiT + if getattr(self.trainer_args, "training_mode", False): + for name, module in self.modules.items(): + if not isinstance(module, torch.nn.Module): + continue + if name == "transformer": + module.requires_grad_(True) + else: + module.requires_grad_(False) + + def post_init(self) -> None: + assert self.trainer_args is not None, "trainer_args must be set" + if self.post_init_called: + return + self.post_init_called = True + if self.trainer_args.training_mode: + assert isinstance(self.trainer_args, TrainingArgs) + self.training_args = self.trainer_args + assert self.training_args is not None + self.initialize_training_pipeline(self.training_args) + if self.training_args.log_validation: + self.initialize_validation_pipeline(self.training_args) + + self.initialize_pipeline(self.trainer_args) + if self.trainer_args.enable_torch_compile: + self.modules["transformer"] = torch.compile( + self.modules["transformer"]) + logger.info("Torch Compile enabled for DiT") + + if not self.trainer_args.training_mode: + logger.info("Creating pipeline stages...") + self.create_pipeline_stages(self.trainer_args) + + def initialize_training_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method") + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError( + "if log_validation is True, the pipeline must implement this method" + ) + + @classmethod + def from_pretrained(cls, + model_path: str, + device: str | None = None, + torch_dtype: torch.dtype | None = None, + pipeline_config: str | PipelineConfig | None = None, + args: argparse.Namespace | None = None, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] + | None = None, + **kwargs) -> "ComposedPipelineBase": + """ + Load a pipeline from a pretrained model. + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + if args is None or args.inference_mode: + + kwargs['model_path'] = model_path + trainer_args = TrainerArgs.from_kwargs(**kwargs) + else: + assert args is not None, "args must be provided for training mode" + trainer_args = TrainingArgs.from_cli_args(args) + # TODO(will): fix this so that its not so ugly + trainer_args.model_path = model_path + for key, value in kwargs.items(): + setattr(trainer_args, key, value) + + trainer_args.dit_cpu_offload = False + # we hijack the precision to be the master weight type so that the + # model is loaded with the correct precision. Subsequently we will + # use FSDP2's MixedPrecisionPolicy to set the precision for the + # fwd, bwd, and other operations' precision. + # assert trainer_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training' + + logger.info("trainer_args in from_pretrained: %s", trainer_args) + + pipe = cls(model_path, + trainer_args, + required_config_modules=required_config_modules, + loaded_modules=loaded_modules) + pipe.post_init() + return pipe + + def get_module(self, module_name: str, default_value: Any = None) -> Any: + if module_name not in self.modules: + return default_value + return self.modules[module_name] + + def add_module(self, module_name: str, module: Any): + self.modules[module_name] = module + + def _load_config(self, model_path: str) -> dict[str, Any]: + model_path = maybe_download_model(self.model_path) + self.model_path = model_path + # trainer_args.downloaded_model_path = model_path + logger.info("Model path: %s", model_path) + config = verify_model_config_and_directory(model_path) + return cast(dict[str, Any], config) + + @property + def required_config_modules(self) -> list[str]: + """ + List of modules that are required by the pipeline. The names should match + the diffusers directory and model_index.json file. These modules will be + loaded using the PipelineComponentLoader and made available in the + modules dictionary. Access these modules using the get_module method. + + class ConcretePipeline(ComposedPipelineBase): + _required_config_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"] + + + @property + def required_config_modules(self): + return self._required_config_modules + """ + return self._required_config_modules + + @property + def stages(self) -> list[PipelineStage]: + """ + List of stages in the pipeline. + """ + return self._stages + + @abstractmethod + def create_pipeline_stages(self, trainer_args: TrainerArgs): + """ + Create the inference pipeline stages. + """ + raise NotImplementedError + + def create_training_stages(self, training_args: TrainingArgs): + """ + Create the training pipeline stages. + """ + raise NotImplementedError + + def initialize_pipeline(self, trainer_args: TrainerArgs): + """ + Initialize the pipeline. + """ + return + + def load_hunyuan_modules( + self, + trainer_args: TrainerArgs, + ) -> dict[str, Any]: + """ + Load the modules from the config. + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + modules = {} + module_name = "transformer" + component_model_path = trainer_args.load_from_dir + trainer_args.module_name = module_name + + module = PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers="diffusers", + trainer_args=trainer_args, + ) + logger.info("Loaded module %s from %s", module_name, + component_model_path) + + if module_name in modules: + logger.warning("Overwriting module %s", module_name) + modules[module_name] = module + return modules + + def load_modules( + self, + trainer_args: TrainerArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None + ) -> dict[str, Any]: + """ + Load the modules from the config. + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + model_index = self._load_config(self.model_path) + logger.info("Loading pipeline modules from config: %s", model_index) + + # remove keys that are not pipeline modules + model_index.pop("_class_name") + model_index.pop("_diffusers_version") + # @TODO(Wei): Temporary hack + # wan2.2 14B才会用上 + if "boundary_ratio" in model_index and model_index[ + "boundary_ratio"] is not None: + logger.info( + "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." + ) + self.required_config_modules.append("transformer_2") + if trainer_args.boundary_ratio is None: + logger.info( + "MoE pipeline detected. Setting boundary ratio to %s", + model_index["boundary_ratio"]) + trainer_args.boundary_ratio = model_index["boundary_ratio"] + + model_index.pop("boundary_ratio", None) + model_index.pop("expand_timesteps", None) + + # some sanity checks + assert len( + model_index + ) > 1, "model_index.json must contain at least one pipeline module" + # print(model_index) + # quit() + for module_name in self.required_config_modules: + if module_name not in model_index and module_name in self._extra_config_module_map: + extra_module_value = self._extra_config_module_map[module_name] + logger.warning( + "model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.", + module_name, module_name, extra_module_value) + if extra_module_value in model_index: + logger.info("Using module %s for %s", extra_module_value, + module_name) + model_index[module_name] = model_index[extra_module_value] + continue + else: + raise ValueError( + f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}" + ) + + # all the component models used by the pipeline + required_modules = self.required_config_modules + logger.info("Loading required modules: %s", required_modules) + + modules = {} + for module_name, (transformers_or_diffusers, + architecture) in model_index.items(): + # print(transformers_or_diffusers) + if transformers_or_diffusers is None: + if module_name in self.required_config_modules: + self.required_config_modules.remove(module_name) + continue + if module_name not in required_modules: + logger.info("Skipping module %s", module_name) + continue + if loaded_modules is not None and module_name in loaded_modules: + logger.info("Using module %s already provided", module_name) + modules[module_name] = loaded_modules[module_name] + continue + + # we load the module from the extra config module map if it exists + if module_name in self._extra_config_module_map: + load_module_name = self._extra_config_module_map[module_name] + else: + load_module_name = module_name + + component_model_path = os.path.join(self.model_path, + load_module_name) + # 为了区分不同的transformer DiT + trainer_args.module_name = module_name + + module = PipelineComponentLoader.load_module( + module_name=load_module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + trainer_args=trainer_args, + ) + logger.info("Loaded module %s from %s", module_name, + component_model_path) + + if module_name in modules: + logger.warning("Overwriting module %s", module_name) + modules[module_name] = module + # quit() + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in modules or modules[module_name] is None: + raise ValueError( + f"Required module key: {module_name} value: {modules.get(module_name)} was not found in loaded modules {modules.keys()}" + ) + + return modules + + def add_stage(self, stage_name: str, stage: PipelineStage): + assert self.modules is not None, "No modules are registered" + self._stages.append(stage) + self._stage_name_mapping[stage_name] = stage + setattr(self, stage_name, stage) + + # TODO(will): don't hardcode no_grad + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + trainer_args: TrainerArgs, + ) -> ForwardBatch: + """ + Generate a video or image using the pipeline. + + Args: + batch: The batch to generate from. + trainer_args: The inference arguments. + Returns: + ForwardBatch: The batch with the generated video or image. + """ + if not self.post_init_called: + self.post_init() + + # Execute each stage + logger.info("Running pipeline stages: %s", + self._stage_name_mapping.keys()) + # logger.info("Batch: %s", batch) + for stage in self.stages: + batch = stage(batch, trainer_args) + + # Return the output + return batch + + def train(self) -> None: + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method") diff --git a/trainer/pipelines/lora_pipeline.py b/trainer/pipelines/lora_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..46b25c21b1751477f0245897a93ff33a5cd81f06 --- /dev/null +++ b/trainer/pipelines/lora_pipeline.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections import defaultdict +from collections.abc import Hashable +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors.torch import load_file +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor + +from trainer.distributed import get_local_torch_device +from trainer.trainer_args import TrainerArgs, TrainingArgs +from trainer.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer, + replace_submodule) +from trainer.logger import init_logger +from trainer.models.loader.utils import get_param_names_mapping +from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase +from trainer.utils import maybe_download_lora + +logger = init_logger(__name__) + + +class LoRAPipeline(ComposedPipelineBase): + """ + Pipeline that supports injecting LoRA adapters into the diffusion transformer. + TODO: support training. + """ + lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict( + dict) # state dicts of loaded lora adapters + cur_adapter_name: str = "" + cur_adapter_path: str = "" + lora_layers: dict[str, BaseLayerWithLoRA] = {} + trainer_args: TrainerArgs | TrainingArgs + exclude_lora_layers: list[str] = [] + device: torch.device = get_local_torch_device() + lora_target_modules: list[str] | None = None + lora_path: str | None = None + lora_nickname: str = "default" + lora_rank: int | None = None + lora_alpha: int | None = None + lora_initialized: bool = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.device = get_local_torch_device() + # self.exclude_lora_layers = self.modules[ + # "transformer"].config.arch_config.exclude_lora_layers + self.lora_target_modules = self.trainer_args.lora_target_modules + self.lora_path = self.trainer_args.lora_path + self.lora_nickname = self.trainer_args.lora_nickname + self.training_mode = self.trainer_args.training_mode + if self.training_mode and getattr(self.trainer_args, "lora_training", + False): + assert isinstance(self.trainer_args, TrainingArgs) + if self.trainer_args.lora_alpha is None: + self.trainer_args.lora_alpha = self.trainer_args.lora_rank + self.lora_rank = self.trainer_args.lora_rank # type: ignore + self.lora_alpha = self.trainer_args.lora_alpha # type: ignore + logger.info("Using LoRA training with rank %d and alpha %d", + self.lora_rank, self.lora_alpha) + if self.lora_target_modules is None: + self.lora_target_modules = [ + "q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k", + "to_v", "to_out", "to_qkv" + ] + self.convert_to_lora_layers() + # Inference + elif not self.training_mode and self.lora_path is not None: + self.convert_to_lora_layers() + self.set_lora_adapter( + self.lora_nickname, # type: ignore + self.lora_path) # type: ignore + + def is_target_layer(self, module_name: str) -> bool: + if self.lora_target_modules is None: + return True + return any(target_name in module_name + for target_name in self.lora_target_modules) + + def set_trainable(self) -> None: + + is_lora_training = self.training_mode and getattr( + self.trainer_args, "lora_training", False) + if not is_lora_training: + super().set_trainable() + return + + self.modules["transformer"].requires_grad_(False) + device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1), + mesh_dim_names=["fake", "replicate"]) + for name, layer in self.lora_layers.items(): + # Enable grads for lora weights only + # Must convert to DTensor for compatibility with other FSDP modules in grad calculation + layer.lora_A.requires_grad_(True) + layer.lora_B.requires_grad_(True) + layer.base_layer.requires_grad_(False) + layer.lora_A = nn.Parameter( + DTensor.from_local(layer.lora_A, device_mesh=device_mesh)) + layer.lora_B = nn.Parameter( + DTensor.from_local(layer.lora_B, device_mesh=device_mesh)) + + def convert_to_lora_layers(self) -> None: + """ + Unified method to convert the transformer to a LoRA transformer. + """ + if self.lora_initialized: + return + self.lora_initialized = True + converted_count = 0 + for name, layer in self.modules["transformer"].named_modules(): + if not self.is_target_layer(name): + continue + + excluded = False + for exclude_layer in self.exclude_lora_layers: + if exclude_layer in name: + excluded = True + break + if excluded: + continue + + layer = get_lora_layer(layer, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + training_mode=self.training_mode) + if layer is not None: + self.lora_layers[name] = layer + replace_submodule(self.modules["transformer"], name, layer) + converted_count += 1 + logger.info("Converted %d layers to LoRA layers", converted_count) + + def set_lora_adapter(self, + lora_nickname: str, + lora_path: str | None = None): # type: ignore + """ + Load a LoRA adapter into the pipeline and merge it into the transformer. + Args: + lora_nickname: The "nick name" of the adapter when referenced in the pipeline. + lora_path: The path to the adapter, either a local path or a Hugging Face repo id. + """ + + if lora_nickname not in self.lora_adapters and lora_path is None: + raise ValueError( + f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it." + ) + if not self.lora_initialized: + self.convert_to_lora_layers() + adapter_updated = False + rank = dist.get_rank() + if lora_path is not None and lora_path != self.cur_adapter_path: + lora_local_path = maybe_download_lora(lora_path) + lora_state_dict = load_file(lora_local_path) + + # Map the hf layer names to our custom layer names + param_names_mapping_fn = get_param_names_mapping( + self.modules["transformer"].param_names_mapping) + lora_param_names_mapping_fn = get_param_names_mapping( + self.modules["transformer"].lora_param_names_mapping) + + to_merge_params: defaultdict[Hashable, + dict[Any, Any]] = defaultdict(dict) + for name, weight in lora_state_dict.items(): + name = name.replace("diffusion_model.", "") + name = name.replace(".weight", "") + name, _, _ = lora_param_names_mapping_fn(name) + target_name, merge_index, num_params_to_merge = param_names_mapping_fn( + name) + # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse + # see param mapping in HunyuanVideoArchConfig + if merge_index is not None and "lora_B" in name: + to_merge_params[target_name][merge_index] = weight + if len(to_merge_params[target_name]) == num_params_to_merge: + # cat at output dim according to the merge_index order + sorted_tensors = [ + to_merge_params[target_name][i] + for i in range(num_params_to_merge) + ] + weight = torch.cat(sorted_tensors, dim=1) + del to_merge_params[target_name] + else: + continue + + if target_name in self.lora_adapters[lora_nickname]: + raise ValueError( + f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]" + ) + self.lora_adapters[lora_nickname][target_name] = weight.to( + self.device) + adapter_updated = True + self.cur_adapter_path = lora_path + logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) + + if not adapter_updated and self.cur_adapter_name == lora_nickname: + return + self.cur_adapter_name = lora_nickname + + # Merge the new adapter + adapted_count = 0 + for name, layer in self.lora_layers.items(): + lora_A_name = name + ".lora_A" + lora_B_name = name + ".lora_B" + if lora_A_name in self.lora_adapters[lora_nickname]\ + and lora_B_name in self.lora_adapters[lora_nickname]: + layer.set_lora_weights( + self.lora_adapters[lora_nickname][lora_A_name], + self.lora_adapters[lora_nickname][lora_B_name], + training_mode=self.trainer_args.training_mode, + lora_path=lora_path) + adapted_count += 1 + else: + if rank == 0: + logger.warning( + "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.", + lora_path, name) + layer.disable_lora = True + logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank, + lora_path, adapted_count) diff --git a/trainer/pipelines/pipeline_batch_info.py b/trainer/pipelines/pipeline_batch_info.py new file mode 100644 index 0000000000000000000000000000000000000000..0a9db2260fb3fe934c6081b217cab10a9ed3d946 --- /dev/null +++ b/trainer/pipelines/pipeline_batch_info.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py +""" +Data structures for functional pipeline processing. + +This module defines the dataclasses used to pass state between pipeline components +in a functional manner, reducing the need for explicit parameter passing. +""" + +import pprint +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Any + +import PIL.Image +import torch + +if TYPE_CHECKING: + from torchcodec.decoders import VideoDecoder + +import time +from collections import OrderedDict + +from trainer.attention import AttentionMetadata +from trainer.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams + + +class PipelineLoggingInfo: + """Simple approach using OrderedDict to track stage metrics.""" + + def __init__(self): + # OrderedDict preserves insertion order and allows easy access + self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict() + + def add_stage_execution_time(self, stage_name: str, execution_time: float): + """Add execution time for a stage.""" + if stage_name not in self.stages: + self.stages[stage_name] = {} + self.stages[stage_name]['execution_time'] = execution_time + self.stages[stage_name]['timestamp'] = time.time() + + def add_stage_metric(self, stage_name: str, metric_name: str, value: Any): + """Add any metric for a stage.""" + if stage_name not in self.stages: + self.stages[stage_name] = {} + self.stages[stage_name][metric_name] = value + + def get_stage_info(self, stage_name: str) -> dict[str, Any]: + """Get all info for a specific stage.""" + return self.stages.get(stage_name, {}) + + def get_execution_order(self) -> list[str]: + """Get stages in execution order.""" + return list(self.stages.keys()) + + def get_total_execution_time(self) -> float: + """Get total pipeline execution time.""" + return sum( + stage.get('execution_time', 0) for stage in self.stages.values()) + + +@dataclass +class ForwardBatch: + """ + Complete state passed through the pipeline execution. + + This dataclass contains all information needed during the diffusion pipeline + execution, allowing methods to update specific components without needing + to manage numerous individual parameters. + """ + # TODO(will): double check that args are separate from trainer_args + # properly. Also maybe think about providing an abstraction for pipeline + # specific arguments. + data_type: str + + generator: torch.Generator | list[torch.Generator] | None = None + + # Image inputs + image_path: str | None = None + image_embeds: list[torch.Tensor] = field(default_factory=list) + pil_image: torch.Tensor | PIL.Image.Image | None = None + preprocessed_image: torch.Tensor | None = None + + # Text inputs + prompt: str | list[str] | None = None + negative_prompt: str | list[str] | None = None + prompt_path: str | None = None + output_path: str = "outputs/" + output_video_name: str | None = None + # Primary encoder embeddings + prompt_embeds: list[torch.Tensor] = field(default_factory=list) + negative_prompt_embeds: list[torch.Tensor] | None = None + prompt_attention_mask: list[torch.Tensor] | None = None + negative_attention_mask: list[torch.Tensor] | None = None + clip_embedding_pos: list[torch.Tensor] | None = None + clip_embedding_neg: list[torch.Tensor] | None = None + + # Additional text-related parameters + max_sequence_length: int | None = None + prompt_template: dict[str, Any] | None = None + do_classifier_free_guidance: bool = False + + # Batch info + batch_size: int | None = None + num_videos_per_prompt: int = 1 + seed: int | None = None + seeds: list[int] | None = None + + # Tracking if embeddings are already processed + is_prompt_processed: bool = False + + # Latent tensors + latents: torch.Tensor | None = None + raw_latent_shape: torch.Tensor | None = None + noise_pred: torch.Tensor | None = None + image_latent: torch.Tensor | None = None + + # Latent dimensions + height_latents: list[int] | int | None = None + width_latents: list[int] | int | None = None + num_frames: list[int] | int = 1 # Default for image models + num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus + + # Original dimensions (before VAE scaling) + height: list[int] | int | None = None + width: list[int] | int | None = None + fps: list[int] | int | None = None + + # Timesteps + timesteps: torch.Tensor | None = None + timestep: torch.Tensor | float | int | None = None + step_index: int | None = None + + # Scheduler parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_scale_2: float | None = None + guidance_rescale: float = 0.0 + eta: float = 0.0 + sigmas: list[float] | None = None + + n_tokens: int | None = None + + # Other parameters that may be needed by specific schedulers + extra_step_kwargs: dict[str, Any] = field(default_factory=dict) + + # Component modules (populated by the pipeline) + modules: dict[str, Any] = field(default_factory=dict) + + # Final output (after pipeline completion) + output: Any = None + + # Extra parameters that might be needed by specific pipeline implementations + extra: dict[str, Any] = field(default_factory=dict) + + # Misc + save_video: bool = True + return_frames: bool = False + + # TeaCache parameters + enable_teacache: bool = False + teacache_params: TeaCacheParams | WanTeaCacheParams | None = None + + # STA parameters + STA_param: list | None = None + is_cfg_negative: bool = False + mask_search_final_result_pos: list[list] | None = None + mask_search_final_result_neg: list[list] | None = None + + # VSA parameters + VSA_sparsity: float = 0.0 + + # Logging info + logging_info: PipelineLoggingInfo = field( + default_factory=PipelineLoggingInfo) + + def __post_init__(self): + """Initialize dependent fields after dataclass initialization.""" + + # Set do_classifier_free_guidance based on guidance scale and negative prompt + if self.guidance_scale > 1.0: + self.do_classifier_free_guidance = True + if self.negative_prompt_embeds is None: + self.negative_prompt_embeds = [] + if self.guidance_scale_2 is None: + self.guidance_scale_2 = self.guidance_scale + + def __str__(self): + return pprint.pformat(asdict(self), indent=2, width=120) + + +@dataclass +class TrainingBatch: + current_timestep: int = 0 + current_vsa_sparsity: float = 0.0 + + # Dataloader batch outputs + latents: torch.Tensor | None = None + prompt_embed: torch.Tensor | None = None + neg_prompt_embed: torch.Tensor | None = None + w2c: torch.Tensor | None = None + intrinsic: torch.Tensor | None = None + action: torch.Tensor | None = None + perspective_embedding: torch.Tensor | None = None + norm_dirs: torch.Tensor | None = None + action_rot_trans: torch.Tensor | None = None + i2v_mask: torch.Tensor | None = None + window_frames: int = 0 + per_seq_length: int = 0 + current_start: int = 0 + current_end: int = 0 + video_path: str = None + stage_one: bool = False + gt_latent: torch.Tensor | None = None + select_window_out_flag: int = 0 + all_channel_latents: torch.Tensor | None = None + selection_frame_id_list = None + + raw_latent_shape: torch.Tensor | None = None + noise_latents: torch.Tensor | None = None + encoder_hidden_states: torch.Tensor | None = None + encoder_attention_mask: torch.Tensor | None = None + + image_cond: torch.Tensor | None = None + vision_states: torch.Tensor | None = None + prompt_mask: torch.Tensor | None = None + byt5_text_states: torch.Tensor | None = None + byt5_text_mask: torch.Tensor | None = None + # i2v + preprocessed_image: torch.Tensor | None = None + image_embeds: torch.Tensor | None = None + image_latents: torch.Tensor | None = None + infos: list[dict[str, Any]] | None = None + mask_lat_size: torch.Tensor | None = None + + # Transformer inputs + noisy_model_input: torch.Tensor | None = None + timesteps: torch.Tensor | None = None + sigmas: torch.Tensor | None = None + noise: torch.Tensor | None = None + + attn_metadata_vsa: AttentionMetadata | None = None + attn_metadata: AttentionMetadata | None = None + + # input kwargs + input_kwargs: dict[str, Any] | None = None + + # Training loss + loss: torch.Tensor | None = None + + # Training outputs + total_loss: float | None = None + grad_norm: float | None = None + generator_grad_norm: float | None = None + fake_grad_norm: float | None = None + + # Distillation-specific attributes + encoder_hidden_states_neg: torch.Tensor | None = None + encoder_attention_mask_neg: torch.Tensor | None = None + conditional_dict: dict[str, Any] | None = None + unconditional_dict: dict[str, Any] | None = None + + # Distillation losses + generator_loss: float = 0.0 + fake_score_loss: float = 0.0 + + dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PreprocessBatch(ForwardBatch): + video_loader: list["VideoDecoder"] = field(default_factory=list) + video_file_name: list[str] = field(default_factory=list) diff --git a/trainer/pipelines/pipeline_registry.py b/trainer/pipelines/pipeline_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..d96c329967342f623d32c23e60c7db98a5cef2ba --- /dev/null +++ b/trainer/pipelines/pipeline_registry.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py +# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py + +import importlib +import pkgutil +from collections.abc import Set +from dataclasses import dataclass, field +from enum import Enum +from functools import lru_cache + +from trainer.trainer_args import WorkloadType +from trainer.logger import init_logger +from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase +from trainer.pipelines.lora_pipeline import LoRAPipeline + +logger = init_logger(__name__) + +# map pipeline name to folder name +_PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { + "WanPipeline": "wan", + "WanDMDPipeline": "wan", + "WanImageToVideoPipeline": "wan", + "StepVideoPipeline": "stepvideo", + "HunyuanVideoPipeline": "hunyuan", +} + +_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { + WorkloadType.I2V: "PreprocessPipelineI2V", + WorkloadType.T2V: "PreprocessPipelineT2V", +} + + +class PipelineType(str, Enum): + """ + Enumeration for different pipeline types. + + Inherits from str to allow string comparison for backward compatibility. + """ + BASIC = "basic" + PREPROCESS = "preprocess" + TRAINING = "training" + + @classmethod + def from_string(cls, value: str) -> "PipelineType": + """Convert string to PipelineType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid pipeline type: {value}. Must be one of: {', '.join([t.value for t in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings.""" + return [pipeline_type.value for pipeline_type in cls] + + +@dataclass +class _PipelineRegistry: + # Keyed by pipeline_type -> architecture -> pipeline_name + # pipelines[pipeline_type][architecture][pipeline_name] = pipeline_cls + pipelines: dict[str, dict[str, dict[str, type[ComposedPipelineBase] + | None]]] = field(default_factory=dict) + + def get_supported_archs(self, pipeline_name_in_config: str, + pipeline_type: PipelineType) -> Set[str]: + """Get supported architectures, optionally filtered by pipeline type and workload type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + return set(self.pipelines[pipeline_type.value][arch].keys()) + + def _load_preprocess_pipeline_cls( + self, workload_type: WorkloadType, + arch: str) -> type[ComposedPipelineBase] | None: + pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[ + workload_type] + + return self.pipelines[ + PipelineType.PREPROCESS.value][arch][pipeline_name] + + def _try_load_pipeline_cls( + self, pipeline_name_in_config: str, pipeline_type: PipelineType, + workload_type: WorkloadType + ) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None: + """Try to load a pipeline class for the given architecture, pipeline type, and workload type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + + if (pipeline_type.value not in self.pipelines + or arch not in self.pipelines[pipeline_type.value]): + return None + + if pipeline_type == PipelineType.PREPROCESS: + return self._load_preprocess_pipeline_cls(workload_type, arch) + elif pipeline_type == PipelineType.BASIC: + return self.pipelines[ + pipeline_type.value][arch][pipeline_name_in_config] + elif pipeline_type == PipelineType.TRAINING: + pass + else: + raise ValueError(f"Invalid pipeline type: {pipeline_type.value}") + + return None + + def resolve_pipeline_cls( + self, + pipeline_name_in_config: str, + pipeline_type: PipelineType, + workload_type: WorkloadType, + ) -> type[ComposedPipelineBase] | type[LoRAPipeline]: + """Resolve pipeline class based on pipeline name in the config, pipeline type, and workload type.""" + if not pipeline_name_in_config: + logger.warning("No pipeline architecture is specified") + + pipeline_cls = self._try_load_pipeline_cls(pipeline_name_in_config, + pipeline_type, workload_type) + if pipeline_cls is not None: + return pipeline_cls + supported_archs = self.get_supported_archs(pipeline_name_in_config, + pipeline_type) + raise ValueError( + f"Pipeline architecture '{pipeline_name_in_config}' is not supported for pipeline type '{pipeline_type.value}' " + f"and workload type '{workload_type.value}'. " + f"Supported architectures: {supported_archs}") + + +@lru_cache +def import_pipeline_classes( + pipeline_types: list[PipelineType] | PipelineType | None = None +) -> dict[str, dict[str, dict[str, type[ComposedPipelineBase] | None]]]: + """ + Import pipeline classes based on the pipeline type and workload type. + + Args: + pipeline_types: The pipeline types to load (basic, preprocess, training). + If None, loads all types. + + Returns: + A three-level nested dictionary: + {pipeline_type: {architecture_name: {pipeline_name: pipeline_cls}}} + e.g., {"basic": {"wan": {"WanPipeline": WanPipeline}}} + """ + type_to_arch_to_pipeline_dict: dict[str, + dict[str, + dict[str, + type[ComposedPipelineBase] + | None]]] = {} + package_name: str = "trainer.pipelines" + + # Determine which pipeline types to scan + if isinstance(pipeline_types, list): + pipeline_types_to_scan = [ + pipeline_type.value for pipeline_type in pipeline_types + ] + elif isinstance(pipeline_types, PipelineType): + pipeline_types_to_scan = [pipeline_types.value] + else: + pipeline_types_to_scan = [pt.value for pt in PipelineType] + + logger.info("Loading pipelines for types: %s", pipeline_types_to_scan) + + for pipeline_type_str in pipeline_types_to_scan: + arch_to_pipeline_dict: dict[str, dict[str, type[ComposedPipelineBase] + | None]] = {} + + # Try to load from pipeline-type-specific directory first + pipeline_type_package_name = f"{package_name}.{pipeline_type_str}" + + try: + pipeline_type_package = importlib.import_module( + pipeline_type_package_name) + logger.debug("Successfully imported %s", pipeline_type_package_name) + + for _, arch, ispkg in pkgutil.iter_modules( + pipeline_type_package.__path__): + pipeline_dict: dict[str, type[ComposedPipelineBase] | None] = {} + + arch_package_name = f"{pipeline_type_package_name}.{arch}" + if ispkg: + arch_package = importlib.import_module(arch_package_name) + for _, module_name, ispkg in pkgutil.walk_packages( + arch_package.__path__, arch_package_name + "."): + if not ispkg: + pipeline_module = importlib.import_module( + module_name) + if hasattr(pipeline_module, "EntryClass"): + if isinstance(pipeline_module.EntryClass, list): + for pipeline in pipeline_module.EntryClass: + pipeline_name = pipeline.__name__ + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[pipeline_name] = pipeline + else: + pipeline_name = pipeline_module.EntryClass.__name__ + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[ + pipeline_name] = pipeline_module.EntryClass + + arch_to_pipeline_dict[arch] = pipeline_dict + + except ImportError as e: + raise ImportError( + f"Could not import {pipeline_type_package_name} when importing pipeline classes: {e}" + ) from None + + type_to_arch_to_pipeline_dict[pipeline_type_str] = arch_to_pipeline_dict + + # Log summary + total_pipelines = sum( + len(pipeline_dict) + for arch_to_pipeline_dict in type_to_arch_to_pipeline_dict.values() + for pipeline_dict in arch_to_pipeline_dict.values()) + logger.info("Loaded %d pipeline classes across %d types", total_pipelines, + len(pipeline_types_to_scan)) + + return type_to_arch_to_pipeline_dict + + +def get_pipeline_registry( + pipeline_type: PipelineType | str | None = None) -> _PipelineRegistry: + """ + Get a pipeline registry for the specified mode, pipeline type, and workload type. + + Args: + pipeline_type: Pipeline type to load. If None and mode is provided, will be derived from mode. + + Returns: + A pipeline registry instance. + """ + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_classes = import_pipeline_classes(pipeline_type) + return _PipelineRegistry(pipeline_classes) diff --git a/trainer/pipelines/training/__init__.py b/trainer/pipelines/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dba3e904bbacbe8eb6ca97515c767d43e9608002 --- /dev/null +++ b/trainer/pipelines/training/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Training pipelines for trainer.v1. + +This package contains pipelines for training diffusion models. +""" diff --git a/trainer/pipelines/validators.py b/trainer/pipelines/validators.py new file mode 100644 index 0000000000000000000000000000000000000000..dac57f001bb4b72488e604c8852ba2dff7343a61 --- /dev/null +++ b/trainer/pipelines/validators.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Common validators for pipeline stage verification. + +This module provides reusable validation functions that can be used across +all pipeline stages for input/output verification. +""" + +from collections.abc import Callable +from typing import Any + +import torch + + +class StageValidators: + """Common validators for pipeline stages.""" + + @staticmethod + def not_none(value: Any) -> bool: + """Check if value is not None.""" + return value is not None + + @staticmethod + def positive_int(value: Any) -> bool: + """Check if value is a positive integer.""" + return isinstance(value, int) and value > 0 + + @staticmethod + def positive_float(value: Any) -> bool: + """Check if value is a positive float.""" + return isinstance(value, int | float) and value > 0 + + @staticmethod + def non_negative_float(value: Any) -> bool: + """Check if value is a non-negative float.""" + return isinstance(value, int | float) and value >= 0 + + @staticmethod + def divisible_by(value: Any, divisor: int) -> bool: + """Check if value is divisible by divisor.""" + return value is not None and isinstance(value, + int) and value % divisor == 0 + + @staticmethod + def is_tensor(value: Any) -> bool: + """Check if value is a torch tensor and doesn't contain NaN values.""" + if not isinstance(value, torch.Tensor): + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_with_dims(value: Any, dims: int) -> bool: + """Check if value is a tensor with specific dimensions and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if value.dim() != dims: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_min_dims(value: Any, min_dims: int) -> bool: + """Check if value is a tensor with at least min_dims dimensions and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if value.dim() < min_dims: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_shape_matches(value: Any, expected_shape: tuple) -> bool: + """Check if tensor shape matches expected shape (None for any size) and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if len(value.shape) != len(expected_shape): + return False + for actual, expected in zip(value.shape, expected_shape, strict=True): + if expected is not None and actual != expected: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def list_not_empty(value: Any) -> bool: + """Check if value is a non-empty list.""" + return isinstance(value, list) and len(value) > 0 + + @staticmethod + def list_length(value: Any, length: int) -> bool: + """Check if list has specific length.""" + return isinstance(value, list) and len(value) == length + + @staticmethod + def list_min_length(value: Any, min_length: int) -> bool: + """Check if list has at least min_length items.""" + return isinstance(value, list) and len(value) >= min_length + + @staticmethod + def string_not_empty(value: Any) -> bool: + """Check if value is a non-empty string.""" + return isinstance(value, str) and len(value.strip()) > 0 + + @staticmethod + def string_or_list_strings(value: Any) -> bool: + """Check if value is a string or list of strings.""" + if isinstance(value, str): + return True + if isinstance(value, list): + return all(isinstance(item, str) for item in value) + return False + + @staticmethod + def bool_value(value: Any) -> bool: + """Check if value is a boolean.""" + return isinstance(value, bool) + + @staticmethod + def generator_or_list_generators(value: Any) -> bool: + """Check if value is a Generator or list of Generators.""" + if isinstance(value, torch.Generator): + return True + if isinstance(value, list): + return all(isinstance(item, torch.Generator) for item in value) + return False + + @staticmethod + def is_list(value: Any) -> bool: + """Check if value is a list (can be empty).""" + return isinstance(value, list) + + @staticmethod + def is_tuple(value: Any) -> bool: + """Check if value is a tuple.""" + return isinstance(value, tuple) + + @staticmethod + def none_or_tensor(value: Any) -> bool: + """Check if value is None or a tensor without NaN values.""" + if value is None: + return True + if not isinstance(value, torch.Tensor): + return False + return not torch.isnan(value).any().item() + + @staticmethod + def list_of_tensors_with_dims(value: Any, dims: int) -> bool: + """Check if value is a non-empty list where all items are tensors with specific dimensions and no NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if item.dim() != dims: + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def list_of_tensors(value: Any) -> bool: + """Check if value is a non-empty list where all items are tensors without NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool: + """Check if value is a non-empty list where all items are tensors with at least min_dims dimensions and no NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if item.dim() < min_dims: + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is None or a tensor with specific dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + if value is None: + return True + if not isinstance(value, torch.Tensor): + return False + if value.dim() != dims: + return False + return not torch.isnan(value).any().item() + + return validator + + @staticmethod + def none_or_list(value: Any) -> bool: + """Check if value is None or a list.""" + return value is None or isinstance(value, list) + + @staticmethod + def none_or_positive_int(value: Any) -> bool: + """Check if value is None or a positive integer.""" + return value is None or (isinstance(value, int) and value > 0) + + # Helper methods that return functions for common patterns + @staticmethod + def with_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if tensor has specific dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.tensor_with_dims(value, dims) + + return validator + + @staticmethod + def min_dims(min_dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if tensor has at least min_dims dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.tensor_min_dims(value, min_dims) + + return validator + + @staticmethod + def divisible(divisor: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is divisible by divisor.""" + + def validator(value: Any) -> bool: + return StageValidators.divisible_by(value, divisor) + + return validator + + @staticmethod + def positive_int_divisible(divisor: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a positive integer divisible by divisor.""" + + def validator(value: Any) -> bool: + return (isinstance(value, int) and value > 0 + and StageValidators.divisible_by(value, divisor)) + + return validator + + @staticmethod + def list_of_tensors_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a list of tensors with specific dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.list_of_tensors_with_dims(value, dims) + + return validator + + @staticmethod + def list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a list of tensors with at least min_dims dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.list_of_tensors_with_min_dims( + value, min_dims) + + return validator + + +class ValidationFailure: + """Details about a specific validation failure.""" + + def __init__(self, + validator_name: str, + actual_value: Any, + expected: str | None = None, + error_msg: str | None = None): + self.validator_name = validator_name + self.actual_value = actual_value + self.expected = expected + self.error_msg = error_msg + + def __str__(self) -> str: + parts = [f"Validator '{self.validator_name}' failed"] + + if self.error_msg: + parts.append(f"Error: {self.error_msg}") + + # Add actual value info (but limit very long representations) + actual_str = self._format_value(self.actual_value) + parts.append(f"Actual: {actual_str}") + + if self.expected: + parts.append(f"Expected: {self.expected}") + + return ". ".join(parts) + + def _format_value(self, value: Any) -> str: + """Format a value for display in error messages.""" + if value is None: + return "None" + elif isinstance(value, torch.Tensor): + return f"tensor(shape={list(value.shape)}, dtype={value.dtype})" + elif isinstance(value, list): + if len(value) == 0: + return "[]" + elif len(value) <= 3: + item_strs = [self._format_value(item) for item in value] + return f"[{', '.join(item_strs)}]" + else: + return f"list(length={len(value)}, first_item={self._format_value(value[0])})" + elif isinstance(value, str): + if len(value) > 50: + return f"'{value[:47]}...'" + else: + return f"'{value}'" + else: + return f"{type(value).__name__}({value})" + + +class VerificationResult: + """Wrapper class for stage verification results.""" + + def __init__(self) -> None: + self._checks: dict[str, bool] = {} + self._failures: dict[str, list[ValidationFailure]] = {} + + def add_check( + self, field_name: str, value: Any, + validators: Callable[[Any], bool] | list[Callable[[Any], bool]] + ) -> 'VerificationResult': + """ + Add a validation check for a field. + + Args: + field_name: Name of the field being checked + value: The actual value to validate + validators: Single validation function or list of validation functions. + Each function will be called with the value as its first argument. + + Returns: + Self for method chaining + + Examples: + # Single validator + result.add_check("tensor", my_tensor, V.is_tensor) + + # Multiple validators (all must pass) + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + + # Using partial functions for parameters + result.add_check("height", batch.height, [V.not_none, V.divisible(8)]) + """ + if not isinstance(validators, list): + validators = [validators] + + failures = [] + all_passed = True + + # Apply all validators and collect detailed failure info + for validator in validators: + try: + passed = validator(value) + if not passed: + all_passed = False + failure = self._create_validation_failure(validator, value) + failures.append(failure) + except Exception as e: + # If any validator raises an exception, consider the check failed + all_passed = False + validator_name = getattr(validator, '__name__', str(validator)) + failure = ValidationFailure( + validator_name=validator_name, + actual_value=value, + error_msg=f"Exception during validation: {str(e)}") + failures.append(failure) + + self._checks[field_name] = all_passed + if not all_passed: + self._failures[field_name] = failures + + return self + + def _create_validation_failure(self, validator: Callable, + value: Any) -> ValidationFailure: + """Create a ValidationFailure with detailed information.""" + validator_name = getattr(validator, '__name__', str(validator)) + + # Try to extract meaningful expected value info based on validator type + expected = None + error_msg = None + + # Handle common validator patterns + if hasattr(validator, '__closure__') and validator.__closure__: + # This is likely a closure (like our helper functions) + if 'dims' in validator_name or 'with_dims' in str(validator): + if isinstance(value, torch.Tensor): + expected = f"tensor with {validator.__closure__[0].cell_contents} dimensions" + else: + expected = "tensor with specific dimensions" + elif 'divisible' in str(validator): + expected = f"integer divisible by {validator.__closure__[0].cell_contents}" + + # Handle specific validator types and check for NaN values + if validator_name == 'is_tensor': + expected = "torch.Tensor without NaN values" + if isinstance(value, + torch.Tensor) and torch.isnan(value).any().item(): + error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values" + elif validator_name == 'positive_int': + expected = "positive integer" + elif validator_name == 'not_none': + expected = "non-None value" + elif validator_name == 'list_not_empty': + expected = "non-empty list" + elif validator_name == 'bool_value': + expected = "boolean value" + elif 'tensor_with_dims' in validator_name or 'tensor_min_dims' in validator_name: + if isinstance(value, torch.Tensor): + if torch.isnan(value).any().item(): + error_msg = f"tensor has {value.dim()} dimensions but contains {torch.isnan(value).sum().item()} NaN values" + else: + error_msg = f"tensor has {value.dim()} dimensions" + elif validator_name == 'is_list': + expected = "list" + elif validator_name == 'none_or_tensor': + expected = "None or tensor without NaN values" + if isinstance(value, + torch.Tensor) and torch.isnan(value).any().item(): + error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values" + elif validator_name == 'list_of_tensors': + expected = "non-empty list of tensors without NaN values" + if isinstance(value, list) and len(value) > 0: + nan_count = 0 + for item in value: + if isinstance( + item, + torch.Tensor) and torch.isnan(item).any().item(): + nan_count += torch.isnan(item).sum().item() + if nan_count > 0: + error_msg = f"list contains tensors with total {nan_count} NaN values" + elif 'list_of_tensors_with_dims' in validator_name: + expected = "non-empty list of tensors with specific dimensions and no NaN values" + if isinstance(value, list) and len(value) > 0: + nan_count = 0 + for item in value: + if isinstance( + item, + torch.Tensor) and torch.isnan(item).any().item(): + nan_count += torch.isnan(item).sum().item() + if nan_count > 0: + error_msg = f"list contains tensors with total {nan_count} NaN values" + + return ValidationFailure(validator_name=validator_name, + actual_value=value, + expected=expected, + error_msg=error_msg) + + def is_valid(self) -> bool: + """Check if all validations passed.""" + return all(self._checks.values()) + + def get_failed_fields(self) -> list[str]: + """Get list of fields that failed validation.""" + return [field for field, passed in self._checks.items() if not passed] + + def get_detailed_failures(self) -> dict[str, list[ValidationFailure]]: + """Get detailed failure information for each failed field.""" + return self._failures.copy() + + def get_failure_summary(self) -> str: + """Get a comprehensive summary of all validation failures.""" + if self.is_valid(): + return "All validations passed" + + summary_parts = [] + for field_name, failures in self._failures.items(): + field_summary = f"\n Field '{field_name}':" + for i, failure in enumerate(failures, 1): + field_summary += f"\n {i}. {failure}" + summary_parts.append(field_summary) + + return "Validation failures:" + "".join(summary_parts) + + def to_dict(self) -> dict: + """Convert to dictionary for backward compatibility.""" + return self._checks.copy() + + +# Alias for convenience +V = StageValidators diff --git a/trainer/platforms/__init__.py b/trainer/platforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a528f24b1fa1e276ee145d7eb173f44d9df102de --- /dev/null +++ b/trainer/platforms/__init__.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/__init__.py + +import traceback +from typing import TYPE_CHECKING + +from trainer.logger import init_logger +# imported by other files, do not remove +from trainer.platforms.interface import AttentionBackendEnum # noqa: F401 +from trainer.platforms.interface import Platform, PlatformEnum +from trainer.utils import resolve_obj_by_qualname + +logger = init_logger(__name__) + + +def cuda_platform_plugin() -> str | None: + is_cuda = False + + try: + from trainer.utils import import_pynvml + pynvml = import_pynvml() # type: ignore[no-untyped-call] + pynvml.nvmlInit() + try: + # NOTE: Edge case: trainer cpu build on a GPU machine. + # Third-party pynvml can be imported in cpu build, + # we need to check if trainer is built with cpu too. + # Otherwise, trainer will always activate cuda plugin + # on a GPU machine, even if in a cpu build. + is_cuda = (pynvml.nvmlDeviceGetCount() > 0) + finally: + pynvml.nvmlShutdown() + except Exception as e: + if "nvml" not in e.__class__.__name__.lower(): + # If the error is not related to NVML, re-raise it. + raise e + + # CUDA is supported on Jetson, but NVML may not be. + import os + + def cuda_is_jetson() -> bool: + return os.path.isfile("/etc/nv_tegra_release") \ + or os.path.exists("/sys/class/tegra-firmware") + + if cuda_is_jetson(): + is_cuda = True + if is_cuda: + logger.info("CUDA is available") + + return "trainer.platforms.cuda.CudaPlatform" if is_cuda else None + + +def mps_platform_plugin() -> str | None: + """Detect if MPS (Metal Performance Shaders) is available on macOS.""" + is_mps = False + + try: + import torch + if torch.backends.mps.is_available(): + is_mps = True + logger.info("MPS (Metal Performance Shaders) is available") + except Exception as e: + logger.info("MPS detection failed: %s", e) + + return "trainer.platforms.mps.MpsPlatform" if is_mps else None + + +def cpu_platform_plugin() -> str | None: + """Detect if CPU platform should be used.""" + # CPU is always available as a fallback + return "trainer.platforms.cpu.CpuPlatform" + + +def rocm_platform_plugin() -> str | None: + is_rocm = False + + try: + import amdsmi + amdsmi.amdsmi_init() + try: + if len(amdsmi.amdsmi_get_processor_handles()) > 0: + is_rocm = True + logger.info("ROCm platform is available") + finally: + amdsmi.amdsmi_shut_down() + except Exception as e: + logger.info("ROCm platform is unavailable: %s", e) + + return "trainer.platforms.rocm.RocmPlatform" if is_rocm else None + + +builtin_platform_plugins = { + 'cuda': cuda_platform_plugin, + 'rocm': rocm_platform_plugin, + 'mps': mps_platform_plugin, + 'cpu': cpu_platform_plugin, +} + + +def resolve_current_platform_cls_qualname() -> str: + # TODO(will): if we need to support other platforms, we should consider if + # vLLM's plugin architecture is suitable for our needs. + + # Try MPS first on macOS + platform_cls_qualname = mps_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to ROCm + platform_cls_qualname = rocm_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CUDA + platform_cls_qualname = cuda_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CPU as last resort + platform_cls_qualname = cpu_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + raise RuntimeError("No platform plugin found. Please check your " + "installation.") + + +_current_platform = None +_init_trace: str = '' + +if TYPE_CHECKING: + current_platform: Platform + + +def __getattr__(name: str): + if name == 'current_platform': + # lazy init current_platform. + # 1. out-of-tree platform plugins need `from trainer.platforms import + # Platform` so that they can inherit `Platform` class. Therefore, + # we cannot resolve `current_platform` during the import of + # `trainer.platforms`. + # 2. when users use out-of-tree platform plugins, they might run + # `import trainer`, some trainer internal code might access + # `current_platform` during the import, and we need to make sure + # `current_platform` is only resolved after the plugins are loaded + # (we have tests for this, if any developer violate this, they will + # see the test failures). + global _current_platform + if _current_platform is None: + platform_cls_qualname = resolve_current_platform_cls_qualname() + _current_platform = resolve_obj_by_qualname(platform_cls_qualname)() + global _init_trace + _init_trace = "".join(traceback.format_stack()) + return _current_platform + elif name in globals(): + return globals()[name] + else: + raise AttributeError( + f"No attribute named '{name}' exists in {__name__}.") + + +__all__ = ['Platform', 'PlatformEnum', 'current_platform', "_init_trace"] diff --git a/trainer/platforms/cpu.py b/trainer/platforms/cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..563c7dbd58e8f5d842121ac659a15273639f74f8 --- /dev/null +++ b/trainer/platforms/cpu.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py + +import platform + +import torch + +from trainer.platforms.interface import CpuArchEnum, Platform, PlatformEnum + + +class CpuPlatform(Platform): + _enum = PlatformEnum.CPU + device_name = "CPU" + device_type = "cpu" + dispatch_key = "CPU" + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture.""" + machine = platform.machine().lower() + if machine in ("x86_64", "amd64", "i386", "i686"): + return CpuArchEnum.X86 + elif machine in ("arm64", "aarch64"): + return CpuArchEnum.ARM + else: + return CpuArchEnum.UNSPECIFIED + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return platform.processor() + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + return platform.machine() + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + # This is a rough estimate for CPU memory + # In practice, you might want to use psutil or similar + return 0 + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + return True + + @classmethod + def get_current_memory_usage(cls, + device: torch.types.Device | None = None + ) -> float: + # For CPU, we can't easily get memory usage without additional libraries + return 0.0 + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "trainer.distributed.device_communicators.cpu_communicator.CpuCommunicator" diff --git a/trainer/platforms/cuda.py b/trainer/platforms/cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..709d036ca9fcaa039e94c6cdce1d9b1790906ad8 --- /dev/null +++ b/trainer/platforms/cuda.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py +"""Code inside this file can safely assume cuda platform, e.g. importing +pynvml. However, it should not initialize cuda context. +""" + +import os +from collections.abc import Callable +from functools import lru_cache, wraps +from typing import TypeVar + +import torch +from typing_extensions import ParamSpec + +import trainer.envs as envs +from trainer.logger import init_logger +from trainer.platforms.interface import (AttentionBackendEnum, + DeviceCapability, Platform, + PlatformEnum) +from trainer.utils import import_pynvml + +logger = init_logger(__name__) + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +pynvml = import_pynvml() # type: ignore[no-untyped-call] + +# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models +# see https://github.com/huggingface/diffusers/issues/9704 for details +torch.backends.cuda.enable_cudnn_sdp(False) + + +def device_id_to_physical_device_id(device_id: int) -> int: + if "CUDA_VISIBLE_DEVICES" in os.environ: + device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",") + if device_ids == [""]: + msg = ( + "CUDA_VISIBLE_DEVICES is set to empty string, which means" + " GPU support is disabled. If you are using ray, please unset" + " the environment variable `CUDA_VISIBLE_DEVICES` inside the" + " worker/actor. " + "Check https://github.com/vllm-project/vllm/issues/8402 for" + " more information.") + raise RuntimeError(msg) + physical_device_id = device_ids[device_id] + return int(physical_device_id) + else: + return device_id + + +def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: + + @wraps(fn) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + pynvml.nvmlInit() + try: + return fn(*args, **kwargs) + finally: + pynvml.nvmlShutdown() + + return wrapper + + +class CudaPlatformBase(Platform): + _enum = PlatformEnum.CUDA + device_name: str = "cuda" + device_type: str = "cuda" + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, + device_id: int = 0) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used") + return False + return True + + @classmethod + def is_full_nvlink(cls, device_ids: list[int]) -> bool: + raise NotImplementedError + + @classmethod + def log_warnings(cls) -> None: + pass + + @classmethod + def get_current_memory_usage(cls, + device: torch.types.Device | None = None + ) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, + head_size: int, dtype: torch.dtype) -> str: + # TODO(will): maybe come up with a more general interface for local attention + # if distributed is False, we always try to use Flash attn + + logger.info("Trying TRAINER_ATTENTION_BACKEND=%s", + envs.TRAINER_ATTENTION_BACKEND) + if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: + try: + from st_attn import sliding_tile_attention # noqa: F401 + + from trainer.attention.backends.sliding_tile_attn import ( # noqa: F401 + SlidingTileAttentionBackend) + logger.info("Using Sliding Tile Attention backend.") + + return "trainer.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Sliding Tile Attention backend: %s", + str(e)) + raise ImportError( + "Sliding Tile Attention backend is not installed. ") from e + elif selected_backend == AttentionBackendEnum.SAGE_ATTN: + try: + from sageattention import sageattn # noqa: F401 + + from trainer.attention.backends.sage_attn import ( # noqa: F401 + SageAttentionBackend) + logger.info("Using Sage Attention backend.") + + return "trainer.attention.backends.sage_attn.SageAttentionBackend" + except ImportError as e: + logger.info(e) + logger.info( + "Sage Attention backend is not installed. Fall back to Flash Attention." + ) + elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + try: + from vsa import block_sparse_attn # noqa: F401 + + from trainer.attention.backends.video_sparse_attn import ( # noqa: F401 + VideoSparseAttentionBackend) + logger.info("Using Video Sparse Attention backend.") + + return "trainer.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Video Sparse Attention backend: %s", + str(e)) + raise ImportError( + "Video Sparse Attention backend is not installed. ") from e + elif selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "trainer.attention.backends.sdpa.SDPABackend" + elif selected_backend == AttentionBackendEnum.FLASH_ATTN or selected_backend is None: + pass + elif selected_backend: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + + target_backend = AttentionBackendEnum.FLASH_ATTN + if not cls.has_device_capability(80): + logger.info( + "Cannot use FlashAttention-2 backend for Volta and Turing " + "GPUs.") + target_backend = AttentionBackendEnum.TORCH_SDPA + elif dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention-2 backend for dtype other than " + "torch.float16 or torch.bfloat16.") + target_backend = AttentionBackendEnum.TORCH_SDPA + + # FlashAttn is valid for the model, checking if the package is + # installed. + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from trainer.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend) + + supported_sizes = \ + FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info("Cannot use FlashAttention-2 backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default).") + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "trainer.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "trainer.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "trainer.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + + +# NVML utils +# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, +# all the related functions work on real physical device ids. +# the major benefit of using NVML is that it will not initialize CUDA +class NvmlCudaPlatform(CudaPlatformBase): + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_capability(cls, + device_id: int = 0) -> DeviceCapability | None: + try: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) + return DeviceCapability(major=major, minor=minor) + except RuntimeError: + return None + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + try: + return bool(super().has_device_capability(capability, device_id)) + except RuntimeError: + return False + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_name(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + return cls._get_physical_device_name(physical_device_id) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_uuid(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return str(pynvml.nvmlDeviceGetUUID(handle)) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_total_memory(cls, device_id: int = 0) -> int: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + + @classmethod + @with_nvml_context + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + """ + query if the set of gpus are fully connected by nvlink (1 hop) + """ + handles = [ + pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids + ] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i < j: + try: + p2p_status = pynvml.nvmlDeviceGetP2PStatus( + handle, + peer_handle, + pynvml.NVML_P2P_CAPS_INDEX_NVLINK, + ) + if p2p_status != pynvml.NVML_P2P_STATUS_OK: + return False + except pynvml.NVMLError: + logger.exception( + "NVLink detection failed. This is normal if" + " your machine has no NVLink equipped.") + return False + return True + + @classmethod + def _get_physical_device_name(cls, device_id: int = 0) -> str: + handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) + return str(pynvml.nvmlDeviceGetName(handle)) + + @classmethod + @with_nvml_context + def log_warnings(cls) -> None: + device_ids: int = pynvml.nvmlDeviceGetCount() + if device_ids > 1: + device_names = [ + cls._get_physical_device_name(i) for i in range(device_ids) + ] + if (len(set(device_names)) > 1 + and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"): + logger.warning( + "Detected different devices in the system: %s. Please" + " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to " + "avoid unexpected behavior.", + ", ".join(device_names), + ) + + +class NonNvmlCudaPlatform(CudaPlatformBase): + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + device_props = torch.cuda.get_device_properties(device_id) + return int(device_props.total_memory) + + @classmethod + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + logger.exception("NVLink detection not possible, as context support was" + " not found. Assuming no NVLink available.") + return False + + +# Autodetect either NVML-enabled or non-NVML platform +# based on whether NVML is available. +nvml_available = False +try: + try: + pynvml.nvmlInit() + nvml_available = True + except Exception: + # On Jetson, NVML is not supported. + nvml_available = False +finally: + if nvml_available: + pynvml.nvmlShutdown() + +CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform + +try: + from sphinx.ext.autodoc.mock import _MockModule + + if not isinstance(pynvml, _MockModule): + CudaPlatform.log_warnings() +except ModuleNotFoundError: + CudaPlatform.log_warnings() diff --git a/trainer/platforms/interface.py b/trainer/platforms/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..5c08e6aff262ad185620c8d406ca382e2b1f9f3b --- /dev/null +++ b/trainer/platforms/interface.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py + +import enum +import random +from typing import NamedTuple + +import numpy as np +import torch + +from trainer.logger import init_logger + +logger = init_logger(__name__) + + +class AttentionBackendEnum(enum.Enum): + FLASH_ATTN = enum.auto() + SLIDING_TILE_ATTN = enum.auto() + TORCH_SDPA = enum.auto() + SAGE_ATTN = enum.auto() + VIDEO_SPARSE_ATTN = enum.auto() + NO_ATTENTION = enum.auto() + + +class PlatformEnum(enum.Enum): + CUDA = enum.auto() + ROCM = enum.auto() + TPU = enum.auto() + CPU = enum.auto() + MPS = enum.auto() + OOT = enum.auto() + UNSPECIFIED = enum.auto() + + +class CpuArchEnum(enum.Enum): + X86 = enum.auto() + ARM = enum.auto() + UNSPECIFIED = enum.auto() + + +class DeviceCapability(NamedTuple): + major: int + minor: int + + def as_version_str(self) -> str: + return f"{self.major}.{self.minor}" + + def to_int(self) -> int: + """ + Express device capability as an integer ````. + + It is assumed that the minor version is always a single digit. + """ + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class Platform: + _enum: PlatformEnum + device_name: str + device_type: str + + # available dispatch keys: + # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa + # use "CPU" as a fallback for platforms not registered in PyTorch + dispatch_key: str = "CPU" + + # The torch.compile backend for compiling simple and + # standalone functions. The default value is "inductor" to keep + # the same behavior as PyTorch. + # NOTE: for the forward part of the model, vLLM has another separate + # compilation strategy. + simple_compile_backend: str = "inductor" + + supported_quantization: list[str] = [] + + def is_cuda(self) -> bool: + return self._enum == PlatformEnum.CUDA + + def is_rocm(self) -> bool: + return self._enum == PlatformEnum.ROCM + + def is_tpu(self) -> bool: + return self._enum == PlatformEnum.TPU + + def is_cpu(self) -> bool: + return self._enum == PlatformEnum.CPU + + def is_out_of_tree(self) -> bool: + return self._enum == PlatformEnum.OOT + + def is_cuda_alike(self) -> bool: + """Stateless version of :func:`torch.cuda.is_available`.""" + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + + def is_mps(self) -> bool: + return self._enum == PlatformEnum.MPS + + @classmethod + def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, + head_size: int, dtype: torch.dtype) -> str: + """Get the attention backend class of a device.""" + return "" + + @classmethod + def get_device_capability( + cls, + device_id: int = 0, + ) -> DeviceCapability | None: + """Stateless version of :func:`torch.cuda.get_device_capability`.""" + return None + + @classmethod + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + """ + Test whether this platform is compatible with a device capability. + + The ``capability`` argument can either be: + + - A tuple ``(major, minor)``. + - An integer ````. (See :meth:`DeviceCapability.to_int`) + """ + current_capability = cls.get_device_capability(device_id=device_id) + if current_capability is None: + return False + + if isinstance(capability, tuple): + return current_capability >= capability + + return current_capability.to_int() >= capability + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + """Get the name of a device.""" + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + """Get the uuid of a device, e.g. the PCI bus ID.""" + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + """Get the total memory of a device in bytes.""" + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + """ + Check if the current platform supports async output. + """ + raise NotImplementedError + + @classmethod + def inference_mode(cls): + """A device-specific wrapper of `torch.inference_mode`. + + This wrapper is recommended because some hardware backends such as TPU + do not support `torch.inference_mode`. In such a case, they will fall + back to `torch.no_grad` by overriding this method. + """ + return torch.inference_mode(mode=True) + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """ + Set the seed of each random module. + `torch.manual_seed` will set seed on all devices. + + Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 + """ + if seed is not None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + @classmethod + def verify_model_arch(cls, model_arch: str) -> None: + """ + Verify whether the current platform supports the specified model + architecture. + + - This will raise an Error or Warning based on the model support on + the current platform. + - By default all models are considered supported. + """ + pass + + @classmethod + def verify_quantization(cls, quant: str) -> None: + """ + Verify whether the quantization is supported by the current platform. + """ + if cls.supported_quantization and \ + quant not in cls.supported_quantization: + raise ValueError( + f"{quant} quantization is currently not supported in " + f"{cls.device_name}.") + + @classmethod + def get_current_memory_usage(cls, + device: torch.types.Device | None = None + ) -> float: + """ + Return the memory usage in bytes. + """ + raise NotImplementedError + + @classmethod + def get_device_communicator_cls(cls) -> str: + """ + Get device specific communicator class for distributed communication. + """ + return "trainer.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture of the current platform.""" + return CpuArchEnum.UNSPECIFIED + + +class UnspecifiedPlatform(Platform): + _enum = PlatformEnum.UNSPECIFIED + device_type = "" diff --git a/trainer/platforms/mps.py b/trainer/platforms/mps.py new file mode 100644 index 0000000000000000000000000000000000000000..35c3fcd95bf6a0061fbca3d920a37a6e932a40bc --- /dev/null +++ b/trainer/platforms/mps.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from trainer.logger import init_logger +from trainer.platforms import AttentionBackendEnum +from trainer.platforms.interface import (DeviceCapability, Platform, + PlatformEnum) + +logger = init_logger(__name__) + + +class MpsPlatform(Platform): + _enum = PlatformEnum.MPS + device_name: str = "mps" + device_type: str = "mps" + dispatch_key: str = "MPS" + device_control_env_var: str = "MPS_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, + device_id: int = 0) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable MPS " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used") + return False + return True + + @classmethod + def get_current_memory_usage(cls, + device: torch.types.Device | None = None + ) -> float: + return 0.0 + + @classmethod + def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, + head_size: int, dtype: torch.dtype) -> str: + # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible + logger.info("Using Torch SDPA backend for MPS.") + return "trainer.attention.backends.sdpa.SDPABackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + # Use base communicator for MPS + return "trainer.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """Set the seed for MPS device.""" + if seed is not None: + import random + + import numpy as np + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + # MPS doesn't have manual_seed_all like CUDA + # The manual_seed above should be sufficient diff --git a/trainer/platforms/rocm.py b/trainer/platforms/rocm.py new file mode 100644 index 0000000000000000000000000000000000000000..4ad279cbedc88731a919359767eb889a38ad589b --- /dev/null +++ b/trainer/platforms/rocm.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py +""" +This file is a platform abstraction for ROCm GPUs, +adjusted to match the structure and interface of `cuda.py`. +""" + +import torch + +import trainer.envs as envs +from trainer.logger import init_logger +from trainer.platforms.interface import (AttentionBackendEnum, + DeviceCapability, Platform, + PlatformEnum) + +logger = init_logger(__name__) + + +# ROCm uses the same torch.cuda interface +class RocmPlatform(Platform): + _enum = PlatformEnum.ROCM + device_name: str = "rocm" + device_type: str = "cuda" # torch uses 'cuda' backend string + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + return torch.cuda.get_device_properties(device_id).total_memory + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA graph. " + "Since enforce-eager is enabled, async output processor cannot be used" + ) + return False + return True + + @classmethod + def log_warnings(cls) -> None: + pass # ROCm-specific warnings can be added here + + @classmethod + def get_current_memory_usage(cls, + device: torch.device | None = None) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, + head_size: int, dtype: torch.dtype) -> str: + logger.info("Trying TRAINER_ATTENTION_BACKEND=%s", + envs.TRAINER_ATTENTION_BACKEND) + + if selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "trainer.attention.backends.sdpa.SDPABackend" + + elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None): + pass + + elif selected_backend in (AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.SAGE_ATTN): + raise ValueError( + f"{selected_backend.name} is not supported on {cls.device_name}." + ) + elif selected_backend: + raise ValueError( + f"Invalid attention backend for {cls.device_name}: {selected_backend}" + ) + + target_backend = AttentionBackendEnum.FLASH_ATTN + if dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention backend for dtype other than " + "torch.float16 or torch.bfloat16.") + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from trainer.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend) + + supported_sizes = \ + FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info("Cannot use FlashAttention backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default).") + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "trainer.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "trainer.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "trainer.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too diff --git a/trainer/third_party/__init__.py b/trainer/third_party/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/third_party/pynvml.py b/trainer/third_party/pynvml.py new file mode 100644 index 0000000000000000000000000000000000000000..0a4be23a09362ab6b11befb8e08c5e840f4f94ec --- /dev/null +++ b/trainer/third_party/pynvml.py @@ -0,0 +1,6139 @@ +# SPDX-License-Identifier: Apache-2.0 +# copied from https://pypi.org/project/nvidia-ml-py +# version 12.570.86 + +##### +# Copyright (c) 2011-2023, NVIDIA Corporation. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA Corporation nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +##### + +## +# Python bindings for the NVML library +## +from ctypes import * +from ctypes.util import find_library +from functools import wraps +import sys +import os +import threading +import string + +## C Type mappings ## +## Enums +_nvmlEnableState_t = c_uint +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 + +_nvmlBrandType_t = c_uint +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 17 + +_nvmlTemperatureThresholds_t = c_uint +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 + +_nvmlTemperatureSensors_t = c_uint +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 + + +_nvmlComputeMode_t = c_uint +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 +NVML_COMPUTEMODE_COUNT = 4 + +_nvmlMemoryLocation_t = c_uint +NVML_MEMORY_LOCATION_L1_CACHE = 0 +NVML_MEMORY_LOCATION_L2_CACHE = 1 +NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 +NVML_MEMORY_LOCATION_DRAM = 2 +NVML_MEMORY_LOCATION_REGISTER_FILE = 3 +NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 +NVML_MEMORY_LOCATION_TEXTURE_SHM = 5 +NVML_MEMORY_LOCATION_CBU = 6 +NVML_MEMORY_LOCATION_SRAM = 7 +NVML_MEMORY_LOCATION_COUNT = 8 + +NVML_NVLINK_MAX_LINKS = 18 + +# For backwards compatibility, maintain the incorrectly-named "LANES" define +NVML_NVLINK_MAX_LANES = NVML_NVLINK_MAX_LINKS + +_nvmlNvLinkErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_REPLAY = 0 +NVML_NVLINK_ERROR_DL_RECOVERY = 1 +NVML_NVLINK_ERROR_DL_CRC_FLIT = 2 +NVML_NVLINK_ERROR_DL_CRC_DATA = 3 +NVML_NVLINK_ERROR_DL_ECC_DATA = 4 +NVML_NVLINK_ERROR_COUNT = 5 + +_nvmlNvLinkEccLaneErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_ECC_LANE0 = 0 +NVML_NVLINK_ERROR_DL_ECC_LANE1 = 1 +NVML_NVLINK_ERROR_DL_ECC_LANE2 = 2 +NVML_NVLINK_ERROR_DL_ECC_LANE3 = 3 +NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 + +_nvmlNvLinkCapability_t = c_uint +NVML_NVLINK_CAP_P2P_SUPPORTED = 0 +NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS= 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 + +_nvmlNvLinkUtilizationCountPktTypes_t = c_uint +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF + +_nvmlNvLinkUtilizationCountUnits_t = c_uint +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 + +_nvmlNvLinkDeviceType_t = c_uint +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF + +# These are deprecated, instead use _nvmlMemoryErrorType_t +_nvmlEccBitType_t = c_uint +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 +NVML_ECC_ERROR_TYPE_COUNT = 2 + +_nvmlEccCounterType_t = c_uint +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 +NVML_ECC_COUNTER_TYPE_COUNT = 2 + +_nvmlMemoryErrorType_t = c_uint +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 + +_nvmlClockType_t = c_uint +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 + +_nvmlClockId_t = c_uint +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 +NVML_CLOCK_ID_COUNT = 4 + +_nvmlDriverModel_t = c_uint +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 + +NVML_MAX_GPU_PERF_PSTATES = 16 + +_nvmlPstates_t = c_uint +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 + +_nvmlInforomObject_t = c_uint +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 + +_nvmlReturn_t = c_uint +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_UNKNOWN = 999 + +_nvmlFanState_t = c_uint +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 + +_nvmlFanControlPolicy_t = c_uint +NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 +NVML_FAN_POLICY_MANUAL = 1 + +_nvmlLedColor_t = c_uint +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 + +_nvmlGpuOperationMode_t = c_uint +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 + +_nvmlPageRetirementCause_t = c_uint +NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 + +_nvmlRestrictedAPI_t = c_uint +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 + +_nvmlBridgeChipType_t = c_uint +NVML_BRIDGE_CHIP_PLX = 0 +NVML_BRIDGE_CHIP_BRO4 = 1 +NVML_MAX_PHYSICAL_BRIDGE = 128 + +_nvmlValueType_t = c_uint +NVML_VALUE_TYPE_DOUBLE = 0 +NVML_VALUE_TYPE_UNSIGNED_INT = 1 +NVML_VALUE_TYPE_UNSIGNED_LONG = 2 +NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 +NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4 +NVML_VALUE_TYPE_SIGNED_INT = 5 +NVML_VALUE_TYPE_UNSIGNED_SHORT = 6 +NVML_VALUE_TYPE_COUNT = 7 + +_nvmlNvlinkVersion_t = c_uint +NVML_NVLINK_VERSION_INVALID = 0 +NVML_NVLINK_VERSION_1_0 = 1 +NVML_NVLINK_VERSION_2_0 = 2 +NVML_NVLINK_VERSION_2_2 = 3 +NVML_NVLINK_VERSION_3_0 = 4 +NVML_NVLINK_VERSION_3_1 = 5 +NVML_NVLINK_VERSION_4_0 = 6 +NVML_NVLINK_VERSION_5_0 = 7 + +_nvmlPerfPolicyType_t = c_uint +NVML_PERF_POLICY_POWER = 0 +NVML_PERF_POLICY_THERMAL = 1 +NVML_PERF_POLICY_SYNC_BOOST = 2 +NVML_PERF_POLICY_BOARD_LIMIT = 3 +NVML_PERF_POLICY_LOW_UTILIZATION = 4 +NVML_PERF_POLICY_RELIABILITY = 5 +NVML_PERF_POLICY_TOTAL_APP_CLOCKS = 10 +NVML_PERF_POLICY_TOTAL_BASE_CLOCKS = 11 +NVML_PERF_POLICY_COUNT = 12 + +_nvmlEncoderQueryType_t = c_uint +NVML_ENCODER_QUERY_H264 = 0 +NVML_ENCODER_QUERY_HEVC = 1 +NVML_ENCODER_QUERY_AV1 = 2 +NVML_ENCODER_QUERY_UNKNOWN = 255 + +_nvmlFBCSessionType_t = c_uint +NVML_FBC_SESSION_TYPE_UNKNOWN = 0 +NVML_FBC_SESSION_TYPE_TOSYS = 1 +NVML_FBC_SESSION_TYPE_CUDA = 2 +NVML_FBC_SESSION_TYPE_VID = 3 +NVML_FBC_SESSION_TYPE_HWENC = 4 + +_nvmlDetachGpuState_t = c_uint +NVML_DETACH_GPU_KEEP = 0 +NVML_DETACH_GPU_REMOVE = 1 + +_nvmlPcieLinkState_t = c_uint +NVML_PCIE_LINK_KEEP = 0 +NVML_PCIE_LINK_SHUT_DOWN = 1 + +_nvmlSamplingType_t = c_uint +NVML_TOTAL_POWER_SAMPLES = 0 +NVML_GPU_UTILIZATION_SAMPLES = 1 +NVML_MEMORY_UTILIZATION_SAMPLES = 2 +NVML_ENC_UTILIZATION_SAMPLES = 3 +NVML_DEC_UTILIZATION_SAMPLES = 4 +NVML_PROCESSOR_CLK_SAMPLES = 5 +NVML_MEMORY_CLK_SAMPLES = 6 +NVML_MODULE_POWER_SAMPLES = 7 +NVML_JPG_UTILIZATION_SAMPLES = 8 +NVML_OFA_UTILIZATION_SAMPLES = 9 +NVML_SAMPLINGTYPE_COUNT = 10 + +_nvmlPcieUtilCounter_t = c_uint +NVML_PCIE_UTIL_TX_BYTES = 0 +NVML_PCIE_UTIL_RX_BYTES = 1 +NVML_PCIE_UTIL_COUNT = 2 + +_nvmlGpuTopologyLevel_t = c_uint +NVML_TOPOLOGY_INTERNAL = 0 +NVML_TOPOLOGY_SINGLE = 10 +NVML_TOPOLOGY_MULTIPLE = 20 +NVML_TOPOLOGY_HOSTBRIDGE = 30 +NVML_TOPOLOGY_NODE = 40 +NVML_TOPOLOGY_CPU = NVML_TOPOLOGY_NODE +NVML_TOPOLOGY_SYSTEM = 50 + +_nvmlGpuP2PCapsIndex_t = c_uint +NVML_P2P_CAPS_INDEX_READ = 0, +NVML_P2P_CAPS_INDEX_WRITE = 1 +NVML_P2P_CAPS_INDEX_NVLINK =2 +NVML_P2P_CAPS_INDEX_ATOMICS = 3 +# +# NVML_P2P_CAPS_INDEX_PROP is deprecated. +# Use NVML_P2P_CAPS_INDEX_PCI instead. +# +NVML_P2P_CAPS_INDEX_PROP = 4 +NVML_P2P_CAPS_INDEX_PCI = 4 +NVML_P2P_CAPS_INDEX_UNKNOWN = 5 + +_nvmlGpuP2PStatus_t = c_uint +NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED +NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY =4 +NVML_P2P_STATUS_NOT_SUPPORTED =5 +NVML_P2P_STATUS_UNKNOWN =6 + +_nvmlDeviceArchitecture_t = c_uint +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_T23X = 11 +NVML_DEVICE_ARCH_UNKNOWN = 0xffffffff + +# PCI bus Types +_nvmlBusType_t = c_uint +NVML_BUS_TYPE_UNKNOWN = 0 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 + +_nvmlPowerSource_t = c_uint +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 + +_nvmlAdaptiveClockInfoStatus_t = c_uint +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000 +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 + +_nvmlClockLimitId_t = c_uint +NVML_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00 +NVML_CLOCK_LIMIT_ID_TDP = 0xffffff01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xffffff02 + +_nvmlPcieLinkMaxSpeed_t = c_uint +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 +NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 +NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 + +_nvmlPcieAtomicsCapability_t = c_uint +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 + +_nvmlAffinityScope_t = c_uint +NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_SOCKET = 1 + +_nvmlDeviceGpuRecoveryAction_t = c_uint +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 + +# C preprocessor defined values +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 + +NVML_MAX_GPC_COUNT = 32 + +# buffer size +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +NVML_PERF_MODES_BUFFER_SIZE = 2048 + +# Format strings +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" + +NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) +NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) + +''' + Field Identifiers. + + All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. +''' +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive + +#ECC Count Totals +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors +#Individual ECC locations +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_DEV = 22 # Device memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_DEV = 23 # Device memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_REG = 24 # Register File single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_REG = 25 # Register File double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TEX = 26 # Texture memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TEX = 27 # Texture memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors + +# Page Retirement +NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors +NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors +NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = 32 # NVLink flow control CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = 33 # NVLink flow control CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = 34 # NVLink flow control CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = 35 # NVLink flow control CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = 36 # NVLink flow control CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = 37 # NVLink flow control CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC Error Counter total for all Lanes + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = 39 # NVLink data CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = 40 # NVLink data CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = 41 # NVLink data CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = 42 # NVLink data CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = 43 # NVLink data CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = 44 # NVLink data CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = 52 # NVLink Replay Error Counter total for all Lanes + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = 53 # NVLink Recovery Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = 54 # NVLink Recovery Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = 55 # NVLink Recovery Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = 56 # NVLink Recovery Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = 57 # NVLink Recovery Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = 58 # NVLink Recovery Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes + +# Perf Policy Counters +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = 78 # Perf Policy Counter for Low GPU Utilization Policy +NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = 80 # Perf Policy Counter for Total App Clock Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81 # Perf Policy Counter for Total Base Clocks Policy + +# Memory temperatures +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device + +# Energy Counter +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 + +# NVLink Link Count +NVML_FI_DEV_NVLINK_LINK_COUNT = 91 + +# Page Retirement pending fields +NVML_FI_DEV_RETIRED_PENDING_SBE = 92 +NVML_FI_DEV_RETIRED_PENDING_DBE = 93 + +# PCIe replay and replay rollover counters +NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 +NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = 96 # NVLink flow control CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = 97 # NVLink flow control CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = 98 # NVLink flow control CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = 99 # NVLink flow control CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = 100 # NVLink flow control CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = 101 # NVLink flow control CRC Error Counter for Lane 11 + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = 102 # NVLink data CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = 103 # NVLink data CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = 104 # NVLink data CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = 105 # NVLink data CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = 106 # NVLink data CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = 107 # NVLink data CRC Error Counter for Lane 11 + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = 112 # NVLink Replay Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = 113 # NVLink Replay Error Counter for Lane 11 + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = 114 # NVLink Recovery Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = 115 # NVLink Recovery Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = 116 # NVLink Recovery Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = 117 # NVLink Recovery Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = 118 # NVLink Recovery Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = 119 # NVLink Recovery Error Counter for Lane 11 + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 + +# NVLink Throughput Counters +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB + +# Row Remapper +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 + +#Remote device NVLink ID +NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 + +# Number of NVLinks connected to NVSwitch +NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 + +# NvLink ECC Data Error Counters +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = 148 #< NVLink data ECC Error Counter for Link 0 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = 149 #< NVLink data ECC Error Counter for Link 1 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = 150 #< NVLink data ECC Error Counter for Link 2 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = 151 #< NVLink data ECC Error Counter for Link 3 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = 152 #< NVLink data ECC Error Counter for Link 4 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = 153 #< NVLink data ECC Error Counter for Link 5 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = 154 #< NVLink data ECC Error Counter for Link 6 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = 155 #< NVLink data ECC Error Counter for Link 7 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = 156 #< NVLink data ECC Error Counter for Link 8 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = 157 #< NVLink data ECC Error Counter for Link 9 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = 158 #< NVLink data ECC Error Counter for Link 10 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = 159 #< NVLink data ECC Error Counter for Link 11 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links + +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 + +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 + +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 + +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 + +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 + +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 + +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 + +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 + +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 + +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 + +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 + +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 + +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 + +NVML_FI_DEV_RESET_STATUS = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 + +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = 258 # HW Circuitry % lifetime remaining +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = 261 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = 262 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = 263 # Ramp down hysteresis value in ms for a given profile +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = 265 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = 266 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = 267 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = 268 # Ramp down hysteresis value in ms for a given profile + +NVML_FI_MAX = 269 # One greater than the largest field ID defined above + +# NVML_FI_DEV_NVLINK_GET_STATE state enums +NVML_NVLINK_STATE_INACTIVE = 0x0 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 + +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS + +## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1 # Device is associated with GPU-Passthorugh +NVML_GPU_VIRTUALIZATION_MODE_VGPU = 2 # Device is associated with vGPU inside virtual machine. +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = 3 # Device is associated with VGX hypervisor in vGPU mode +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = 4 # Device is associated with VGX hypervisor in vSGA mode + +## Lib loading ## +nvmlLib = None +libLoadLock = threading.Lock() +_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown + +## vGPU Management +_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuInstance_t = c_uint + +_nvmlVgpuVmIdType_t = c_uint +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 + +_nvmlGridLicenseFeatureCode_t = c_uint +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 + +_nvmlGridLicenseExpiryStatus_t = c_uint8 +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = 0, # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = 1, # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = 2, # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = 3, # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = 4, # Permanent expiry + +_nvmlVgpuCapability_t = c_uint +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +NVML_VGPU_CAP_COUNT = 5 + +_nvmlVgpuDriverCapability_t = c_uint +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = 0 # Supports mixing of different vGPU profiles within one guest VM +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 + +_nvmlDeviceVgpuCapability_t = c_uint +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = 6 # Set/Get support of mini-quarter vGPU profiles +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = 7 # Set/Get support for compute media engine vGPU profiles +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = 8 # Query whether the GPU supports FSR and warm update +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_COUNT = 10 + +_nvmlVgpuGuestInfoState_t = c_uint +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 + +_nvmlVgpuVmCompatibility_t = c_uint +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 + +_nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 + +_nvmlHostVgpuMode_t = c_uint +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 + +_nvmlConfComputeGpusReadyState_t = c_uint +NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 +NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1 + +_nvmlConfComputeGpuCaps_t = c_uint +NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0 +NVML_CC_SYSTEM_GPUS_CC_CAPABLE = 1 + +_nvmlConfComputeCpuCaps_t = c_uint +NVML_CC_SYSTEM_CPU_CAPS_NONE = 0 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1 +NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV_SNP = 3 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SNP_VTOM = 4 + +_nvmlConfComputeDevToolsMode_t = c_uint +NVML_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0 +NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 + +NVML_CC_SYSTEM_MULTIGPU_NONE = 0 +NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 + +NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 +NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 +NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 + +_nvmlConfComputeCcFeature_t = c_uint +NVML_CC_SYSTEM_FEATURE_DISABLED = 0 +NVML_CC_SYSTEM_FEATURE_ENABLED = 1 + +_nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50 +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 + +# GSP firmware +NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + +class NVMLLibraryMismatchError(Exception): + pass + +## Error Checking ## +class NVMLError(Exception): + _valClassMapping = dict() + # List of currently known error codes + _errcode_to_string = { + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + def __new__(typ, value): + ''' + Maps value to a proper subclass of NVMLError. + See _extractNVMLErrorsAsClasses function for more details + ''' + if typ == NVMLError: + typ = NVMLError._valClassMapping.get(value, typ) + obj = Exception.__new__(typ) + obj.value = value + return obj + def __str__(self): + try: + if self.value not in NVMLError._errcode_to_string: + NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value)) + return NVMLError._errcode_to_string[self.value] + except NVMLError: + return "NVML Error with code %d" % self.value + def __eq__(self, other): + return self.value == other.value + +def nvmlExceptionClass(nvmlErrorCode): + if nvmlErrorCode not in NVMLError._valClassMapping: + raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode) + return NVMLError._valClassMapping[nvmlErrorCode] + +def _extractNVMLErrorsAsClasses(): + ''' + Generates a hierarchy of classes on top of NVMLError class. + + Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate + exceptions more easily. + + NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. + e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized + ''' + this_module = sys.modules[__name__] + nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")] + for err_name in nvmlErrorsNames: + # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized + class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "") + err_val = getattr(this_module, err_name) + def gen_new(val): + def new(typ): + obj = NVMLError.__new__(typ, val) + return obj + return new + new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)}) + new_error_class.__module__ = __name__ + setattr(this_module, class_name, new_error_class) + NVMLError._valClassMapping[err_val] = new_error_class +_extractNVMLErrorsAsClasses() + +def _nvmlCheckReturn(ret): + if (ret != NVML_SUCCESS): + raise NVMLError(ret) + return ret + +## Function access ## +_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking +def _nvmlGetFunctionPointer(name): + global nvmlLib + + if name in _nvmlGetFunctionPointer_cache: + return _nvmlGetFunctionPointer_cache[name] + + libLoadLock.acquire() + try: + # ensure library was loaded + if (nvmlLib == None): + raise NVMLError(NVML_ERROR_UNINITIALIZED) + try: + _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) + return _nvmlGetFunctionPointer_cache[name] + except AttributeError: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + +## Alternative object +# Allows the object to be printed +# Allows mismatched types to be assigned +# - like None when the Structure variant requires c_uint +class nvmlFriendlyObject(object): + def __init__(self, dictionary): + for x in dictionary: + setattr(self, x, dictionary[x]) + def __str__(self): + return self.__dict__.__str__() + +def nvmlStructToFriendlyObject(struct): + d = {} + for x in struct._fields_: + key = x[0] + value = getattr(struct, key) + # only need to convert from bytes if bytes, no need to check python version. + d[key] = value.decode() if isinstance(value, bytes) else value + obj = nvmlFriendlyObject(d) + return obj + +# pack the object so it can be passed to the NVML library +def nvmlFriendlyObjectToStruct(obj, model): + for x in model._fields_: + key = x[0] + value = obj.__dict__[key] + # any c_char_p in python3 needs to be bytes, default encoding works fine. + if sys.version_info >= (3,): + setattr(model, key, value.encode()) + else: + setattr(model, key, value) + return model + +## Unit structures +class struct_c_nvmlUnit_t(Structure): + pass # opaque handle +c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + +class _PrintableStructure(Structure): + """ + Abstract class that produces nicer __str__ output than ctypes.Structure. + e.g. instead of: + >>> print str(obj) + + this class will print + class_name(field_name: formatted_value, field_name: formatted_value) + + _fmt_ dictionary of -> + e.g. class that has _field_ 'hex_value', c_uint could be formatted with + _fmt_ = {"hex_value" : "%08X"} + to produce nicer output. + Default fomratting string for all fields can be set with key "" like: + _fmt_ = {"" : "%d MHz"} # e.g all values are numbers in MHz. + If not set it's assumed to be just "%s" + + Exact format of returned str from this class is subject to change in the future. + """ + _fmt_ = {} + def __str__(self): + result = [] + for x in self._fields_: + key = x[0] + value = getattr(self, key) + fmt = "%s" + if key in self._fmt_: + fmt = self._fmt_[key] + elif "" in self._fmt_: + fmt = self._fmt_[""] + result.append(("%s: " + fmt) % (key, value)) + return self.__class__.__name__ + "(" + ", ".join(result) + ")" + + def __getattribute__(self, name): + res = super(_PrintableStructure, self).__getattribute__(name) + # need to convert bytes to unicode for python3 don't need to for python2 + # Python 2 strings are of both str and bytes + # Python 3 strings are not of type bytes + # ctypes should convert everything to the correct values otherwise + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + def __setattr__(self, name, value): + if isinstance(value, str): + # encoding a python2 string returns the same value, since python2 strings are bytes already + # bytes passed in python3 will be ignored. + value = value.encode() + super(_PrintableStructure, self).__setattr__(name, value) + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ('name', c_char * 96), + ('id', c_char * 96), + ('serial', c_char * 96), + ('firmwareVersion', c_char * 96), + ] + +class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('isC2cEnabled', c_uint) + ] + +nvmlC2cModeInfo_v1 = 0x1000008; + +class c_nvmlLedState_t(_PrintableStructure): + _fields_ = [ + ('cause', c_char * 256), + ('color', _nvmlLedColor_t), + ] + +class c_nvmlPSUInfo_t(_PrintableStructure): + _fields_ = [ + ('state', c_char * 256), + ('current', c_uint), + ('voltage', c_uint), + ('power', c_uint), + ] + +class c_nvmlUnitFanInfo_t(_PrintableStructure): + _fields_ = [ + ('speed', c_uint), + ('state', _nvmlFanState_t), + ] + +class c_nvmlUnitFanSpeeds_t(_PrintableStructure): + _fields_ = [ + ('fans', c_nvmlUnitFanInfo_t * 24), + ('count', c_uint) + ] + +## Device structures +class struct_c_nvmlDevice_t(Structure): + pass # opaque handle +c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + +class nvmlPciInfoExt_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + ('pciSubSystemId', c_uint), + ('baseClass', c_uint), + ('subClass', c_uint), + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + 'version' : "0x%04X", + 'domain' : "0x%04X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + 'baseClass' : "0x%01X", + 'subClass' : "0x%01X", + } + +nvmlPciInfoExt_v1 = 0x1000040 + +# Legacy pciInfo used for _v1 and _v2 +class nvmlPciInfo_v2_t(_PrintableStructure): + _fields_ = [ + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + + # Added in 2.285 + ('pciSubSystemId', c_uint), + ('reserved0', c_uint), + ('reserved1', c_uint), + ('reserved2', c_uint), + ('reserved3', c_uint), + ] + _fmt_ = { + 'domain' : "0x%04X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + } + +class nvmlPciInfo_t(_PrintableStructure): + _fields_ = [ + # Moved to the new busId location below + ('busIdLegacy', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + + # Added in 2.285 + ('pciSubSystemId', c_uint), + # New busId replaced the long deprecated and reserved fields with a + # field of the same size in 9.0 + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + 'domain' : "0x%08X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + } + +class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ] + +SystemDriverBranchInfo_v1 = 0x1000054 + +class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): + _fields_ = [ + ('pci', nvmlPciInfo_t), + ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE) + ] + +class nvmlNvLinkUtilizationControl_t(_PrintableStructure): + _fields_ = [ + ('units', _nvmlNvLinkUtilizationCountUnits_t), + ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t), + ] + +class c_nvmlMemory_t(_PrintableStructure): + _fields_ = [ + ('total', c_ulonglong), + ('free', c_ulonglong), + ('used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +class c_nvmlMemory_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('total', c_ulonglong), + ('reserved', c_ulonglong), + ('free', c_ulonglong), + ('used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +nvmlMemory_v2 = 0x02000028 + +class c_nvmlBAR1Memory_t(_PrintableStructure): + _fields_ = [ + ('bar1Total', c_ulonglong), + ('bar1Free', c_ulonglong), + ('bar1Used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +class nvmlClkMonFaultInfo_t(Structure): + _fields_ = [("clkApiDomain", c_uint), + ("clkDomainFaultMask", c_uint) + ] + +MAX_CLK_DOMAINS = 32 + +class nvmlClkMonStatus_t(Structure): + _fields_ = [("bGlobalStatus", c_uint), + ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS) + ] + +# On Windows with the WDDM driver, usedGpuMemory is reported as None +# Code that processes this structure should check for None, I.E. +# +# if (info.usedGpuMemory == None): +# # TODO handle the error +# pass +# else: +# print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) +# endif +# +# See NVML documentation for more information +class c_nvmlProcessInfo_v2_t(_PrintableStructure): + _fields_ = [ + ('pid', c_uint), + ('usedGpuMemory', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint), + ] + _fmt_ = {'usedGpuMemory': "%d B"} + +c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t + +c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t + +_nvmlProcessMode_t = c_uint +NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_GRAPHICS = 1 +NVML_PROCESS_MODE_MPS = 2 + +class c_nvmlProcessDetail_v1_t(Structure): + _fields_ = [ + ('pid', c_uint), + ('usedGpuMemory', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint), + ('usedGpuCcProtectedMemory', c_ulonglong), + ] + +class c_nvmlProcessDetailList_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('mode', _nvmlProcessMode_t), + ('numProcArrayEntries', c_uint), + ('procArray', POINTER(c_nvmlProcessDetail_v1_t)), + ] + _fmt_ = {'numProcArrayEntries': "%d B"} + +c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t + +nvmlProcessDetailList_v1 = 0x1000018 + +class c_nvmlBridgeChipInfo_t(_PrintableStructure): + _fields_ = [ + ('type', _nvmlBridgeChipType_t), + ('fwVersion', c_uint), + ] + +class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): + _fields_ = [ + ('bridgeCount', c_uint), + ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), + ] + +class c_nvmlEccErrorCounts_t(_PrintableStructure): + _fields_ = [ + ('l1Cache', c_ulonglong), + ('l2Cache', c_ulonglong), + ('deviceMemory', c_ulonglong), + ('registerFile', c_ulonglong), + ] + +class c_nvmlUtilization_t(_PrintableStructure): + _fields_ = [ + ('gpu', c_uint), + ('memory', c_uint), + ] + _fmt_ = {'': "%d %%"} + +# Added in 2.285 +class c_nvmlHwbcEntry_t(_PrintableStructure): + _fields_ = [ + ('hwbcId', c_uint), + ('firmwareVersion', c_char * 32), + ] + +class c_nvmlValue_t(Union): + _fields_ = [ + ('dVal', c_double), + ('uiVal', c_uint), + ('ulVal', c_ulong), + ('ullVal', c_ulonglong), + ('sllVal', c_longlong), + ('siVal', c_int), + ('usVal', c_ushort), + ] + +class c_nvmlSample_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('sampleValue', c_nvmlValue_t), + ] + +class c_nvmlViolationTime_t(_PrintableStructure): + _fields_ = [ + ('referenceTime', c_ulonglong), + ('violationTime', c_ulonglong), + ] + +class c_nvmlFieldValue_t(_PrintableStructure): + _fields_ = [ + ('fieldId', c_uint32), + ('scopeId', c_uint32), + ('timestamp', c_int64), + ('latencyUsec', c_int64), + ('valueType', _nvmlValueType_t), + ('nvmlReturn', _nvmlReturn_t), + ('value', c_nvmlValue_t) + ] + +NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 + +nvmlNvlinkSupportedBwModes_v1 = 0x100001c +class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ('totalBwModes', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(version=nvmlNvlinkSupportedBwModes_v1) + +nvmlNvlinkGetBwMode_v1 = 0x100000c +class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bIsBest', c_uint), + ('bwMode', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) + +nvmlNvlinkSetBwMode_v1 = 0x100000c +class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bSetBest', c_uint), + ('bwMode', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) + +class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('mode', c_uint), + ] + +VgpuHeterogeneousMode_v1 = 0x1000008 + +class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('placementId', c_uint), + ] + +VgpuPlacementId_v1 = 0x1000008 + +class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('count', c_uint), + ('placementSize', c_uint), + ('placementIds', POINTER(c_uint)), + ] + +VgpuPlacementList_v1 = 0x1000018 + +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + +class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('placementSize', c_uint), + ('count', c_uint), + ('placementIds', POINTER(c_uint)), + ('mode', c_uint), + ] + +VgpuPlacementList_v2 = 0x2000020 + +class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bar1Size', c_ulonglong), + ] + +VgpuTypeBar1Info_v1 = 0x1000010 + +class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('vgpuInstance', _nvmlVgpuInstance_t), + ('timeStamp', c_ulonglong), + ('smUtil', c_nvmlValue_t), + ('memUtil', c_nvmlValue_t), + ('encUtil', c_nvmlValue_t), + ('decUtil', c_nvmlValue_t), + ] + +class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('smUtil', c_nvmlValue_t), + ('memUtil', c_nvmlValue_t), + ('encUtil', c_nvmlValue_t), + ('decUtil', c_nvmlValue_t), + ('jpgUtil', c_nvmlValue_t), + ('ofaUtil', c_nvmlValue_t), + ] + +class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('sampleValType', _nvmlValueType_t), + ('vgpuInstanceCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), + ] + +VgpuInstancesUtilizationInfo_v1 = 0x01000020 + +class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('vgpuInstance', _nvmlVgpuInstance_t), + ('pid', c_uint), + ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ('timeStamp', c_ulonglong), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ] + +class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ('timeStamp', c_ulonglong), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('pid', c_uint), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ('jpgUtil', c_uint), + ('ofaUtil', c_uint), + ] + +class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuProcessCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), + ] + +VgpuProcessesUtilizationInfo_v1 = 0x01000018 + +class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('size', c_ulonglong), + ] + +VgpuRuntimeState_v1 = 0x1000010 + +class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), + ] + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 +NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + +class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): + _fields_ = [ + ('isLicensed', c_uint8), + ('licenseExpiry', c_nvmlVgpuLicenseExpiry_t), + ('currentState', c_uint), + ] + +class c_nvmlEncoderSession_t(_PrintableStructure): + _fields_ = [ + ('sessionId', c_uint), + ('pid', c_uint), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('codecType', c_uint), + ('hResolution', c_uint), + ('vResolution', c_uint), + ('averageFps', c_uint), + ('encodeLatency', c_uint), + ] + +class c_nvmlProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('pid', c_uint), + ('timeStamp', c_ulonglong), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ] + +class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('pid', c_uint), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ('jpgUtil', c_uint), + ('ofaUtil', c_uint), + ] + +class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('processSamplesCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)), + ] + +ProcessesUtilizationInfo_v1 = 0x01000018 + +class c_nvmlGridLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), + ] + +class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('featureEnabled', c_uint), + ('licenseExpiry', c_nvmlGridLicenseExpiry_t), + ] + +class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('featureEnabled', c_uint), + ] + +class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + +class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + +class c_nvmlGridLicensableFeatures_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlMarginTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('marginTemperature', c_int), + ] + +nvmlMarginTemperature_v1 = 0x1000008 + +## Event structures +class struct_c_nvmlEventSet_t(Structure): + pass # opaque handle +c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) + +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = ( + nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction + ) + +## Clock Event Reasons defines +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = ( + nvmlClocksEventReasonNone | + nvmlClocksEventReasonGpuIdle | + nvmlClocksEventReasonApplicationsClocksSetting | + nvmlClocksEventReasonSwPowerCap | + nvmlClocksEventReasonHwSlowdown | + nvmlClocksEventReasonSyncBoost | + nvmlClocksEventReasonSwThermalSlowdown | + nvmlClocksEventReasonHwThermalSlowdown | + nvmlClocksEventReasonHwPowerBrakeSlowdown | + nvmlClocksEventReasonDisplayClockSetting + ) + +## Following have been deprecated +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = ( + nvmlClocksThrottleReasonNone | + nvmlClocksThrottleReasonGpuIdle | + nvmlClocksThrottleReasonApplicationsClocksSetting | + nvmlClocksThrottleReasonSwPowerCap | + nvmlClocksThrottleReasonHwSlowdown | + nvmlClocksThrottleReasonSyncBoost | + nvmlClocksThrottleReasonSwThermalSlowdown | + nvmlClocksThrottleReasonHwThermalSlowdown | + nvmlClocksThrottleReasonHwPowerBrakeSlowdown | + nvmlClocksThrottleReasonDisplayClockSetting + ) + +class c_nvmlEventData_t(_PrintableStructure): + _fields_ = [ + ('device', c_nvmlDevice_t), + ('eventType', c_ulonglong), + ('eventData', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint) + ] + _fmt_ = {'eventType': "0x%08X"} + +class c_nvmlAccountingStats_t(_PrintableStructure): + _fields_ = [ + ('gpuUtilization', c_uint), + ('memoryUtilization', c_uint), + ('maxMemoryUsage', c_ulonglong), + ('time', c_ulonglong), + ('startTime', c_ulonglong), + ('isRunning', c_uint), + ('reserved', c_uint * 5) + ] + +class c_nvmlVgpuVersion_t(Structure): + _fields_ = [("minVersion", c_uint), + ("maxVersion", c_uint) + ] + +class c_nvmlVgpuMetadata_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("revision", c_uint), + ("guestInfoState", _nvmlVgpuGuestInfoState_t), + ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), + ("guestVgpuVersion", c_uint), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE) + ] + +class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("revision", c_uint), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("pgpuVirtualizationCaps", c_uint), + ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) + ] + +class c_nvmlVgpuPgpuCompatibility_t(Structure): + _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t) + ] + +## vGPU scheduler policy defines +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 + +## Supported vGPU scheduler policy count +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 + +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 + +class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): + _fields_ = [ + ('avgFactor', c_uint), + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedData_t(_PrintableStructure): + _fields_ = [ + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedulerParams_t(Union): + _fields_ = [ + ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t), + ('vgpuSchedData', c_nvmlVgpuSchedData_t), + ] + +class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): + _fields_ = [ + ('timestamp', c_ulonglong), + ('timeRunTotal', c_ulonglong), + ('timeRun', c_ulonglong), + ('swRunlistId', c_uint), + ('targetTimeSlice', c_ulonglong), + ('cumulativePreemptionTime', c_ulonglong), + ] + +class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): + _fields_ = [ + ('engineId', c_uint), + ('schedulerPolicy', c_uint), + ('arrMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), + ('entriesCount', c_uint), + ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), + ] + +class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): + _fields_ = [ + ('schedulerPolicy', c_uint), + ('arrMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), + ] + +class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): + _fields_ = [ + ('avgFactor', c_uint), + ('frequency', c_uint), + ] + +class c_nvmlVgpuSchedSetData_t(_PrintableStructure): + _fields_ = [ + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedulerSetParams_t(Union): + _fields_ = [ + ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t), + ('vgpuSchedData', c_nvmlVgpuSchedSetData_t), + ] + +class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): + _fields_ = [ + ('schedulerPolicy', c_uint), + ('enableARRMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t), + ] + +class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): + _fields_ = [ + ('supportedSchedulers', c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), + ('maxTimeslice', c_uint), + ('minTimeslice', c_uint), + ('isArrModeSupported', c_uint), + ('maxFrequencyForARR', c_uint), + ('minFrequencyForARR', c_uint), + ('maxAvgFactorForARR', c_uint), + ('minAvgFactorForARR', c_uint), + ] + +class c_nvmlFBCStats_t(Structure): + _fields_ = [("sessionsCount", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint) + ] + +class c_nvmlFBCSession_t(_PrintableStructure): + _fields_ = [ + ('sessionId', c_uint), + ('pid', c_uint), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('displayOrdinal', c_uint), + ('sessionType', c_uint), + ('sessionFlags', c_uint), + ('hMaxResolution', c_uint), + ('vMaxResolution', c_uint), + ('hResolution', c_uint), + ('vResolution', c_uint), + ('averageFPS', c_uint), + ('averageLatency', c_uint), + ] + +NVML_DEVICE_MIG_DISABLE = 0x0 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD + +class c_nvmlGpuInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), + ("size", c_uint) + ] + +class c_nvmlGpuInstanceProfileInfo_t(Structure): + _fields_ = [("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + +nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + +class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) + ] + + def __init__(self): + super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2) + +class c_nvmlGpuInstanceInfo_t(Structure): + _fields_ = [("device", c_nvmlDevice_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t) + ] + +class struct_c_nvmlGpuInstance_t(Structure): + pass # opaque handle +c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) + +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 + +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + +class c_nvmlComputeInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), + ("size", c_uint) + ] + +class c_nvmlComputeInstanceProfileInfo_t(Structure): + _fields_ = [("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint) + ] + +nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + +class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) + ] + + def __init__(self): + super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2) + +class c_nvmlComputeInstanceInfo_t(Structure): + _fields_ = [("device", c_nvmlDevice_t), + ("gpuInstance", c_nvmlGpuInstance_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t) + ] + +NVML_MAX_GPU_UTILIZATIONS = 8 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 +class c_nvmlGpuDynamicPstatesUtilization_t(Structure): + _fields_ = [("bIsPresent", c_uint, 1), + ("percentage", c_uint), + ("incThreshold", c_uint), + ("decThreshold", c_uint)] +class c_nvmlGpuDynamicPstatesInfo_t(Structure): + _fields_ = [("flags", c_uint), + ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)] + +NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 + +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 + +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + +class c_nvmlGpuThermalSensor_t(Structure): + _fields_ = [("controller", c_int), + ("defaultMinTemp", c_int), + ("defaultMaxTemp", c_int), + ("currentTemp", c_int), + ("target", c_int)] +class c_nvmlGpuThermalSettings_t(Structure): + _fields_ = [("count", c_uint), + ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)] + +_nvmlCoolerControl_t = c_uint +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 + +_nvmlCoolerTarget_t = c_uint +NVML_THERMAL_COOLER_TARGET_NONE = (1 << 0) +NVML_THERMAL_COOLER_TARGET_GPU = (1 << 1) +NVML_THERMAL_COOLER_TARGET_MEMORY = (1 << 2) +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = (1 << 3) +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY) + +class c_nvmlCoolerInfo_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("index", c_uint), + ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t) + ] + +nvmlCoolerInfo_v1 = 0x1000010 + +def nvmlDeviceGetCoolerInfo(handle): + c_coolerInfo = c_nvmlCoolerInfo_t() + c_coolerInfo.version = nvmlCoolerInfo_v1 + c_coolerInfo.index = 0 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCoolerInfo") + ret = fn(handle, byref(c_coolerInfo)) + _nvmlCheckReturn(ret) + return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + +class struct_c_nvmlComputeInstance_t(Structure): + pass # opaque handle +c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + +class c_nvmlDeviceAttributes(Structure): + _fields_ = [("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + +class c_nvmlRowRemapperHistogramValues(Structure): + _fields_ = [("max", c_uint), + ("high", c_uint), + ("partial", c_uint), + ("low", c_uint), + ("none", c_uint) + ] + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 +NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 +NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + +class c_nvmlConfComputeSystemState_t(Structure): + _fields_ = [('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ] + +nvmlSystemConfComputeSettings_v1 = 0x1000014 + +class c_nvmlSystemConfComputeSettings_v1_t(Structure): + _fields_ = [('version', c_uint), + ('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ('multiGpuMode', c_uint), + ] + def __init__(self): + super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1) + +class c_nvmlConfComputeSystemCaps_t(Structure): + _fields_ = [('cpuCaps', c_uint), + ('gpusCaps', c_uint), + ] + +class c_nvmlConfComputeMemSizeInfo_t(Structure): + _fields_ = [('protectedMemSizeKib', c_ulonglong), + ('unprotectedMemSizeKib', c_ulonglong), + ] + +class c_nvmlConfComputeGpuCertificate_t(Structure): + _fields_ = [('certChainSize', c_uint), + ('attestationCertChainSize', c_uint), + ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), + ] + +class c_nvmlConfComputeGpuAttestationReport_t(Structure): + _fields_ = [('isCecAttestationReportPresent', c_uint), + ('attestationReportSize', c_uint), + ('cecAttestationReportSize', c_uint), + ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), + ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), + ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), + ] + +class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): + _fields_ = [('version', c_uint), + ('maxAttackerAdvantage', c_ulong), + ] +ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + +class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): + _fields_ = [('version', c_uint), + ('attackerAdvantage', c_ulong), + ] +ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 + + +## string/bytes conversion for ease of use +def convertStrBytes(func): + ''' + In python 3, strings are unicode instead of bytes, and need to be converted for ctypes + Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) + Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> + ---- + Returned from function: b'returned string' + Returned to caller: 'returned string' + ''' + @wraps(func) + def wrapper(*args, **kwargs): + # encoding a str returns bytes in python 2 and 3 + args = [arg.encode() if isinstance(arg, str) else arg for arg in args] + res = func(*args, **kwargs) + # In python 2, str and bytes are the same + # In python 3, str is unicode and should be decoded. + # Ctypes handles most conversions, this only effects c_char and char arrays. + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + if sys.version_info >= (3,): + return wrapper + return func + +def throwOnVersionMismatch(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except NVMLError_FunctionNotFound: + raise NVMLLibraryMismatchError("Unversioned function called and the " + "pyNVML version does not match the NVML lib version. " + "Either use matching pyNVML and NVML lib versions or " + "use a versioned function such as " + func.__name__ + "_v2") + return wrapper + +## C function wrappers ## +def nvmlInitWithFlags(flags): + _LoadNvmlLibrary() + + # + # Initialize the library + # + fn = _nvmlGetFunctionPointer("nvmlInitWithFlags") + ret = fn(flags) + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + _nvmlLib_refcount += 1 + libLoadLock.release() + return None + +def nvmlInit(): + nvmlInitWithFlags(0) + return None + +def _LoadNvmlLibrary(): + ''' + Load the library if it isn't loaded already + ''' + global nvmlLib + + if (nvmlLib == None): + # lock to ensure only one caller loads the library + libLoadLock.acquire() + + try: + # ensure the library still isn't loaded + if (nvmlLib == None): + try: + if (sys.platform[:3] == "win"): + # cdecl calling convention + try: + # Check for nvml.dll in System32 first for DCH drivers + nvmlLib = CDLL(os.path.join(os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll")) + except OSError as ose: + # If nvml.dll is not found in System32, it should be in ProgramFiles + # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll + nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll")) + else: + # assume linux + nvmlLib = CDLL("libnvidia-ml.so.1") + except OSError as ose: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + if (nvmlLib == None): + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + +def nvmlShutdown(): + # + # Leave the library loaded, but shutdown the interface + # + fn = _nvmlGetFunctionPointer("nvmlShutdown") + ret = fn() + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + if (0 < _nvmlLib_refcount): + _nvmlLib_refcount -= 1 + libLoadLock.release() + return None + +# Added in 2.285 +@convertStrBytes +def nvmlErrorString(result): + fn = _nvmlGetFunctionPointer("nvmlErrorString") + fn.restype = c_char_p # otherwise return is an int + ret = fn(result) + return ret + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetNVMLVersion(): + c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +def nvmlSystemGetCudaDriverVersion(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + +def nvmlSystemGetCudaDriverVersion_v2(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetProcessName(pid): + c_name = create_string_buffer(1024) + fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") + ret = fn(c_uint(pid), c_name, c_uint(1024)) + _nvmlCheckReturn(ret) + return c_name.value + +@convertStrBytes +def nvmlSystemGetDriverVersion(): + c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 2.285 +def nvmlSystemGetHicVersion(): + c_count = c_uint(0) + hics = None + fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") + + # get the count + ret = fn(byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # If there are no hics + if (c_count.value == 0): + return [] + + hic_array = c_nvmlHwbcEntry_t * c_count.value + hics = hic_array() + ret = fn(byref(c_count), hics) + _nvmlCheckReturn(ret) + return hics + +def nvmlSystemGetDriverBranch(): + c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) + c_branchInfo.version = SystemDriverBranchInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_branchInfo + +## Unit get functions +def nvmlUnitGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlUnitGetHandleByIndex(index): + c_index = c_uint(index) + unit = c_nvmlUnit_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") + ret = fn(c_index, byref(unit)) + _nvmlCheckReturn(ret) + return unit + +def nvmlUnitGetUnitInfo(unit): + c_info = c_nvmlUnitInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlUnitGetLedState(unit): + c_state = c_nvmlLedState_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") + ret = fn(unit, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + +def nvmlUnitGetPsuInfo(unit): + c_info = c_nvmlPSUInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlUnitGetTemperature(unit, type): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") + ret = fn(unit, c_uint(type), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlUnitGetFanSpeedInfo(unit): + c_speeds = c_nvmlUnitFanSpeeds_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") + ret = fn(unit, byref(c_speeds)) + _nvmlCheckReturn(ret) + return c_speeds + +# added to API +def nvmlUnitGetDeviceCount(unit): + c_count = c_uint(0) + # query the unit to determine device count + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), None) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = NVML_SUCCESS + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlUnitGetDevices(unit): + c_count = c_uint(nvmlUnitGetDeviceCount(unit)) + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return c_devices + +## Device get functions +def nvmlDeviceGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetHandleByIndex(index): + c_index = c_uint(index) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") + ret = fn(c_index, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleBySerial(serial): + c_serial = c_char_p(serial) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") + ret = fn(c_serial, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleByUUID(uuid): + c_uuid = c_char_p(uuid) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") + ret = fn(c_uuid, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleByPciBusId(pciBusId): + c_busId = c_char_p(pciBusId) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") + ret = fn(c_busId, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetName(handle): + c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") + ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_name.value + +class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + +nvmlDevicePerfModes_v1 = 0x1000804 + +@convertStrBytes +def nvmlDeviceGetPerformanceModes(handle): + perfModes = c_nvmlDevicePerfModes_v1_t() + perfModes.version = nvmlDevicePerfModes_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceModes") + ret = fn(handle, byref(perfModes)) + _nvmlCheckReturn(ret) + return perfModes.str + +class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + +nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + +@convertStrBytes +def nvmlDeviceGetCurrentClockFreqs(handle): + currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() + currentClockFreqs.version = nvmlDeviceCurrentClockFreqs_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClockFreqs") + ret = fn(handle, byref(currentClockFreqs)) + _nvmlCheckReturn(ret) + return currentClockFreqs.str + +def nvmlDeviceGetBoardId(handle): + c_id = c_uint(); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") + ret = fn(handle, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + +def nvmlDeviceGetMultiGpuBoard(handle): + c_multiGpu = c_uint(); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") + ret = fn(handle, byref(c_multiGpu)) + _nvmlCheckReturn(ret) + return c_multiGpu.value + +def nvmlDeviceGetBrand(handle): + c_type = _nvmlBrandType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") + ret = fn(handle, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + +def nvmlDeviceGetC2cModeInfoV1(handle): + c_info = c_nvmlC2cModeInfo_v1_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlDeviceGetC2cModeInfoV(handle): + return nvmlDeviceGetC2cModeInfoV1(handle) + +@convertStrBytes +def nvmlDeviceGetBoardPartNumber(handle): + c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardPartNumber") + ret = fn(handle, c_part_number, c_uint(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_part_number.value + +@convertStrBytes +def nvmlDeviceGetSerial(handle): + c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") + ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_serial.value + +def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): + isReference = type(moduleId) is not c_uint + moduleIdRef = moduleId if isReference else byref(moduleId) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetModuleId") + ret = fn(handle, moduleIdRef) + if isReference: + return ret + else: + _nvmlCheckReturn(ret) + return moduleId.value + +def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): + affinity_array = c_ulonglong * nodeSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryAffinity") + ret = fn(handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinityWithinScope") + ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") + ret = fn(handle, cpuSetSize, byref(c_affinity)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceSetCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetNumaNodeId(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") + node = c_int() + ret = fn(handle, byref(node)) + _nvmlCheckReturn(ret) + return node.value + +def nvmlDeviceGetMinorNumber(handle): + c_minor_number = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") + ret = fn(handle, byref(c_minor_number)) + _nvmlCheckReturn(ret) + return c_minor_number.value + +@convertStrBytes +def nvmlDeviceGetUUID(handle): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") + ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlDeviceGetInforomVersion(handle, infoRomObject): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") + ret = fn(handle, _nvmlInforomObject_t(infoRomObject), + c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 4.304 +@convertStrBytes +def nvmlDeviceGetInforomImageVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 4.304 +def nvmlDeviceGetInforomConfigurationChecksum(handle): + c_checksum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") + ret = fn(handle, byref(c_checksum)) + _nvmlCheckReturn(ret) + return c_checksum.value + +# Added in 4.304 +def nvmlDeviceValidateInforom(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetLastBBXFlushTime(handle): + c_timestamp = c_ulonglong() + c_durationUs = c_ulong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetLastBBXFlushTime") + ret = fn(handle, byref(c_timestamp), byref(c_durationUs)) + _nvmlCheckReturn(ret) + return [c_timestamp.value, c_durationUs.value] + +def nvmlDeviceGetDisplayMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceGetDisplayActive(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetPersistenceMode(handle): + c_state = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") + ret = fn(handle, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +def nvmlDeviceGetPciInfoExt(handle, c_info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") + ret = fn(handle, c_info) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetPciInfo_v3(handle): + c_info = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlDeviceGetPciInfo(handle): + return nvmlDeviceGetPciInfo_v3(handle) + +def nvmlDeviceGetClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 2.285 +def nvmlDeviceGetMaxClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 4.304 +def nvmlDeviceGetApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +def nvmlDeviceGetMaxCustomerBoostClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +def nvmlDeviceGetClock(handle, type, id): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") + ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 5.319 +def nvmlDeviceGetDefaultApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 4.304 +def nvmlDeviceGetSupportedMemoryClocks(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no clocks + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + +# Added in 4.304 +def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no clocks + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetFanSpeed(handle): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") + ret = fn(handle, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetFanSpeed_v2(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +class c_nvmlFanSpeedInfo_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('fan', c_uint), + ('speed', c_uint), + ] + +nvmlFanSpeedInfo_v1 = 0x100000C + +def nvmlDeviceGetFanSpeedRPM(handle): + c_fanSpeed = c_nvmlFanSpeedInfo_t() + c_fanSpeed.fan = 0 + c_fanSpeed.version = nvmlFanSpeedInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeedRPM") + ret = fn(handle, byref(c_fanSpeed)) + _nvmlCheckReturn(ret) + return c_fanSpeed.speed + +def nvmlDeviceGetTargetFanSpeed(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetNumFans(device): + c_numFans = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") + ret = fn(device, byref(c_numFans)) + _nvmlCheckReturn(ret) + return c_numFans.value + +def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2"); + ret = fn(handle, index) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): + isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) + minSpeedRef = minSpeed if isReference else byref(minSpeed) + maxSpeedRef = maxSpeed if isReference else byref(maxSpeed) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxFanSpeed") + ret = fn(handle, minSpeedRef, maxSpeedRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + +def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): + isReference = type(fanControlPolicy) is not c_uint + fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanControlPolicy_v2") + ret = fn(handle, fan, fanControlPolicyRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else fanControlPolicy.value + +def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") + ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +class c_nvmlTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('sensorType', _nvmlTemperatureSensors_t), + ('temperature', c_int), + ] +nvmlTemperature_v1 = 0x100000C + +def nvmlDeviceGetTemperatureV1(handle, sensor): + c_temp = c_nvmlTemperature_v1_t() + c_temp.version = nvmlTemperature_v1 + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") + ret = fn(handle, byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.temperature + +def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): + if version == nvmlTemperature_v1: + return nvmlDeviceGetTemperatureV1(handle, sensor) + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + +# DEPRECATED use nvmlDeviceGetTemperatureV instead +def nvmlDeviceGetTemperature(handle, sensor): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") + ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlDeviceGetTemperatureThreshold(handle, threshold): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): + c_temp = c_uint() + c_temp.value = temp + fn = _nvmlGetFunctionPointer("nvmlDeviceSetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetMarginTemperature(handle): + c_marginTempInfo = c_nvmlMarginTemperature_v1_t() + c_marginTempInfo.version = nvmlMarginTemperature_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMarginTemperature") + ret = fn(handle, byref(c_marginTempInfo)) + _nvmlCheckReturn(ret) + return c_marginTempInfo.marginTemperature + +# DEPRECATED use nvmlDeviceGetPerformanceState +def nvmlDeviceGetPowerState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + +def nvmlDeviceGetPerformanceState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + +def nvmlDeviceGetPowerManagementMode(handle): + c_pcapMode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") + ret = fn(handle, byref(c_pcapMode)) + _nvmlCheckReturn(ret) + return c_pcapMode.value + +def nvmlDeviceGetPowerManagementLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + +# Added in 4.304 +def nvmlDeviceGetPowerManagementLimitConstraints(handle): + c_minLimit = c_uint() + c_maxLimit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") + ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) + _nvmlCheckReturn(ret) + return [c_minLimit.value, c_maxLimit.value] + +# Added in 4.304 +def nvmlDeviceGetPowerManagementDefaultLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 331 +def nvmlDeviceGetEnforcedPowerLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + +def nvmlDeviceGetPowerUsage(handle): + c_watts = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") + ret = fn(handle, byref(c_watts)) + _nvmlCheckReturn(ret) + return c_watts.value + +def nvmlDeviceGetTotalEnergyConsumption(handle): + c_millijoules = c_uint64() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") + ret = fn(handle, byref(c_millijoules)) + _nvmlCheckReturn(ret) + return c_millijoules.value + +# Added in 4.304 +def nvmlDeviceGetGpuOperationMode(handle): + c_currState = _nvmlGpuOperationMode_t() + c_pendingState = _nvmlGpuOperationMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + +# Added in 4.304 +def nvmlDeviceGetCurrentGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[0] + +# Added in 4.304 +def nvmlDeviceGetPendingGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[1] + +def nvmlDeviceGetMemoryInfo(handle, version=None): + if not version: + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") + else: + c_memory = c_nvmlMemory_v2_t() + c_memory.version = version + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo_v2") + ret = fn(handle, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + +def nvmlDeviceGetBAR1MemoryInfo(handle): + c_bar1_memory = c_nvmlBAR1Memory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") + ret = fn(handle, byref(c_bar1_memory)) + _nvmlCheckReturn(ret) + return c_bar1_memory + +def nvmlDeviceGetComputeMode(handle): + c_mode = _nvmlComputeMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceGetCudaComputeCapability(handle): + c_major = c_int() + c_minor = c_int() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") + ret = fn(handle, byref(c_major), byref(c_minor)) + _nvmlCheckReturn(ret) + return (c_major.value, c_minor.value) + +def nvmlDeviceGetEccMode(handle): + c_currState = _nvmlEnableState_t() + c_pendingState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + +# added to API +def nvmlDeviceGetCurrentEccMode(handle): + return nvmlDeviceGetEccMode(handle)[0] + +# added to API +def nvmlDeviceGetPendingEccMode(handle): + return nvmlDeviceGetEccMode(handle)[1] + +def nvmlDeviceGetDefaultEccMode(handle): + c_defaultState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") + ret = fn(handle, byref(c_defaultState)) + _nvmlCheckReturn(ret) + return [c_defaultState.value] + +def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter +def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): + c_counts = c_nvmlEccErrorCounts_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), byref(c_counts)) + _nvmlCheckReturn(ret) + return c_counts + +# Added in 4.304 +def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") + ret = fn(handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), + byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetUtilizationRates(handle): + c_util = c_nvmlUtilization_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") + ret = fn(handle, byref(c_util)) + _nvmlCheckReturn(ret) + return c_util + +def nvmlDeviceGetEncoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetDecoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetJpgUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetJpgUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetOfaUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetOfaUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetPcieReplayCounter(handle): + c_replay = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") + ret = fn(handle, byref(c_replay)) + _nvmlCheckReturn(ret) + return c_replay.value + +def nvmlDeviceGetDriverModel(handle): + c_currModel = _nvmlDriverModel_t() + c_pendingModel = _nvmlDriverModel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") + ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) + _nvmlCheckReturn(ret) + return [c_currModel.value, c_pendingModel.value] + +# added to API +def nvmlDeviceGetCurrentDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[0] + +# added to API +def nvmlDeviceGetPendingDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[1] + +# Added in 2.285 +@convertStrBytes +def nvmlDeviceGetVbiosVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +@throwOnVersionMismatch +def nvmlDeviceGetComputeRunningProcesses(handle): + return nvmlDeviceGetComputeRunningProcesses_v3(handle) + +def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +@throwOnVersionMismatch +def nvmlDeviceGetGraphicsRunningProcesses(handle): + return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + +@throwOnVersionMismatch +def nvmlDeviceGetMPSComputeRunningProcesses(handle): + return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + +def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): + c_processDetailList = c_nvmlProcessDetailList_t() + c_processDetailList.version = version + c_processDetailList.mode = mode + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRunningProcessDetailList") + + # first call to get the size + ret = fn(handle, byref(c_processDetailList)) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries + c_processDetailList.procArray = cast((c_procs)(), POINTER(c_nvmlProcessDetail_v1_t)) + + # make the call again + ret = fn(handle, byref(c_processDetailList)) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_processDetailList.numProcArrayEntries): + # use an alternative struct for this object + obj = c_processDetailList.procArray[i] + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + obj.usedGpuMemory = None + if (obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + obj.usedGpuCcProtectedMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetAutoBoostedClocksEnabled(handle): + c_isEnabled = _nvmlEnableState_t() + c_defaultIsEnabled = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") + ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) + _nvmlCheckReturn(ret) + return [c_isEnabled.value, c_defaultIsEnabled.value] + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +## Set functions +def nvmlUnitSetLedState(unit, color): + fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") + ret = fn(unit, _nvmlLedColor_t(color)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetPersistenceMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetComputeMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") + ret = fn(handle, _nvmlComputeMode_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetEccMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearEccErrorCounts(handle, counterType): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") + ret = fn(handle, _nvmlEccCounterType_t(counterType)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetDriverModel(handle, model): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") + ret = fn(handle, _nvmlDriverModel_t(model)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled)) + _nvmlCheckReturn(ret) + return None + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) + _nvmlCheckReturn(ret) + return None + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") + ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceResetGpuLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") + ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceResetMemoryLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): + isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t + c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClkMonStatus") + ret = fn(handle, c_clkMonInfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_clkMonInfo + +# Added in 4.304 +def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") + ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +def nvmlDeviceResetApplicationsClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +def nvmlDeviceSetPowerManagementLimit(handle, limit): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") + ret = fn(handle, c_uint(limit)) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +def nvmlDeviceSetGpuOperationMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") + ret = fn(handle, _nvmlGpuOperationMode_t(mode)) + _nvmlCheckReturn(ret) + return None + +# Added in 2.285 +def nvmlEventSetCreate(): + fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") + eventSet = c_nvmlEventSet_t() + ret = fn(byref(eventSet)) + _nvmlCheckReturn(ret) + return eventSet + +# Added in 2.285 +def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): + fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") + ret = fn(handle, c_ulonglong(eventTypes), eventSet) + _nvmlCheckReturn(ret) + return None + +# Added in 2.285 +def nvmlDeviceGetSupportedEventTypes(handle): + c_eventTypes = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") + ret = fn(handle, byref(c_eventTypes)) + _nvmlCheckReturn(ret) + return c_eventTypes.value + +# raises NVML_ERROR_TIMEOUT exception on timeout +def nvmlEventSetWait_v2(eventSet, timeoutms): + fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") + data = c_nvmlEventData_t() + ret = fn(eventSet, byref(data), c_uint(timeoutms)) + _nvmlCheckReturn(ret) + return data + +def nvmlEventSetWait(eventSet, timeoutms): + return nvmlEventSetWait_v2(eventSet, timeoutms) + +# Added in 2.285 +def nvmlEventSetFree(eventSet): + fn = _nvmlGetFunctionPointer("nvmlEventSetFree") + ret = fn(eventSet) + _nvmlCheckReturn(ret) + return None + +# Added in 3.295 +def nvmlDeviceOnSameBoard(handle1, handle2): + fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") + onSameBoard = c_int() + ret = fn(handle1, handle2, byref(onSameBoard)) + _nvmlCheckReturn(ret) + return (onSameBoard.value != 0) + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + +def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 4.304 +def nvmlDeviceGetSupportedClocksThrottleReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +def nvmlDeviceGetSupportedClocksEventReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +# Added in 4.304 +def nvmlDeviceGetCurrentClocksThrottleReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +def nvmlDeviceGetCurrentClocksEventReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +# Added in 5.319 +def nvmlDeviceGetIndex(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") + c_index = c_uint() + ret = fn(handle, byref(c_index)) + _nvmlCheckReturn(ret) + return c_index.value + +# Added in 5.319 +def nvmlDeviceGetAccountingMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceSetAccountingMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearAccountingPids(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetAccountingStats(handle, pid): + stats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") + ret = fn(handle, c_uint(pid), byref(stats)) + _nvmlCheckReturn(ret) + if (stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + stats.maxMemoryUsage = None + return stats + +def nvmlDeviceGetAccountingPids(handle): + count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) + pids = (c_uint * count.value)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") + ret = fn(handle, byref(count), pids) + _nvmlCheckReturn(ret) + return list(map(int, pids[0:count.value])) + +def nvmlDeviceGetAccountingBufferSize(handle): + bufferSize = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") + ret = fn(handle, byref(bufferSize)) + _nvmlCheckReturn(ret) + return int(bufferSize.value) + +def nvmlDeviceGetRetiredPages(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + ret = fn(device, c_source, byref(c_count), c_pages) + _nvmlCheckReturn(ret) + return list(map(int, c_pages[0:c_count.value])) + +def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages_v2") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + times_array = c_ulonglong * c_count.value + c_times = times_array() + ret = fn(device, c_source, byref(c_count), c_pages, c_times) + _nvmlCheckReturn(ret) + return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ]; + +def nvmlDeviceGetRetiredPagesPendingStatus(device): + c_pending = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") + ret = fn(device, byref(c_pending)) + _nvmlCheckReturn(ret) + return int(c_pending.value) + +def nvmlDeviceGetAPIRestriction(device, apiType): + c_permission = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") + ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) + _nvmlCheckReturn(ret) + return int(c_permission.value) + +def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") + ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetBridgeChipInfo(handle): + bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") + ret = fn(handle, byref(bridgeHierarchy)) + _nvmlCheckReturn(ret) + return bridgeHierarchy + +def nvmlDeviceGetSamples(device, sampling_type, timeStamp): + c_sampling_type = _nvmlSamplingType_t(sampling_type) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_count = c_uint(0) + c_sample_value_type = _nvmlValueType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") + + ## First Call gets the size + ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None) + + # Stop if this fails + if (ret != NVML_SUCCESS): + raise NVMLError(ret) + + sampleArray = c_sample_count.value * c_nvmlSample_t + c_samples = sampleArray() + ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) + _nvmlCheckReturn(ret) + return (c_sample_value_type.value, c_samples[0:c_sample_count.value]) + +def nvmlDeviceGetViolationStatus(device, perfPolicyType): + c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) + c_violTime = c_nvmlViolationTime_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") + + ## Invoke the method to get violation time + ret = fn(device, c_perfPolicy_type, byref(c_violTime)) + _nvmlCheckReturn(ret) + return c_violTime + +def nvmlDeviceGetPcieThroughput(device, counter): + c_util = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") + ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) + _nvmlCheckReturn(ret) + return c_util.value + +def nvmlSystemGetTopologyGpuSet(cpuNumber): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") + + # First call will get the size + ret = fn(cpuNumber, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(cpuNumber, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0:c_count.value]) + +def nvmlDeviceGetTopologyNearestGpus(device, level): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") + + # First call will get the size + ret = fn(device, level, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(device, level, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0:c_count.value]) + +def nvmlDeviceGetTopologyCommonAncestor(device1, device2): + c_level = _nvmlGpuTopologyLevel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") + ret = fn(device1, device2, byref(c_level)) + _nvmlCheckReturn(ret) + return c_level.value + +def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): + c_rxcounter = c_ulonglong() + c_txcounter = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationCounter") + ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter)) + _nvmlCheckReturn(ret) + return (c_rxcounter.value, c_txcounter.value) + +def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): + fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") + ret = fn(device, link, counter, freeze) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") + ret = fn(device, link, counter) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(control), reset) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): + c_control = nvmlNvLinkUtilizationControl_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(c_control)) + _nvmlCheckReturn(ret) + return c_control + +def nvmlDeviceGetNvLinkCapability(device, link, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") + ret = fn(device, link, capability, byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): + c_result = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") + ret = fn(device, link, counter, byref(c_result)) + _nvmlCheckReturn(ret) + return c_result.value + +def nvmlDeviceResetNvLinkErrorCounters(device, link): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") + ret = fn(device, link) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetNvLinkRemotePciInfo(device, link): + c_pci = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") + ret = fn(device, link, byref(c_pci)) + _nvmlCheckReturn(ret) + return c_pci + +def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): + c_type = _nvmlNvLinkDeviceType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") + ret = fn(handle, link, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + +def nvmlDeviceGetNvLinkState(device, link): + c_isActive = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") + ret = fn(device, link, byref(c_isActive)) + _nvmlCheckReturn(ret) + return c_isActive.value + +def nvmlDeviceGetNvLinkVersion(device, link): + c_version = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") + ret = fn(device, link, byref(c_version)) + _nvmlCheckReturn(ret) + return c_version.value + +def nvmlDeviceModifyDrainState(pciInfo, newState): + fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") + ret = fn(pointer(pciInfo), newState) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceQueryDrainState(pciInfo): + c_newState = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") + ret = fn(pointer(pciInfo), byref(c_newState)) + _nvmlCheckReturn(ret) + return c_newState.value + +def nvmlDeviceRemoveGpu(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceDiscoverGpus(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + +def nvmlDeviceClearFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceClearFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + +def nvmlDeviceGetVirtualizationMode(handle): + c_virtualization_mode = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") + ret = fn(handle, byref(c_virtualization_mode)) + _nvmlCheckReturn(ret) + return c_virtualization_mode.value + +def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") + return fn(handle, virtualization_mode) + +def nvmlDeviceGetVgpuHeterogeneousMode(handle): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return c_vgpuHeterogeneousMode.mode + +def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + c_vgpuHeterogeneousMode.mode = heterogeneous_mode + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlVgpuInstanceGetPlacementId(vgpuInstance): + c_placement = c_nvmlVgpuPlacementId_v1_t(0) + c_placement.version = VgpuPlacementId_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetPlacementId") + ret = fn(vgpuInstance, byref(c_placement)) + _nvmlCheckReturn(ret) + return c_placement.placementId + +def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + c_vgpu_placements.mode = mode + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + +def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + +def nvmlGetVgpuDriverCapabilities(capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") + ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceGetVgpuCapabilities(handle, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceSetVgpuCapabilities(handle, capability, state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetSupportedVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no supported vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetCreatableVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no supported vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): + c_profile_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + ret = fn(vgpuTypeId, byref(c_profile_id)) + _nvmlCheckReturn(ret) + return (c_profile_id.value) + +@convertStrBytes +def nvmlVgpuTypeGetClass(vgpuTypeId): + c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_class.value + +@convertStrBytes +def nvmlVgpuTypeGetName(vgpuTypeId): + c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_name.value + +def nvmlVgpuTypeGetDeviceID(vgpuTypeId): + c_device_id = c_ulonglong(0) + c_subsystem_id = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) + _nvmlCheckReturn(ret) + return (c_device_id.value, c_subsystem_id.value) + +def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): + c_fb_size = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + ret = fn(vgpuTypeId, byref(c_fb_size)) + _nvmlCheckReturn(ret) + return c_fb_size.value + +def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): + c_num_heads = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + ret = fn(vgpuTypeId, byref(c_num_heads)) + _nvmlCheckReturn(ret) + return c_num_heads.value + +def nvmlVgpuTypeGetResolution(vgpuTypeId): + c_xdim = c_uint(0) + c_ydim = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) + _nvmlCheckReturn(ret) + return (c_xdim.value, c_ydim.value) + +@convertStrBytes +def nvmlVgpuTypeGetLicense(vgpuTypeId): + c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + ret = fn(vgpuTypeId, c_license, c_buffer_size) + _nvmlCheckReturn(ret) + return c_license.value + +def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): + c_frl_config = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + ret = fn(vgpuTypeId, byref(c_frl_config)) + _nvmlCheckReturn(ret) + return c_frl_config.value + +def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): + c_gsp_heap = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + ret = fn(vgpuTypeId, byref(c_gsp_heap)) + _nvmlCheckReturn(ret) + return c_gsp_heap.value + +def nvmlVgpuTypeGetFbReservation(vgpuTypeId): + c_fb_reservation = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + ret = fn(vgpuTypeId, byref(c_fb_reservation)) + _nvmlCheckReturn(ret) + return c_fb_reservation.value + +def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): + c_runtime_state = nvmlVgpuRuntimeState_v1_t() + c_runtime_state.version = VgpuRuntimeState_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + ret = fn(vgpuInstance, byref(c_runtime_state)) + _nvmlCheckReturn(ret) + return c_runtime_state + +def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + return c_max_instances.value + +def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): + c_max_instances_per_vm = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) + _nvmlCheckReturn(ret) + return c_max_instances_per_vm.value + +def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): + c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) + c_bar1Info.version = VgpuTypeBar1Info_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + ret = fn(vgpuTypeId, byref(c_bar1Info)) + _nvmlCheckReturn(ret) + return c_bar1Info + +def nvmlDeviceGetActiveVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value + c_vgpu_instances = vgpu_instance_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_instances[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +@convertStrBytes +def nvmlVgpuInstanceGetVmID(vgpuInstance): + c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) + _nvmlCheckReturn(ret) + return (c_vm_id.value, c_vm_id_type.value) + +@convertStrBytes +def nvmlVgpuInstanceGetUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): + c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) + _nvmlCheckReturn(ret) + return c_driver_version.value + +def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): + c_license_status = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + ret = fn(vgpuInstance, byref(c_license_status)) + _nvmlCheckReturn(ret) + return c_license_status.value + +def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + c_license_info = c_nvmlVgpuLicenseInfo_t() + ret = fn(vgpuInstance, byref(c_license_info)) + _nvmlCheckReturn(ret) + return c_license_info + +def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): + return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + +def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): + c_frl = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + ret = fn(vgpuInstance, byref(c_frl)) + _nvmlCheckReturn(ret) + return c_frl.value + +def nvmlVgpuInstanceGetEccMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlVgpuInstanceGetType(vgpuInstance): + c_vgpu_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + ret = fn(vgpuInstance, byref(c_vgpu_type)) + _nvmlCheckReturn(ret) + return c_vgpu_type.value + +def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): + c_encoder_capacity = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + ret = fn(vgpuInstance, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + +def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + return fn(vgpuInstance, encoder_capacity) + +def nvmlVgpuInstanceGetFbUsage(vgpuInstance): + c_fb_usage = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + ret = fn(vgpuInstance, byref(c_fb_usage)) + _nvmlCheckReturn(ret) + return c_fb_usage.value + +def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): + c_cap_result = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) + _nvmlCheckReturn(ret) + return (c_cap_result.value) + +def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): + c_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + ret = fn(vgpuInstance, byref(c_id)) + _nvmlCheckReturn(ret) + return (c_id.value) + +@convertStrBytes +def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): + c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") + ret = fn(vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE))) + _nvmlCheckReturn(ret) + return c_vgpuPciId.value + +def nvmlDeviceGetVgpuUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_value_type = _nvmlValueType_t() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0) + c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1 + c_vgpuUtilInfo.sampleValType = _nvmlValueType_t() + c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) + c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + ret = fn(handle, byref(c_vgpuUtilInfo)) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t + c_samples = sampleArray() + c_vgpuUtilInfo.vgpuUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpuUtilInfo.vgpuInstanceCount] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): + c_p2pstatus = _nvmlGpuP2PStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") + ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus)) + _nvmlCheckReturn(ret) + return c_p2pstatus.value + +def nvmlDeviceGetGridLicensableFeatures_v4(handle): + c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") + ret = fn(handle, byref(c_get_grid_licensable_features)) + _nvmlCheckReturn(ret) + + return (c_get_grid_licensable_features) + +def nvmlDeviceGetGridLicensableFeatures(handle): + return nvmlDeviceGetGridLicensableFeatures_v4(handle) + +def nvmlDeviceGetGspFirmwareVersion(handle, version=None): + isUserDefined = version is not None + if not isUserDefined: + version = (c_char * NVML_GSP_FIRMWARE_VERSION_BUF_SIZE)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareVersion") + ret = fn(handle, version) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else version.value + +def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): + isReference = type(isEnabled) is not c_uint + isEnabledRef = isEnabled if isReference else byref(isEnabled) + defaultModeRef = defaultMode if isReference else byref(defaultMode) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareMode") + ret = fn(handle, isEnabledRef, defaultModeRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + +def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): + c_encoder_capacity = c_ulonglong(0) + c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderCapacity") + ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + +def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0) + c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1 + c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) + c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t + c_samples = sampleArray() + c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpuProcUtilInfo.vgpuProcessCount] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetEncoderStats(handle): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderStats") + ret = fn(handle, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + +def nvmlDeviceGetEncoderSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + ret = fn(handle, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetFBCStats(handle): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") + ret = fn(handle, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + +def nvmlDeviceGetFBCSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + ret = fn(handle, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + +def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetFBCStats(vgpuInstance): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") + ret = fn(vgpuInstance, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + +def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetProcessUtilization(handle, timeStamp): + # first call to get the size + c_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + ret = fn(handle, None, byref(c_count), c_time_stamp) + + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_samples, byref(c_count), c_time_stamp) + _nvmlCheckReturn(ret) + + return c_samples[0:c_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0) + c_processesUtilInfo.version = ProcessesUtilizationInfo_v1 + c_processesUtilInfo.processSamplesCount = c_uint(0) + c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + ret = fn(handle, byref(c_processesUtilInfo)) + + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t + c_samples = sampleArray() + c_processesUtilInfo.procUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_processesUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_processesUtilInfo.processSamplesCount] + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetMetadata(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") + c_vgpuMetadata = c_nvmlVgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuMetadata + +def nvmlDeviceGetVgpuMetadata(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") + c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuPgpuMetadata + +def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): + fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") + c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() + ret = fn(byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility)) + _nvmlCheckReturn(ret) + return c_vgpuPgpuCompatibility + +@convertStrBytes +def nvmlDeviceGetPgpuMetadataString(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") + c_pgpuMetadata = create_string_buffer(NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pgpuMetadata.value, c_bufferSize.value) + +def nvmlDeviceGetVgpuSchedulerLog(handle): + c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") + ret = fn(handle, byref(c_vgpu_sched_log)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_log + +def nvmlDeviceGetVgpuSchedulerState(handle): + c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") + ret = fn(handle, byref(c_vgpu_sched_state)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_state + +def nvmlDeviceGetVgpuSchedulerCapabilities(handle): + c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") + ret = fn(handle, byref(c_vgpu_sched_caps)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_caps + +def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") + ret = fn(handle, byref(sched_state)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSetVgpuVersion(vgpuVersion): + fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") + ret = fn(byref(vgpuVersion)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGetVgpuVersion(supported=None, current=None): + isUserDefined = (supported is not None) or (current is not None) + if not isUserDefined: + supported = c_nvmlVgpuVersion_t() + current = c_nvmlVgpuVersion_t() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") + ret = fn(byref(supported), byref(current)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else [(supported.minVersion, + supported.maxVersion), + (current.minVersion, + current.maxVersion)] + +def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): + c_pidCount = c_uint() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") + ret = fn(vgpuInstance, byref(c_pidCount), None) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + sampleArray = c_pidCount.value * c_uint + c_pidArray = sampleArray() + ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pidCount, c_pidArray) + +def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): + c_accountingStats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") + ret = fn(vgpuInstance, pid, byref(c_accountingStats)) + _nvmlCheckReturn(ret) + return c_accountingStats + +def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") + ret = fn(vgpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGetExcludedDeviceCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlGetExcludedDeviceInfoByIndex(index): + c_index = c_uint(index) + info = c_nvmlExcludedDeviceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceInfoByIndex") + ret = fn(c_index, byref(info)) + _nvmlCheckReturn(ret) + return info + +def nvmlDeviceGetHostVgpuMode(handle): + c_host_vgpu_mode = _nvmlHostVgpuMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") + ret = fn(handle, byref(c_host_vgpu_mode)) + _nvmlCheckReturn(ret) + return c_host_vgpu_mode.value + +def nvmlDeviceSetMigMode(device, mode): + c_activationStatus = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") + ret = fn(device, mode, byref(c_activationStatus)) + _nvmlCheckReturn(ret) + return c_activationStatus.value + +def nvmlDeviceGetMigMode(device): + c_currentMode = c_uint() + c_pendingMode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigMode") + ret = fn(device, byref(c_currentMode), byref(c_pendingMode)) + _nvmlCheckReturn(ret) + return [c_currentMode.value, c_pendingMode.value] + +def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): + if version == 2: + c_info = c_nvmlGpuInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlGpuInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +# Define function alias for the API exposed by NVML +nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo + +def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") + ret = fn(device, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") + ret = fn(device, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceCreateGpuInstance(device, profileId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") + ret = fn(device, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") + ret = fn(device, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceDestroy(gpuInstance): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") + ret = fn(gpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") + ret = fn(device, profileId, gpuInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") + ret = fn(device, gpuInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceGetInfo(gpuInstance): + c_info = c_nvmlGpuInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") + ret = fn(gpuInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2): + if version == 2: + c_info = c_nvmlComputeInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlComputeInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, engProfile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +# Define function alias for the API exposed by NVML +nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo + +def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceRemainingCapacity") + ret = fn(gpuInstance, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") + ret = fn(gpuInstance, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") + ret = fn(gpuInstance, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") + ret = fn(gpuInstance, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlComputeInstanceDestroy(computeInstance): + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") + ret = fn(computeInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") + ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") + ret = fn(gpuInstance, computeInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlComputeInstanceGetInfo_v2(computeInstance): + c_info = c_nvmlComputeInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") + ret = fn(computeInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlComputeInstanceGetInfo(computeInstance): + return nvmlComputeInstanceGetInfo_v2(computeInstance) + +def nvmlDeviceIsMigDeviceHandle(device): + c_isMigDevice = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") + ret = fn(device, byref(c_isMigDevice)) + _nvmlCheckReturn(ret) + return c_isMigDevice + +def nvmlDeviceGetGpuInstanceId(device): + c_gpuInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") + ret = fn(device, byref(c_gpuInstanceId)) + _nvmlCheckReturn(ret) + return c_gpuInstanceId.value + +def nvmlDeviceGetComputeInstanceId(device): + c_computeInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") + ret = fn(device, byref(c_computeInstanceId)) + _nvmlCheckReturn(ret) + return c_computeInstanceId.value + +def nvmlDeviceGetMaxMigDeviceCount(device): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") + ret = fn(device, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetMigDeviceHandleByIndex(device, index): + c_index = c_uint(index) + migDevice = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigDeviceHandleByIndex") + ret = fn(device, c_index, byref(migDevice)) + _nvmlCheckReturn(ret) + return migDevice + +def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") + ret = fn(migDevice, byref(device)) + _nvmlCheckReturn(ret) + return device + +def nvmlDeviceGetAttributes_v2(device): + c_attrs = c_nvmlDeviceAttributes() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") + ret = fn(device, byref(c_attrs)) + _nvmlCheckReturn(ret) + return c_attrs + +def nvmlDeviceGetAttributes(device): + return nvmlDeviceGetAttributes_v2(device) + +def nvmlDeviceGetRemappedRows(device): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") + c_corr = c_uint() + c_unc = c_uint() + c_bpending = c_uint() + c_bfailure = c_uint() + ret = fn(device, byref(c_corr), byref(c_unc), byref(c_bpending), byref(c_bfailure)) + _nvmlCheckReturn(ret) + return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + +def nvmlDeviceGetRowRemapperHistogram(device): + c_vals = c_nvmlRowRemapperHistogramValues() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") + ret = fn(device, byref(c_vals)) + _nvmlCheckReturn(ret) + return c_vals + +def nvmlDeviceGetArchitecture(device): + arch = _nvmlDeviceArchitecture_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") + ret = fn(device, byref(arch)) + _nvmlCheckReturn(ret) + return arch.value + +def nvmlDeviceGetBusType(device): + c_busType = _nvmlBusType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") + ret = fn(device, byref(c_busType)) + _nvmlCheckReturn(ret) + return c_busType.value + +def nvmlDeviceGetIrqNum(device): + c_irqNum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") + ret = fn(device, byref(c_irqNum)) + _nvmlCheckReturn(ret) + return c_irqNum.value + +def nvmlDeviceGetNumGpuCores(device): + c_numCores = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") + ret = fn(device, byref(c_numCores)) + _nvmlCheckReturn(ret) + return c_numCores.value + +def nvmlDeviceGetPowerSource(device): + c_powerSource = _nvmlPowerSource_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") + ret = fn(device, byref(c_powerSource)) + _nvmlCheckReturn(ret) + return c_powerSource.value + +def nvmlDeviceGetMemoryBusWidth(device): + c_memBusWidth = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") + ret = fn(device, byref(c_memBusWidth)) + _nvmlCheckReturn(ret) + return c_memBusWidth.value + +def nvmlDeviceGetPcieLinkMaxSpeed(device): + c_speed = _nvmlPcieLinkMaxSpeed_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetAdaptiveClockInfoStatus(device): + c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") + ret = fn(device, byref(c_adaptiveClockInfoStatus)) + _nvmlCheckReturn(ret) + return c_adaptiveClockInfoStatus.value + +def nvmlDeviceGetPcieSpeed(device): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()): + isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t + dynamicpstatesinfoRef = c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo"); + ret = fn(device, dynamicpstatesinfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + +def nvmlDeviceSetFanSpeed_v2(handle, index, speed): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2"); + ret = fn(handle, index, speed) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()): + isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t + thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings"); + ret = fn(device, sensorindex, thermalsettingsRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] + +def nvmlDeviceGetMinMaxClockOfPState(device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()): + isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) + minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) + maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState"); + ret = fn(device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + +class c_nvmlClockOffset_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('type', _nvmlClockType_t), + ('pstate', _nvmlPstates_t), + ('clockOffsetMHz', c_int), + ('minClockOffsetMHz', c_int), + ('maxClockOffsetMHz', c_int), + ] + +nvmlClockOffset_v1 = 0x1000018 + +def nvmlDeviceGetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets"); + ret = fn(device, info) + return NVML_SUCCESS + +def nvmlDeviceSetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets"); + ret = fn(device, info) + return NVML_SUCCESS + +def nvmlDeviceGetSupportedPerformanceStates(device): + pstates = [] + c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) + c_size = sizeof(c_uint)*c_count.value + + # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. + pstates_array = _nvmlPstates_t * c_count.value + c_pstates = pstates_array() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedPerformanceStates") + ret = fn(device, c_pstates, c_size) + _nvmlCheckReturn(ret) + + for value in c_pstates: + if value != NVML_PSTATE_UNKNOWN: + pstates.append(value) + + return pstates + +def nvmlDeviceGetGpcClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + +def nvmlDeviceSetGpcClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + +def nvmlDeviceGetMemClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + +def nvmlDeviceSetMemClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + +def nvmlSystemSetConfComputeGpusReadyState(state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") + ret = fn(c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetConfComputeGpusReadyState(): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +def nvmlSystemGetConfComputeCapabilities(): + c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") + ret = fn(byref(c_ccSysCaps)) + _nvmlCheckReturn(ret) + return c_ccSysCaps + +def nvmlSystemGetConfComputeState(): + c_state = c_nvmlConfComputeSystemState_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + +def nvmlSystemGetConfComputeSettings(settings): + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") + return fn(settings) + +def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") + ret = fn(device, c_ccMemSize) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetConfComputeMemSizeInfo(device): + c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") + ret = fn(device, byref(c_ccMemSize)) + _nvmlCheckReturn(ret) + return c_ccMemSize + +def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") + ret = fn(device, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + +def nvmlDeviceGetConfComputeGpuCertificate(device): + c_cert = c_nvmlConfComputeGpuCertificate_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") + ret = fn(device, byref(c_cert)) + _nvmlCheckReturn(ret) + return c_cert + +def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): + c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() + c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) + setattr(c_attestReport, 'nonce', c_nonce_arr) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") + ret = fn(device, byref(c_attestReport)) + _nvmlCheckReturn(ret) + return c_attestReport + +def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): + c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 + c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): + c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return c_keyRotationThrInfo + +## GPM ## +######### + +## Enums/defines + +#### GPM Metric Identifiers +NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DRAM_BW_UTIL = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = 63 # NvLink write bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = 65 # NvLink write bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = 67 # NvLink write bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = 69 # NvLink write bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = 71 # NvLink write bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = 73 # NvLink write bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = 75 # NvLink write bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = 77 # NvLink write bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = 79 # NvLink write bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = 81 # NvLink write bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = 82 # NvLink read bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = 83 # NvLink write bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = 84 # NvLink read bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = 85 # NvLink write bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = 86 # NvLink read bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = 87 # NvLink write bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = 88 # NvLink read bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = 89 # NvLink write bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = 90 # NvLink read bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = 91 # NvLink write bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = 92 # NvLink read bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = 93 # NvLink write bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = 94 # NvLink read bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = 95 # NvLink write bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = 96 # NvLink read bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = 97 # NvLink write bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_MAX = 98 + +## Structs + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ('name', c_char * 96), + ('id', c_char * 96), + ('serial', c_char * 96), + ('firmwareVersion', c_char * 96), + ] + +class struct_c_nvmlGpmSample_t(Structure): + pass # opaque handle +c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + +class c_metricInfo_t(Structure): + _fields_ = [ + ("shortName", c_char_p), + ("longName", c_char_p), + ("unit", c_char_p), + ] + +class c_nvmlGpmMetric_t(_PrintableStructure): + _fields_ = [ + ('metricId', c_uint), + ('nvmlReturn', _nvmlReturn_t), + ('value', c_double), + ('metricInfo', c_metricInfo_t) + ] + +class c_nvmlGpmMetricsGet_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('numMetrics', c_uint), + ('sample1', c_nvmlGpmSample_t), + ('sample2', c_nvmlGpmSample_t), + ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX) + ] + +NVML_GPM_METRICS_GET_VERSION = 1 + +class c_nvmlGpmSupport_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('isSupportedDevice', c_uint), + ] + +NVML_GPM_SUPPORT_VERSION = 1 + +## Functions + +def nvmlGpmMetricsGet(metricsGet): + fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") + ret = fn(byref(metricsGet)) + _nvmlCheckReturn(ret) + return metricsGet + +def nvmlGpmSampleFree(gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") + ret = fn(gpmSample) + _nvmlCheckReturn(ret) + return + +def nvmlGpmSampleAlloc(): + gpmSample = c_nvmlGpmSample_t() + fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") + ret = fn(byref(gpmSample)) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmSampleGet(device, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") + ret = fn(device, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") + ret = fn(device, gpuInstanceId, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmQueryDeviceSupport(device): + gpmSupport = c_nvmlGpmSupport_t() + gpmSupport.version = NVML_GPM_SUPPORT_VERSION + fn = _nvmlGetFunctionPointer("nvmlGpmQueryDeviceSupport") + ret = fn(device, byref(gpmSupport)) + _nvmlCheckReturn(ret) + return gpmSupport + +def nvmlGpmSetStreamingEnabled(device, state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") + ret = fn(device, c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpmQueryIfStreamingEnabled(device): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") + ret = fn(device, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +# Low Power Structure and Function + +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 + +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + +class c_nvmlNvLinkPowerThres_t(Structure): + _fields_ = [ + ("lowPwrThreshold", c_uint), + ] + +def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): + c_info = c_nvmlNvLinkPowerThres_t() + c_info.lowPwrThreshold = l1threshold + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkDeviceLowPowerThreshold") + ret = fn(device, byref(c_info)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +NVML_GPU_FABRIC_UUID_LEN = 16 + +_nvmlGpuFabricState_t = c_uint +NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + +class c_nvmlGpuFabricInfo_t(_PrintableStructure): + _fields_ = [ + ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t) + ] + +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 + +nvmlGpuFabricInfo_v2 = 0x02000024 + +class c_nvmlGpuFabricInfoV_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ("healthMask", c_uint32) + ] + + def __init__(self): + super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2) + +def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo"); + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV"); + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +###################### +## Enums/defines +#### NVML GPU NVLINK BW MODE +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + +def nvmlSystemSetNvlinkBwMode(mode): + fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") + ret = fn(mode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetNvlinkBwMode(): + mode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") + ret = fn(byref(mode)) + _nvmlCheckReturn(ret) + return mode.value + +_nvmlPowerScopeType_t = c_uint +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + +class c_nvmlPowerValue_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('powerScope', _nvmlPowerScopeType_t), + ('powerValueMw', c_uint), + ] + _fmt_ = {'': "%d B"} + +nvmlPowerValue_v2 = 0x0200000C + +def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2): + c_powerScope = _nvmlPowerScopeType_t(powerScope) + c_powerValue = c_nvmlPowerValue_v2_t() + c_powerValue.version = c_uint(version) + c_powerValue.powerScope = c_powerScope + c_powerValue.powerValueMw = c_uint(powerLimit) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit_v2") + ret = fn(device, byref(c_powerValue)) + return NVML_SUCCESS + +class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('aggregateUncParity', c_ulonglong), + ('aggregateUncSecDed', c_ulonglong), + ('aggregateCor', c_ulonglong), + ('volatileUncParity', c_ulonglong), + ('volatileUncSecDed', c_ulonglong), + ('volatileCor', c_ulonglong), + ('aggregateUncBucketL2', c_ulonglong), + ('aggregateUncBucketSm', c_ulonglong), + ('aggregateUncBucketPcie', c_ulonglong), + ('aggregateUncBucketMcu', c_ulonglong), + ('aggregateUncBucketOther', c_ulonglong), + ('bThresholdExceeded', c_uint) + ] + + def __init__(self): + super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1) + +nvmlEccSramErrorStatus_v1 = 0x1000068 +def nvmlDeviceGetSramEccErrorStatus(device, status): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") + ret = fn(device, status) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +NVML_DEV_CAP_EGM = (1 << 0) +nvmlDeviceCapabilities_v1 = 0x1000008 + +class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('capMask', c_uint), + ] + + def __init__(self): + super(c_nvmlDeviceCapabilities_v1_t, self).__init__(version=nvmlDeviceCapabilities_v1) + + +def nvmlDeviceGetCapabilities(device, caps): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") + return fn(device, caps) + +class c_nvmlPlatformInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('ibGuid', c_char * 16), + ('rackGuid', c_char * 16), + ('chassisPhysicalSlotNumber', c_char), + ('computeSlotIndex', c_char), + ('nodeIndex', c_char), + ('peerType', c_char), + ('moduleId', c_char) + ] + + def __init__(self): + super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) + +nvmlPlatformInfo_v1 = 0x100002c +def nvmlDeviceGetPlatformInfo(device, platformInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") + ret = fn(device, platformInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +class c_nvmlMask255_t(_PrintableStructure): + _fields_ = [ + ('mask', c_uint * 8), + ] + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 + +nvmlWorkloadPowerProfileInfo_v1 = 0x100002c +class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('profileId', c_uint), + ('priority', c_uint), + ('conflictingmask', c_nvmlMask255_t) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileInfo_v1) + +nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002bf8 +class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('perfProfilesMask', c_nvmlMask255_t), + ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1) + +nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 +class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('perfProfilesMask', c_nvmlMask255_t), + ('requestedProfilesMask', c_nvmlMask255_t), + ('enforcedProfilesMask', c_nvmlMask255_t) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1) + +nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 +class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('requestedProfilesMask', c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1) + +def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") + ret = fn(device, profilesInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") + ret = fn(device, currentProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") + ret = fn(device, supportedBwModes) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetNvlinkBwMode(device, getBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") + ret = fn(device, getBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceSetNvlinkBwMode(device, setBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") + ret = fn(device, setBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +nvmlDramEncryptionInfo_v1 = 0x01000008 + +class c_nvmlDramEncryptionInfo_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('encryptionState', _nvmlEnableState_t), + ] + + def __init__(self): + super(c_nvmlDramEncryptionInfo_t, self).__init__(version=nvmlDramEncryptionInfo_v1) + +def nvmlDeviceGetDramEncryptionMode(handle): + c_currState = c_nvmlDramEncryptionInfo_t() + c_pendingState = c_nvmlDramEncryptionInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDramEncryptionMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.encryptionState, c_pendingState.encryptionState] + +# added to API +def nvmlDeviceGetCurrentDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[0] + +# added to API +def nvmlDeviceGetPendingDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[1] + +def nvmlDeviceSetDramEncryptionMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") + c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() + c_dramEncryptionMode.encryptionState = mode; + ret = fn(handle, byref(c_dramEncryptionMode)) + _nvmlCheckReturn(ret) + return None + +# Power Smoothing defines +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 + +nvmlPowerSmoothingState_v1=0x1000008 +class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('state', c_uint), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingState_v1_t, self).__init__(version=nvmlPowerSmoothingState_v1) + +nvmlPowerSmoothingProfile_v1=0x1000018 +class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('profileId', c_uint), + ('paramId', c_uint), + ('value', c_double), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(version=nvmlPowerSmoothingProfile_v1) + +def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + +def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + +def nvmlDevicePowerSmoothingSetState(device, state): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") + ret = fn(device, state) + _nvmlCheckReturn(ret) + diff --git a/trainer/trainer_args.py b/trainer/trainer_args.py new file mode 100644 index 0000000000000000000000000000000000000000..0d8876cd19b6eecf78ea51e082991128bce37dfb --- /dev/null +++ b/trainer/trainer_args.py @@ -0,0 +1,1097 @@ +# SPDX-License-Identifier: Apache-2.0 +# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py +"""The arguments of Trainer Inference.""" + +import argparse +import dataclasses +from contextlib import contextmanager +from dataclasses import field +from enum import Enum +from typing import Any + +from trainer.configs.configs import PreprocessConfig +from trainer.configs.pipelines.base import PipelineConfig, STA_Mode +from trainer.configs.utils import clean_cli_args +from trainer.logger import init_logger +from trainer.platforms import current_platform +from trainer.utils import FlexibleArgumentParser, StoreBoolean + +logger = init_logger(__name__) + + +class ExecutionMode(str, Enum): + """ + Enumeration for different pipeline modes. + + Inherits from str to allow string comparison for backward compatibility. + """ + INFERENCE = "inference" + PREPROCESS = "preprocess" + FINETUNING = "finetuning" + DISTILLATION = "distillation" + + @classmethod + def from_string(cls, value: str) -> "ExecutionMode": + """Convert string to ExecutionMode enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid mode: {value}. Must be one of: {', '.join([m.value for m in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings for argparse.""" + return [mode.value for mode in cls] + + +class WorkloadType(str, Enum): + """ + Enumeration for different workload types. + + Inherits from str to allow string comparison for backward compatibility. + """ + I2V = "i2v" # Image to Video + T2V = "t2v" # Text to Video + T2I = "t2i" # Text to Image + I2I = "i2i" # Image to Image + + @classmethod + def from_string(cls, value: str) -> "WorkloadType": + """Convert string to WorkloadType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid workload type: {value}. Must be one of: {', '.join([m.value for m in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings for argparse.""" + return [workload.value for workload in cls] + + +# args for trainer framework +@dataclasses.dataclass +class TrainerArgs: + # Model and path configuration (for convenience) + model_path: str + + # Running mode + mode: ExecutionMode = ExecutionMode.INFERENCE + + # Workload type + workload_type: WorkloadType = WorkloadType.T2V + + # Cache strategy + cache_strategy: str = "none" + + # Distributed executor backend + distributed_executor_backend: str = "mp" + + inference_mode: bool = True # if False == training mode + + # HuggingFace specific parameters + trust_remote_code: bool = False + revision: str | None = None + + # Parallelism + num_gpus: int = 1 + tp_size: int = -1 + sp_size: int = -1 + hsdp_replicate_dim: int = 1 + hsdp_shard_dim: int = -1 + dist_timeout: int | None = None # timeout for torch.distributed + + pipeline_config: PipelineConfig = field(default_factory=PipelineConfig) + preprocess_config: PreprocessConfig | None = None + + # LoRA parameters + # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. + lora_path: str | None = None + lora_nickname: str = "default" # for swapping adapters in the pipeline + # can restrict layers to adapt, e.g. ["q_proj"] + # Will adapt only q, k, v, o by default. + lora_target_modules: list[str] | None = None + + output_type: str = "pil" + + # CPU offload parameters + dit_cpu_offload: bool = True + use_fsdp_inference: bool = True + text_encoder_cpu_offload: bool = True + image_encoder_cpu_offload: bool = True + vae_cpu_offload: bool = True + pin_cpu_memory: bool = True + + # STA (Sliding Tile Attention) parameters + mask_strategy_file_path: str | None = None + STA_mode: STA_Mode = STA_Mode.STA_INFERENCE + skip_time_steps: int = 15 + + # Compilation + enable_torch_compile: bool = False + + disable_autocast: bool = False + + # VSA parameters + VSA_sparsity: float = 0.0 # inference/validation sparsity + + # Master port for distributed training/inference + master_port: int | None = None + + # Stage verification + enable_stage_verification: bool = True + + # Prompt text file for batch processing + prompt_txt: str | None = None + + # model paths for correct deallocation + model_paths: dict[str, str] = field(default_factory=dict) + model_loaded: dict[str, bool] = field(default_factory=lambda: { + "transformer": True, + "vae": True, + }) + + cls_name: str | None = None + load_from_dir: str | None = None + module_name: str| None = None + ar_action_load_from_dir: str | None = None + + # # DMD parameters + # dmd_denoising_steps: List[int] | None = field(default=None) + + # MoE parameters used by Wan2.2 + boundary_ratio: float | None = None + + @property + def training_mode(self) -> bool: + return not self.inference_mode + + def __post_init__(self): + self.check_trainer_args() + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + # Model and path configuration + parser.add_argument( + "--model-path", + type=str, + help= + "The path of the model weights. This can be a local folder or a Hugging Face repo ID.", + ) + parser.add_argument( + "--model-dir", + type=str, + help="Directory containing StepVideo model", + ) + + # Running mode + parser.add_argument( + "--mode", + type=str, + choices=ExecutionMode.choices(), + default=TrainerArgs.mode.value, + help="The mode to run Trainer", + ) + + # Workload type + parser.add_argument( + "--workload-type", + type=str, + choices=WorkloadType.choices(), + default=TrainerArgs.workload_type.value, + help="The workload type", + ) + + # distributed_executor_backend + parser.add_argument( + "--distributed-executor-backend", + type=str, + choices=["mp"], + default=TrainerArgs.distributed_executor_backend, + help="The distributed executor backend to use", + ) + + parser.add_argument( + "--inference-mode", + action=StoreBoolean, + default=TrainerArgs.inference_mode, + help="Whether to use inference mode", + ) + + # HuggingFace specific parameters + parser.add_argument( + "--trust-remote-code", + action=StoreBoolean, + default=TrainerArgs.trust_remote_code, + help="Trust remote code when loading HuggingFace models", + ) + parser.add_argument( + "--revision", + type=str, + default=TrainerArgs.revision, + help= + "The specific model version to use (can be a branch name, tag name, or commit id)", + ) + + # Parallelism + parser.add_argument( + "--num-gpus", + type=int, + default=TrainerArgs.num_gpus, + help="The number of GPUs to use.", + ) + parser.add_argument( + "--tp-size", + type=int, + default=TrainerArgs.tp_size, + help="The tensor parallelism size.", + ) + parser.add_argument( + "--sp-size", + type=int, + default=TrainerArgs.sp_size, + help="The sequence parallelism size.", + ) + parser.add_argument( + "--hsdp-replicate-dim", + type=int, + default=TrainerArgs.hsdp_replicate_dim, + help="The data parallelism size.", + ) + parser.add_argument( + "--hsdp-shard-dim", + type=int, + default=TrainerArgs.hsdp_shard_dim, + help="The data parallelism shards.", + ) + parser.add_argument( + "--dist-timeout", + type=int, + default=TrainerArgs.dist_timeout, + help="Set timeout for torch.distributed initialization.", + ) + + # Output type + parser.add_argument( + "--output-type", + type=str, + default=TrainerArgs.output_type, + choices=["pil"], + help="Output type for the generated video", + ) + + # Prompt text file for batch processing + parser.add_argument( + "--prompt-txt", + type=str, + default=TrainerArgs.prompt_txt, + help= + "Path to a text file containing prompts (one per line) for batch processing", + ) + + # STA (Sliding Tile Attention) parameters + parser.add_argument( + "--STA-mode", + type=str, + default=TrainerArgs.STA_mode.value, + choices=[mode.value for mode in STA_Mode], + help= + "STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None", + ) + parser.add_argument( + "--skip-time-steps", + type=int, + default=TrainerArgs.skip_time_steps, + help="Number of time steps to warmup (full attention) for STA", + ) + parser.add_argument( + "--mask-strategy-file-path", + type=str, + help="Path to mask strategy JSON file for STA", + ) + parser.add_argument( + "--enable-torch-compile", + action=StoreBoolean, + default=TrainerArgs.enable_torch_compile, + help="Use torch.compile to speed up DiT inference." + + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)", + ) + + parser.add_argument( + "--dit-cpu-offload", + action=StoreBoolean, + help= + "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.", + ) + parser.add_argument( + "--use-fsdp-inference", + action=StoreBoolean, + help= + "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.", + ) + parser.add_argument( + "--text-encoder-cpu-offload", + action=StoreBoolean, + help= + "Use CPU offload for text encoder. Enable if run out of memory.", + ) + parser.add_argument( + "--image-encoder-cpu-offload", + action=StoreBoolean, + help= + "Use CPU offload for image encoder. Enable if run out of memory.", + ) + parser.add_argument( + "--vae-cpu-offload", + action=StoreBoolean, + help="Use CPU offload for VAE. Enable if run out of memory.", + ) + parser.add_argument( + "--pin-cpu-memory", + action=StoreBoolean, + help= + "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". " + "Should be enabled in almost all cases", + ) + parser.add_argument( + "--disable-autocast", + action=StoreBoolean, + help= + "Disable autocast for denoising loop and vae decoding in pipeline sampling", + ) + + # VSA parameters + parser.add_argument( + "--VSA-sparsity", + type=float, + default=TrainerArgs.VSA_sparsity, + help="Validation sparsity for VSA", + ) + + # Master port for distributed training/inference + parser.add_argument( + "--master-port", + type=int, + default=TrainerArgs.master_port, + help="Master port for distributed training/inference", + ) + + # Stage verification + parser.add_argument( + "--enable-stage-verification", + action=StoreBoolean, + default=TrainerArgs.enable_stage_verification, + help="Enable input/output verification for pipeline stages", + ) + + parser.add_argument( + "--cls-name", + type=str, + help="model package name", + ) + parser.add_argument( + "--load-from-dir", + type=str, + help="ar model checkpoint directory", + ) + parser.add_argument( + "--ar-action-load-from-dir", + type=str, + help="ar action model checkpoint directory", + ) + # Add pipeline configuration arguments + PipelineConfig.add_cli_args(parser) + + # Add preprocessing configuration arguments + PreprocessConfig.add_cli_args(parser) + + return parser + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "TrainerArgs": + provided_args = clean_cli_args(args) + # Get all fields from the dataclass + attrs = [attr.name for attr in dataclasses.fields(cls)] + + # Create a dictionary of attribute values, with defaults for missing attributes + kwargs: dict[str, Any] = {} + for attr in attrs: + if attr == 'pipeline_config': + if "HunyuanTransformer" in provided_args["cls_name"]: continue + pipeline_config = PipelineConfig.from_kwargs(provided_args) + kwargs['pipeline_config'] = pipeline_config + elif attr == 'preprocess_config': + preprocess_config = PreprocessConfig.from_kwargs(provided_args) + kwargs['preprocess_config'] = preprocess_config + elif attr == 'mode': + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, TrainerArgs.mode.value) + kwargs['mode'] = ExecutionMode.from_string( + mode_value) if isinstance(mode_value, str) else mode_value + elif attr == 'workload_type': + # Convert string to WorkloadType enum + workload_type_value = getattr(args, 'workload_type', + TrainerArgs.workload_type.value) + kwargs['workload_type'] = WorkloadType.from_string( + workload_type_value) if isinstance( + workload_type_value, str) else workload_type_value + # Use getattr with default value from the dataclass for potentially missing attributes + else: + # Get the field to check if it has a default_factory + field = dataclasses.fields(cls)[next( + i for i, f in enumerate(dataclasses.fields(cls)) + if f.name == attr)] + if field.default_factory is not dataclasses.MISSING: + # Use the default_factory to create the default value + default_value = field.default_factory() + else: + default_value = getattr(cls, attr, None) + value = getattr(args, attr, default_value) + kwargs[attr] = value # type: ignore + + return cls(**kwargs) # type: ignore + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> "TrainerArgs": + # Convert mode string to enum if necessary + if 'mode' in kwargs and isinstance(kwargs['mode'], str): + kwargs['mode'] = ExecutionMode.from_string(kwargs['mode']) + + # Convert workload_type string to enum if necessary + if 'workload_type' in kwargs and isinstance(kwargs['workload_type'], + str): + kwargs['workload_type'] = WorkloadType.from_string( + kwargs['workload_type']) + + kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs) + kwargs['preprocess_config'] = PreprocessConfig.from_kwargs(kwargs) + return cls(**kwargs) + + def check_trainer_args(self) -> None: + """Validate inference arguments for consistency""" + if current_platform.is_mps(): + self.use_fsdp_inference = False + + # Validate mode and inference_mode consistency + assert isinstance( + self.mode, ExecutionMode + ), f"Mode must be an ExecutionMode enum, got {type(self.mode)}" + assert self.mode in ExecutionMode.choices( + ), f"Invalid execution mode: {self.mode}" + + # Validate workload type + assert isinstance( + self.workload_type, WorkloadType + ), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}" + assert self.workload_type in WorkloadType.choices( + ), f"Invalid workload type: {self.workload_type}" + + if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING + ] and self.inference_mode: + logger.warning( + "Mode is 'training' but inference_mode is True. Setting inference_mode to False." + ) + self.inference_mode = False + elif self.mode in [ExecutionMode.INFERENCE, ExecutionMode.PREPROCESS + ] and not self.inference_mode: + logger.warning( + "Mode is '%s' but inference_mode is False. Setting inference_mode to True.", + self.mode) + self.inference_mode = True + + if not self.inference_mode: + assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training" + assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training" + assert self.sp_size != -1, "sp_size must be set for training" + + if self.tp_size == -1: + self.tp_size = 1 + if self.sp_size == -1: + self.sp_size = self.num_gpus + if self.hsdp_shard_dim == -1: + self.hsdp_shard_dim = self.num_gpus + + assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size" + assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim" + assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim" + + if self.num_gpus < max(self.tp_size, self.sp_size): + self.num_gpus = max(self.tp_size, self.sp_size) + + if self.pipeline_config is None: + raise ValueError("pipeline_config is not set in TrainerArgs") + + self.pipeline_config.check_pipeline_config() + + # Add preprocessing config validation if needed + if self.mode == ExecutionMode.PREPROCESS: + if self.preprocess_config is None: + raise ValueError( + "preprocess_config is not set in TrainerArgs when mode is PREPROCESS" + ) + if self.preprocess_config.model_path == "": + self.preprocess_config.model_path = self.model_path + if not self.pipeline_config.vae_config.load_encoder: + self.pipeline_config.vae_config.load_encoder = True + self.preprocess_config.check_preprocess_config() + + +_current_trainer_args = None + + +def prepare_trainer_args(argv: list[str]) -> TrainerArgs: + """ + Prepare the inference arguments from the command line arguments. + + Args: + argv: The command line arguments. Typically, it should be `sys.argv[1:]` + to ensure compatibility with `parse_args` when no arguments are passed. + + Returns: + The inference arguments. + """ + parser = FlexibleArgumentParser() + TrainerArgs.add_cli_args(parser) + raw_args = parser.parse_args(argv) + trainer_args = TrainerArgs.from_cli_args(raw_args) + global _current_trainer_args + _current_trainer_args = trainer_args + return trainer_args + + +@contextmanager +def set_current_trainer_args(trainer_args: TrainerArgs): + """ + Temporarily set the current trainer config. + Used during model initialization. + We save the current trainer config in a global variable, + so that all modules can access it, e.g. custom ops + can access the trainer config to determine how to dispatch. + """ + global _current_trainer_args + old_trainer_args = _current_trainer_args + try: + _current_trainer_args = trainer_args + yield + finally: + _current_trainer_args = old_trainer_args + + +def get_current_trainer_args() -> TrainerArgs: + if _current_trainer_args is None: + # in ci, usually when we test custom ops/modules directly, + # we don't set the trainer config. In that case, we set a default + # config. + # TODO(will): may need to handle this for CI. + raise ValueError("Current trainer args is not set.") + return _current_trainer_args + + +@dataclasses.dataclass +class TrainingArgs(TrainerArgs): + """ + Training arguments. Inherits from TrainerArgs and adds training-specific + arguments. If there are any conflicts, the training arguments will take + precedence. + """ + data_path: str = "" + dataloader_num_workers: int = 0 + num_height: int = 0 + num_width: int = 0 + num_frames: int = 0 + + train_batch_size: int = 0 + num_latent_t: int = 0 + group_frame: bool = False + group_resolution: bool = False + + # text encoder & vae & diffusion model + pretrained_model_name_or_path: str = "" + dit_model_name_or_path: str = "" + + # diffusion setting + ema_decay: float = 0.0 + ema_start_step: int = 0 + training_cfg_rate: float = 0.0 + precondition_outputs: bool = False + + # validation & logs + validation_dataset_file: str = "" + validation_preprocessed_path: str = "" + validation_sampling_steps: str = "" + validation_guidance_scale: str = "" + validation_steps: float = 0.0 + log_validation: bool = False + tracker_project_name: str = "" + wandb_run_name: str = "" + seed: int | None = None + + # output + output_dir: str = "" + checkpoints_total_limit: int = 0 + checkpointing_steps: int = 0 + resume_from_checkpoint: str = "" # specify the checkpoint folder to resume from + + # optimizer & scheduler + num_train_epochs: int = 0 + max_train_steps: int = 0 + gradient_accumulation_steps: int = 0 + learning_rate: float = 0.0 + scale_lr: bool = False + lr_scheduler: str = "constant" + lr_warmup_steps: int = 0 + max_grad_norm: float = 0.0 + enable_gradient_checkpointing_type: str | None = None + selective_checkpointing: float = 0.0 + mixed_precision: str = "" + train_sp_batch_size: int = 0 + fsdp_sharding_startegy: str = "" + + weighting_scheme: str = "" + logit_mean: float = 0.0 + logit_std: float = 1.0 + mode_scale: float = 0.0 + + num_euler_timesteps: int = 0 + lr_num_cycles: int = 0 + lr_power: float = 0.0 + min_lr_ratio: float = 0.5 # minimum learning rate ratio for cosine_with_min_lr scheduler + not_apply_cfg_solver: bool = False + distill_cfg: float = 0.0 + scheduler_type: str = "" + linear_quadratic_threshold: float = 0.0 + linear_range: float = 0.0 + weight_decay: float = 0.0 + use_ema: bool = False + multi_phased_distill_schedule: str = "" + pred_decay_weight: float = 0.0 + pred_decay_type: str = "" + hunyuan_teacher_disable_cfg: bool = False + + # master_weight_type + master_weight_type: str = "" + + # VSA training decay parameters + VSA_decay_rate: float = 0.01 # decay rate -> 0.02 + VSA_decay_interval_steps: int = 1 # decay interval steps -> 50 + + # LoRA training parameters + lora_rank: int | None = None + lora_alpha: int | None = None + lora_training: bool = False + + # camera dataset + json_path: str = "" + causal: bool = False + window_frames: int = 9 + wandb_key: str = "f8f8f614d325bd04540f8517b142ca942b21017a" + wandb_entity: str = "zhyzhy" + action: bool = False + use_memory: bool = False + i2v_rate: float = 0.0 + train_time_shift: float = 1.0 + + # distillation args + generator_update_interval: int = 5 + min_timestep_ratio: float = 0.2 + max_timestep_ratio: float = 0.98 + real_score_guidance_scale: float = 3.5 + fake_score_learning_rate: float = 0.0 # separate learning rate for fake_score_transformer, if 0.0, use learning_rate + fake_score_lr_scheduler: str = "constant" # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler + training_state_checkpointing_steps: int = 0 # for resuming training + weight_only_checkpointing_steps: int = 0 # for inference + log_visualization: bool = False + # simulate generator forward to match inference + simulate_generator_forward: bool = False + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": + provided_args = clean_cli_args(args) + # Get all fields from the dataclass + attrs = [attr.name for attr in dataclasses.fields(cls)] + logger.info(provided_args) + # Create a dictionary of attribute values, with defaults for missing attributes + kwargs: dict[str, Any] = {} + for attr in attrs: + if attr == 'pipeline_config': + if "HunyuanTransformer" in provided_args["cls_name"]: continue + pipeline_config = PipelineConfig.from_kwargs(provided_args) + kwargs[attr] = pipeline_config + elif attr == 'mode': + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value) + kwargs[attr] = ExecutionMode.from_string( + mode_value) if isinstance(mode_value, str) else mode_value + elif attr == 'workload_type': + # Convert string to WorkloadType enum + workload_type_value = getattr(args, 'workload_type', + WorkloadType.T2V.value) + kwargs[attr] = WorkloadType.from_string( + workload_type_value) if isinstance( + workload_type_value, str) else workload_type_value + # Use getattr with default value from the dataclass for potentially missing attributes + else: + # Get the field to check its default value + field = dataclasses.fields(cls)[next( + i for i, f in enumerate(dataclasses.fields(cls)) + if f.name == attr)] + + # Check if the attribute is provided in args + if hasattr(args, attr): + value = getattr(args, attr) + else: + # Use the field's default value + if field.default_factory is not dataclasses.MISSING: + value = field.default_factory() + elif field.default is not dataclasses.MISSING: + value = field.default + else: + # No default value, use None + value = None + + kwargs[attr] = value + + return cls(**kwargs) # type: ignore + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument("--data-path", + type=str, + required=True, + help="Path to parquet files") + parser.add_argument("--dataloader-num-workers", + type=int, + required=True, + help="Number of workers for dataloader") + parser.add_argument("--num-height", + type=int, + required=True, + help="Number of heights") + parser.add_argument("--num-width", + type=int, + required=True, + help="Number of widths") + parser.add_argument("--num-frames", + type=int, + required=True, + help="Number of frames") + + # Training batch and model configuration + parser.add_argument("--train-batch-size", + type=int, + required=True, + help="Training batch size") + parser.add_argument("--num-latent-t", + type=int, + required=True, + help="Number of latent time steps") + parser.add_argument("--group-frame", + action=StoreBoolean, + help="Whether to group frames during training") + parser.add_argument("--group-resolution", + action=StoreBoolean, + help="Whether to group resolutions during training") + + # Model paths + parser.add_argument("--pretrained-model-name-or-path", + type=str, + required=True, + help="Path to pretrained model or model name") + parser.add_argument("--dit-model-name-or-path", + type=str, + required=False, + help="Path to DiT model or model name") + parser.add_argument("--cache-dir", + type=str, + help="Directory to cache models") + + # Diffusion settings + parser.add_argument("--ema-decay", + type=float, + default=0.999, + help="EMA decay rate") + parser.add_argument("--ema-start-step", + type=int, + default=0, + help="Step to start EMA") + parser.add_argument("--training-cfg-rate", + type=float, + help="Classifier-free guidance scale") + parser.add_argument( + "--precondition-outputs", + action=StoreBoolean, + help="Whether to precondition the outputs of the model") + + # Validation and logging + parser.add_argument("--validation-dataset-file", + type=str, + help="Path to unprocessed validation dataset") + parser.add_argument("--validation-preprocessed-path", + type=str, + help="Path to processed validation dataset") + parser.add_argument("--validation-sampling-steps", + type=str, + help="Validation sampling steps") + parser.add_argument("--validation-guidance-scale", + type=str, + help="Validation guidance scale") + parser.add_argument("--validation-steps", + type=float, + help="Number of validation steps") + parser.add_argument("--log-validation", + action=StoreBoolean, + help="Whether to log validation results") + parser.add_argument("--tracker-project-name", + type=str, + help="Project name for tracking") + parser.add_argument("--wandb-run-name", + type=str, + help="Run name for wandb") + parser.add_argument("--seed", + type=int, + default=42, + help="Seed for deterministic training") + + # Output configuration + parser.add_argument("--output-dir", + type=str, + required=True, + help="Output directory for checkpoints and logs") + parser.add_argument("--checkpoints-total-limit", + type=int, + help="Maximum number of checkpoints to keep") + parser.add_argument("--checkpointing-steps", + type=int, + help="Steps between checkpoints") + parser.add_argument( + "--training-state-checkpointing-steps", + type=int, + help= + "Steps between training state checkpoints (for resuming training)") + parser.add_argument( + "--weight-only-checkpointing-steps", + type=int, + help="Steps between weight-only checkpoints (for inference)") + parser.add_argument("--resume-from-checkpoint", + type=str, + help="Path to checkpoint to resume from") + parser.add_argument("--logging-dir", + type=str, + help="Directory for logging") + + # Training configuration + parser.add_argument("--num-train-epochs", + type=int, + help="Number of training epochs") + parser.add_argument("--max-train-steps", + type=int, + help="Maximum number of training steps") + parser.add_argument("--gradient-accumulation-steps", + type=int, + help="Number of steps to accumulate gradients") + parser.add_argument("--learning-rate", + type=float, + required=True, + help="Learning rate") + parser.add_argument("--scale-lr", + action=StoreBoolean, + help="Whether to scale learning rate") + parser.add_argument("--lr-scheduler", + type=str, + default="constant", + help="Learning rate scheduler type") + parser.add_argument("--lr-warmup-steps", + type=int, + default=10, + help="Number of warmup steps for learning rate") + parser.add_argument("--max-grad-norm", + type=float, + help="Maximum gradient norm") + parser.add_argument("--enable-gradient-checkpointing-type", + type=str, + choices=["full", "ops", "block_skip"], + default=None, + help="Gradient checkpointing type") + parser.add_argument("--selective-checkpointing", + type=float, + help="Selective checkpointing threshold") + parser.add_argument("--mixed-precision", + type=str, + help="Mixed precision training type") + parser.add_argument("--train-sp-batch-size", + type=int, + help="Training spatial parallelism batch size") + + parser.add_argument("--fsdp-sharding-strategy", + type=str, + help="FSDP sharding strategy") + + parser.add_argument( + "--weighting-scheme", + type=str, + default="uniform", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"], + ) + parser.add_argument( + "--logit-mean", + type=float, + default=0.0, + help="mean to use when using the `'logit_normal'` weighting scheme.", + ) + parser.add_argument( + "--logit-std", + type=float, + default=1.0, + help="std to use when using the `'logit_normal'` weighting scheme.", + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help= + "Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + + # Additional training parameters + parser.add_argument("--num-euler-timesteps", + type=int, + help="Number of Euler timesteps") + parser.add_argument("--lr-num-cycles", + type=int, + help="Number of learning rate cycles") + parser.add_argument("--lr-power", + type=float, + help="Learning rate power") + parser.add_argument( + "--min-lr-ratio", + type=float, + default=TrainingArgs.min_lr_ratio, + help="Minimum learning rate ratio for cosine_with_min_lr scheduler") + parser.add_argument("--not-apply-cfg-solver", + action=StoreBoolean, + help="Whether to not apply CFG solver") + parser.add_argument("--distill-cfg", + type=float, + help="Distillation CFG scale") + parser.add_argument("--scheduler-type", type=str, help="Scheduler type") + parser.add_argument("--linear-quadratic-threshold", + type=float, + help="Linear quadratic threshold") + parser.add_argument("--linear-range", type=float, help="Linear range") + parser.add_argument("--weight-decay", type=float, help="Weight decay") + parser.add_argument("--use-ema", + action=StoreBoolean, + help="Whether to use EMA") + parser.add_argument("--multi-phased-distill-schedule", + type=str, + help="Multi-phased distillation schedule") + parser.add_argument("--pred-decay-weight", + type=float, + help="Prediction decay weight") + parser.add_argument("--pred-decay-type", + type=str, + help="Prediction decay type") + parser.add_argument("--hunyuan-teacher-disable-cfg", + action=StoreBoolean, + help="Whether to disable CFG for Hunyuan teacher") + parser.add_argument("--master-weight-type", + type=str, + help="Master weight type") + + # VSA parameters for training with dense to sparse adaption + parser.add_argument( + "--VSA-decay-rate", # decay rate, how much sparsity you want to decay each step + type=float, + default=TrainingArgs.VSA_decay_rate, + help="VSA decay rate") + parser.add_argument( + "--VSA-decay-interval-steps", # how many steps for training with current sparsity + type=int, + default=TrainingArgs.VSA_decay_interval_steps, + help="VSA decay interval steps") + parser.add_argument("--lora-training", + action=StoreBoolean, + help="Whether to use LoRA training") + parser.add_argument("--lora-rank", type=int, help="LoRA rank") + parser.add_argument("--lora-alpha", type=int, help="LoRA alpha") + + # Distillation arguments + parser.add_argument("--generator-update-interval", + type=int, + default=TrainingArgs.generator_update_interval, + help="Ratio of student updates to critic updates.") + parser.add_argument("--min-timestep-ratio", + type=float, + default=TrainingArgs.min_timestep_ratio, + help="Minimum step ratio") + parser.add_argument("--max-timestep-ratio", + type=float, + default=TrainingArgs.max_timestep_ratio, + help="Maximum step ratio") + parser.add_argument("--real-score-guidance-scale", + type=float, + default=TrainingArgs.real_score_guidance_scale, + help="Teacher guidance scale") + parser.add_argument("--fake-score-learning-rate", + type=float, + default=TrainingArgs.fake_score_learning_rate, + help="Learning rate for fake score transformer") + parser.add_argument( + "--fake-score-lr-scheduler", + type=str, + default=TrainingArgs.fake_score_lr_scheduler, + help="Learning rate scheduler for fake score transformer") + parser.add_argument("--log-visualization", + action=StoreBoolean, + help="Whether to log visualization") + parser.add_argument( + "--simulate-generator-forward", + action=StoreBoolean, + help="Whether to simulate generator forward to match inference") + + # camera + parser.add_argument("--json-path", + type=str, + default="/all_data.json", + help="camera json path") + parser.add_argument("--causal", + action=StoreBoolean, + help="training ar model") + parser.add_argument("--window-frames", + type=int, + default=9, + help="window size") + parser.add_argument("--wandb-key", + type=str, + default="", + help="key") + parser.add_argument("--wandb-entity", + type=str, + default="", + help="entity") + parser.add_argument("--action", + action=StoreBoolean, + help="training ar model") + + parser.add_argument("--use-memory", + action=StoreBoolean, + help="training mem model") + parser.add_argument("--i2v-rate", + type=float, + default=0.0, + help="training rating for i2v") + parser.add_argument("--train-time-shift", + type=float, + default=1.0, + help="training time shift") + return parser + + +def parse_int_list(value: str) -> list[int]: + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] diff --git a/trainer/training/__init__.py b/trainer/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/training/activation_checkpoint.py b/trainer/training/activation_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..874d013f739c0dfa720ae4c51fbcf0a8b9ac3d6b --- /dev/null +++ b/trainer/training/activation_checkpoint.py @@ -0,0 +1,90 @@ +import collections +from enum import Enum + +import torch +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + checkpoint_wrapper) + +TRANSFORMER_BLOCK_NAMES = [ + "blocks", + "double_blocks", + "single_blocks", + "transformer_blocks", + "temporal_transformer_blocks", + "transformer_double_blocks", + "transformer_single_blocks", +] + + +class CheckpointType(str, Enum): + FULL = "full" + OPS = "ops" + BLOCK_SKIP = "block_skip" + + +_SELECTIVE_ACTIVATION_CHECKPOINTING_OPS = { + torch.ops.aten.mm.default, + torch.ops.aten._scaled_dot_product_efficient_attention.default, + torch.ops.aten._scaled_dot_product_flash_attention.default, + torch.ops._c10d_functional.reduce_scatter_tensor.default, +} + + +def apply_activation_checkpointing( + module: torch.nn.Module, + checkpointing_type: str = CheckpointType.FULL, + n_layer: int = 1) -> torch.nn.Module: + if checkpointing_type == CheckpointType.FULL: + module = _apply_activation_checkpointing_blocks(module) + elif checkpointing_type == CheckpointType.OPS: + module = _apply_activation_checkpointing_ops( + module, _SELECTIVE_ACTIVATION_CHECKPOINTING_OPS) + elif checkpointing_type == CheckpointType.BLOCK_SKIP: + module = _apply_activation_checkpointing_blocks(module, n_layer) + else: + raise ValueError( + f"Checkpointing type '{checkpointing_type}' not supported. Supported types are {CheckpointType.__members__.keys()}" + ) + return module + + +def _apply_activation_checkpointing_blocks(module: torch.nn.Module, + n_layer: int | None = None + ) -> torch.nn.Module: + for transformer_block_name in TRANSFORMER_BLOCK_NAMES: + blocks: torch.nn.Module = getattr(module, transformer_block_name, None) + if blocks is None: + continue + for index, (layer_id, block) in enumerate(blocks.named_children()): + if n_layer is None or index % n_layer == 0: + block = checkpoint_wrapper(block, preserve_rng_state=False) + blocks.register_module(layer_id, block) + return module + + +def _apply_activation_checkpointing_ops(module: torch.nn.Module, + ops) -> torch.nn.Module: + from torch.utils.checkpoint import (CheckpointPolicy, + create_selective_checkpoint_contexts) + + def _get_custom_policy(meta: dict[str, int]) -> CheckpointPolicy: + + def _custom_policy(ctx, func, *args, **kwargs): + mode = "recompute" if ctx.is_recompute else "forward" + mm_count_key = f"{mode}_mm_count" + if func == torch.ops.aten.mm.default: + meta[mm_count_key] += 1 + # Saves output of all compute ops, except every second mm + to_save = func in ops and not (func == torch.ops.aten.mm.default + and meta[mm_count_key] % 2 == 0) + return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE + + return _custom_policy + + def selective_checkpointing_context_fn(): + meta: dict[str, int] = collections.defaultdict(int) + return create_selective_checkpoint_contexts(_get_custom_policy(meta)) + + return checkpoint_wrapper(module, + context_fn=selective_checkpointing_context_fn, + preserve_rng_state=False) diff --git a/trainer/training/ar_hunyuan_mem_training_pipeline.py b/trainer/training/ar_hunyuan_mem_training_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..0d54950633e39e70cc68a3d7f6895412422bd52c --- /dev/null +++ b/trainer/training/ar_hunyuan_mem_training_pipeline.py @@ -0,0 +1,851 @@ +# SPDX-License-Identifier: Apache-2.0 +import dataclasses +import math +import os +import time +from abc import ABC, abstractmethod +from collections import deque +from collections.abc import Iterator +from typing import Any + +import imageio +import numpy as np +import torch +import torch.distributed as dist +import torchvision +from diffusers import FlowMatchEulerDiscreteScheduler +from einops import rearrange +from torch.utils.data import DataLoader +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm.auto import tqdm + + +from trainer.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) + +import trainer.envs as envs +from trainer.attention.backends.video_sparse_attn import ( + VideoSparseAttentionMetadataBuilder) +from trainer.configs.sample import SamplingParam +from trainer.dataset import build_ar_camera_hunyuan_w_mem_dataloader +from trainer.dataset.dataloader.schema import pyarrow_schema_t2v +from trainer.dataset.validation_dataset import ValidationDataset +from trainer.distributed import (cleanup_dist_env_and_memory, + get_local_torch_device, get_sp_group, + get_world_group) +from trainer.trainer_args import TrainerArgs, TrainingArgs +from trainer.forward_context import set_forward_context +from trainer.logger import init_logger +from trainer.pipelines import (ComposedPipelineBase, ForwardBatch, + LoRAPipeline, TrainingBatch) +from trainer.training.activation_checkpoint import ( + apply_activation_checkpointing) +from trainer.training.training_utils import ( + clip_grad_norm_while_handling_failing_dtensor_cases, + compute_density_for_timestep_sampling, get_scheduler, get_sigmas, + load_checkpoint, normalize_dit_input, save_checkpoint, +) +from trainer.utils import is_vsa_available, set_random_seed, shallow_asdict +# import muon optimizer +from trainer.training.muon import get_muon_optimizer + +import wandb # isort: skip + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +def _get_trainable_params(model: torch.nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def merge_tensor_by_mask(tensor_1, tensor_2, mask, dim): + assert tensor_1.shape == tensor_2.shape + # Mask is a 0/1 vector. Choose tensor_2 when the value is 1; otherwise, tensor_1 + masked_indices = torch.nonzero(mask).squeeze(1) + tmp = tensor_1.clone() + if dim == 0: + tmp[masked_indices] = tensor_2[masked_indices] + elif dim == 1: + tmp[:, masked_indices] = tensor_2[:, masked_indices] + elif dim == 2: + tmp[:, :, masked_indices] = tensor_2[:, :, masked_indices] + return tmp + +class TrainingPipeline(LoRAPipeline, ABC): + """ + A pipeline for training a model. All training pipelines should inherit from this class. + All reusable components and code should be implemented in this class. + """ + _required_config_modules = ["scheduler", "transformer"] + validation_pipeline: ComposedPipelineBase + train_dataloader: StatefulDataLoader + train_loader_iter: Iterator[dict[str, Any]] + current_epoch: int = 0 + + def __init__( + self, + model_path: str, + trainer_args: TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None) -> None: + trainer_args.inference_mode = False + self.lora_training = trainer_args.lora_training + if self.lora_training and trainer_args.lora_rank is None: + raise ValueError("lora rank must be set when using lora training") + + set_random_seed(trainer_args.seed) # for lora param init + super().__init__(model_path, trainer_args, required_config_modules, + loaded_modules) # type: ignore + + def create_pipeline_stages(self, trainer_args: TrainerArgs): + raise RuntimeError( + "create_pipeline_stages should not be called for training pipeline") + + def set_schemas(self) -> None: + self.train_dataset_schema = pyarrow_schema_t2v + + def initialize_training_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing training pipeline...") + self.device = get_local_torch_device() + self.training_args = training_args + world_group = get_world_group() + self.world_size = world_group.world_size + self.global_rank = world_group.rank + self.sp_group = get_sp_group() + self.rank_in_sp_group = self.sp_group.rank_in_group + self.sp_world_size = self.sp_group.world_size + self.local_rank = world_group.local_rank + self.transformer = self.get_module("transformer") + self.seed = training_args.seed + self.set_schemas() + self.action = training_args.action + # add the causal option + self.causal = training_args.causal + self.train_time_shift = training_args.train_time_shift + + # Set random seeds for deterministic training + assert self.seed is not None, "seed must be set" + set_random_seed(self.seed) + self.transformer.train() + if training_args.enable_gradient_checkpointing_type is not None: + self.transformer = apply_activation_checkpointing( + self.transformer, + checkpointing_type=training_args. + enable_gradient_checkpointing_type) + + self.set_trainable() + params_to_optimize = self.transformer.parameters() + params_to_optimize = list( + filter(lambda p: p.requires_grad, params_to_optimize)) + self.optimizer = get_muon_optimizer( + model=self.transformer, + lr=training_args.learning_rate, # Learning rate + weight_decay=training_args.weight_decay, # Weight decay + adamw_betas=(0.9, 0.999), # AdamW betas for 1D parameters + adamw_eps=1e-8, # AdamW epsilon + ) + + + self.init_steps = 0 + logger.info("optimizer: %s", self.optimizer) + + self.lr_scheduler = get_scheduler( + training_args.lr_scheduler, + optimizer=self.optimizer, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) + + self.train_dataset, self.train_dataloader = build_ar_camera_hunyuan_w_mem_dataloader( + json_path=training_args.json_path, + causal=training_args.causal, + window_frames=training_args.window_frames, + batch_size=training_args.train_batch_size, + num_data_workers=training_args.dataloader_num_workers, + drop_last=True, + drop_first_row=False, + seed=self.seed, + cfg_rate=training_args.training_cfg_rate, + i2v_rate=training_args.i2v_rate, + ) + + self.num_update_steps_per_epoch = math.ceil( + len(self.train_dataloader) / + training_args.gradient_accumulation_steps * training_args.sp_size / + training_args.train_sp_batch_size) + self.num_train_epochs = math.ceil(training_args.max_train_steps / + self.num_update_steps_per_epoch) + + # TODO(will): is there a cleaner way to track epochs? + self.current_epoch = 0 + + if self.global_rank == 0: + project = training_args.tracker_project_name or "trainer" + wandb_config = dataclasses.asdict(training_args) + wandb.login(key=training_args.wandb_key) + wandb.init( + config=wandb_config, + name=training_args.wandb_run_name, + entity=training_args.wandb_entity, + project=project, + ) + + @abstractmethod + def initialize_validation_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError( + "Training pipelines must implement this method") + + def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + return training_batch + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + latents = batch["latent"] + prompt_embed = batch["prompt_embed"] + + if self.action: + w2c = batch['w2c'] + intrinsic = batch['intrinsic'] + action = batch['action'] + + video_path = batch['video_path'] + image_cond = batch['image_cond'] + vision_states = batch['vision_states'] + prompt_mask = batch['prompt_mask'] + byt5_text_states = batch['byt5_text_states'] + byt5_text_mask = batch['byt5_text_mask'] + # add an indicator for memory training + select_window_out_flag = batch['select_window_out_flag'] + i2v_mask = batch['i2v_mask'] + else: + video_path = batch['path'] + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.prompt_embed = prompt_embed.to( + get_local_torch_device(), dtype=torch.bfloat16) + if self.action: + training_batch.w2c = w2c.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.intrinsic = intrinsic.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.action = action.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.video_path = video_path[0] + + training_batch.image_cond = image_cond.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.vision_states = vision_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.prompt_mask = prompt_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.byt5_text_states = byt5_text_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.byt5_text_mask = byt5_text_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.select_window_out_flag = select_window_out_flag[0] + training_batch.i2v_mask = i2v_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) # i2v mask only works for memory training + else: + training_batch.video_path = video_path[0] + + return training_batch + + def _normalize_dit_input(self, + training_batch: TrainingBatch) -> TrainingBatch: + # TODO(will): support other models + training_batch.latents = normalize_dit_input('wan', + training_batch.latents, + self.get_module("vae")) + return training_batch + + def timestep_transform(self, t, shift=1.0, num_timesteps=1000.0): + t = t / num_timesteps + t = shift * t / (1 + (shift - 1) * t) + t = t * num_timesteps + return t + + def _prepare_ar_dit_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + latents = training_batch.latents + batch_size = latents.shape[0] + latent_t = latents.shape[2] + latent_h = latents.shape[3] + latent_w = latents.shape[4] + noise = torch.randn(latents.shape, + generator=self.noise_gen_cuda, + device=latents.device, + dtype=latents.dtype) + + # add a parameter: chunk_latent_num means number of latent in one chunk + chunk_latent_num = 4 + first_chunk_num = 4 + u = compute_density_for_timestep_sampling( + weighting_scheme=self.training_args.weighting_scheme, + batch_size=batch_size * ((latent_t - first_chunk_num) // chunk_latent_num + 1), + generator=self.noise_random_generator, + logit_mean=self.training_args.logit_mean, + logit_std=self.training_args.logit_std, + mode_scale=self.training_args.mode_scale, + ) + u = u.reshape(batch_size, -1) + if first_chunk_num == 1: + first_u = u[:, :first_chunk_num] + rest_u = u[:, first_chunk_num:] + # Replicate the rest timesteps chunk_latent_num times + rest_u = rest_u.unsqueeze(-1).repeat_interleave(chunk_latent_num, dim=-1).reshape(batch_size, -1) + # Concatenate the first and rest sigmas + u = torch.cat([first_u, rest_u], dim=1).reshape(-1) + else: + u = u.unsqueeze(-1).repeat_interleave(chunk_latent_num, dim=-1).reshape(batch_size, -1).reshape(-1) + indices = (u * self.noise_scheduler.config.num_train_timesteps).long() + indices = (self.noise_scheduler.config.num_train_timesteps - self.timestep_transform(indices, self.train_time_shift)).long() + # TO DO: change the noise, for outside window, only add noise to the outside part + if training_batch.select_window_out_flag == 1: # select as memory training + for i in range(0, indices.shape[0] - 4, 4): + rand_val = torch.randint(500, 985, (1, ), device=latents.device) + indices[i:i + 4] = rand_val + + timesteps = self.noise_scheduler.timesteps[indices].to(device=self.device) + if self.training_args.sp_size > 1: + # Make sure that the timesteps are the same across all sp processes. + sp_group = get_sp_group() + sp_group.broadcast(timesteps, src=0) + + sigmas = get_sigmas( + self.noise_scheduler, + latents.device, + timesteps, + n_dim=latents.ndim, + dtype=latents.dtype, + ) + sigmas = rearrange(sigmas, '(B D) C T H W -> B C (D T) H W', D=latent_t) + noisy_model_input = (1.0 - + sigmas) * training_batch.latents + sigmas * noise + training_batch.noisy_model_input = noisy_model_input + training_batch.timesteps = timesteps + training_batch.sigmas = sigmas + training_batch.noise = noise + training_batch.raw_latent_shape = training_batch.latents.shape + + return training_batch + + def _build_attention_metadata( + self, training_batch: TrainingBatch) -> TrainingBatch: + latents_shape = training_batch.raw_latent_shape + patch_size = self.training_args.pipeline_config.dit_config.patch_size + current_vsa_sparsity = training_batch.current_vsa_sparsity + assert latents_shape is not None + assert training_batch.timesteps is not None + if vsa_available and envs.TRAINER_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder( # type: ignore + ).build( # type: ignore + raw_latent_shape=latents_shape[2:5], + current_timestep=training_batch.timesteps, + patch_size=patch_size, + VSA_sparsity=current_vsa_sparsity, + device=get_local_torch_device()) + else: + training_batch.attn_metadata = None + + return training_batch + + def _build_rope_idx(self, + training_batch: TrainingBatch) -> TrainingBatch: + rank_in_sp_group = get_sp_parallel_rank() + per_sp_seq_length = training_batch.latents.shape[2] * training_batch.per_seq_length + + training_batch.current_start = per_sp_seq_length * rank_in_sp_group + training_batch.current_end = per_sp_seq_length * rank_in_sp_group + per_sp_seq_length + return training_batch + + def _prepare_cond_latents(self, task_type, cond_latents, latents, multitask_mask): + """ + Prepare conditional latents and mask for multitask training. + + Args: + task_type: Type of task ("i2v" or "t2v"). + cond_latents: Conditional latents tensor. + latents: Main latents tensor. + multitask_mask: Multitask mask tensor. + + Returns: + tuple[torch.Tensor, torch.Tensor]: + - latents_concat: Concatenated conditional latents. + - mask_concat: Concatenated mask tensor. + """ + latents_concat = None + mask_concat = None + + if cond_latents is not None and task_type == 'i2v': + latents_concat = cond_latents.repeat(1, 1, latents.shape[2], 1, 1) + latents_concat[:, :, 1:, :, :] = 0.0 + else: + latents_concat = torch.zeros(latents.shape[0], latents.shape[1], latents.shape[2], latents.shape[3], + latents.shape[4]).to(latents.device) + + mask_zeros = torch.zeros(latents.shape[0], 1, latents.shape[2], latents.shape[3], latents.shape[4]) + mask_ones = torch.ones(latents.shape[0], 1, latents.shape[2], latents.shape[3], latents.shape[4]) + mask_concat = merge_tensor_by_mask(mask_zeros.cpu(), mask_ones.cpu(), mask=multitask_mask.cpu(), dim=2).to( + device=latents.device) + + cond_latents = torch.concat([latents_concat, mask_concat], dim=1) + + return cond_latents + + def get_task_mask(self, task_type, latent_target_length): + if task_type == "t2v": + mask = torch.zeros(latent_target_length) + elif task_type == "i2v": + mask = torch.zeros(latent_target_length) + mask[0] = 1.0 + else: + raise ValueError(f"{task_type} is not supported !") + return mask + + def _build_input_kwargs(self, + training_batch: TrainingBatch) -> TrainingBatch: + extra_kwargs = { + "byt5_text_states": training_batch.byt5_text_states, + "byt5_text_mask": training_batch.byt5_text_mask, + } + + multitask_mask = self.get_task_mask("i2v", training_batch.noisy_model_input.shape[2]).to(self.device) + cond_latents = self._prepare_cond_latents( + "i2v", training_batch.image_cond, training_batch.noisy_model_input, multitask_mask + ) + + latents_concat = torch.concat([training_batch.noisy_model_input, cond_latents], dim=1) + training_batch.input_kwargs = { + "hidden_states": + latents_concat, + "timestep": + training_batch.timesteps.to(get_local_torch_device(), + dtype=torch.bfloat16), + "timestep_txt": torch.tensor(0).unsqueeze(0).to(get_local_torch_device(), + dtype=torch.bfloat16), # for ar model, we set txt timestep to 0 + "text_states": + training_batch.prompt_embed, + "text_states_2": None, + "encoder_attention_mask": training_batch.prompt_mask, + "timestep_r": None, + "vision_states": training_batch.vision_states, + "mask_type": "i2v", + "guidance": None, + "extra_kwargs": extra_kwargs, + + "viewmats": training_batch.w2c, + "Ks": training_batch.intrinsic, + "action": training_batch.action.reshape(-1), + "return_dict": False, + + } + return training_batch + + def _transformer_forward_and_compute_loss( + self, training_batch: TrainingBatch) -> TrainingBatch: + if vsa_available and envs.TRAINER_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + assert training_batch.attn_metadata is not None + else: + assert training_batch.attn_metadata is None + input_kwargs = training_batch.input_kwargs + + with set_forward_context( + current_timestep=training_batch.current_timestep, + attn_metadata=training_batch.attn_metadata): + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): + model_pred = self.transformer(**input_kwargs)[0] + + if self.training_args.precondition_outputs: + assert training_batch.sigmas is not None + model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas + assert training_batch.latents is not None + assert training_batch.noise is not None + target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents + i2v_mask = training_batch.i2v_mask + if training_batch.select_window_out_flag == 1 and self.causal: + i2v_mask[:,:,:-4,...] = 0 # only compute the last chunk for outside window training + assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}" + + diff = (model_pred.float() * i2v_mask - target.float() * i2v_mask) ** 2 + loss = diff.sum() / max(i2v_mask.sum(), 1) / self.training_args.gradient_accumulation_steps + + loss.backward() + avg_loss = loss.detach().clone() + + dist.all_reduce(avg_loss, op=dist.ReduceOp.MAX) + training_batch.total_loss += avg_loss.item() + + return training_batch + + def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: + max_grad_norm = self.training_args.max_grad_norm + + # TODO(will): perhaps move this into transformer api so that we can do + # the following: + # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + if max_grad_norm is not None: + model_parts = [self.transformer] + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + [p for m in model_parts for p in m.parameters()], + max_grad_norm, + foreach=None, + ) + assert grad_norm is not float('nan') or grad_norm is not float( + 'inf') + grad_norm = grad_norm.item() if grad_norm is not None else 0.0 + else: + grad_norm = 0.0 + training_batch.grad_norm = grad_norm + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + training_batch = self._prepare_training(training_batch) + + for _ in range(self.training_args.gradient_accumulation_steps): + training_batch = self._get_next_batch(training_batch) + + training_batch = self._prepare_ar_dit_inputs(training_batch) + + training_batch = self._build_input_kwargs(training_batch) + + training_batch = self._transformer_forward_and_compute_loss( + training_batch) + + training_batch = self._clip_grad_norm(training_batch) + grad_norm = torch.tensor(training_batch.grad_norm).to(get_local_torch_device()) + dist.all_reduce(grad_norm, op=dist.ReduceOp.MAX) + training_batch.grad_norm = grad_norm.item() + + if self.global_rank == 0 and training_batch.grad_norm >= 10.0: + print(self.global_rank, training_batch.grad_norm, training_batch.current_timestep, training_batch.video_path) + + if training_batch.grad_norm < 10.0 or (not self.action): + self.optimizer.step() + self.lr_scheduler.step() + + training_batch.total_loss = training_batch.total_loss + training_batch.grad_norm = training_batch.grad_norm + return training_batch + + def _resume_from_checkpoint(self) -> None: + logger.info("Loading checkpoint from %s", + self.training_args.resume_from_checkpoint) + resumed_step = load_checkpoint( + self.transformer, self.global_rank, + self.training_args.resume_from_checkpoint, self.optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + if resumed_step > 0: + self.init_steps = resumed_step + logger.info("Successfully resumed from step %s", resumed_step) + else: + logger.warning("Failed to load checkpoint, starting from step 0") + self.init_steps = 0 + + def train(self) -> None: + assert self.seed is not None, "seed must be set" + set_random_seed(self.seed + self.global_rank) + logger.info('rank: %s: start training', + self.global_rank, + local_main_process_only=False) + if not self.post_init_called: + self.post_init() + num_trainable_params = _get_trainable_params(self.transformer) + logger.info("Starting training with %s B trainable parameters", + round(num_trainable_params / 1e9, 3)) + + # Set random seeds for deterministic training + self.noise_random_generator = torch.Generator(device="cpu").manual_seed( + self.seed) + self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed( + self.seed) + self.validation_random_generator = torch.Generator( + device="cpu").manual_seed(self.seed) + logger.info("Initialized random seeds with seed: %s", self.seed) + + self.noise_scheduler = FlowMatchEulerDiscreteScheduler() + + if self.training_args.resume_from_checkpoint: + self._resume_from_checkpoint() + + self.train_loader_iter = iter(self.train_dataloader) + + step_times: deque[float] = deque(maxlen=100) + + self._log_training_info() + + # Train! + progress_bar = tqdm( + range(0, self.training_args.max_train_steps), + initial=self.init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + disable=self.local_rank > 0, + ) + + for step in range(self.init_steps + 1, + self.training_args.max_train_steps + 1): + + self.train_dataset.update_max_frames(step) + + start_time = time.perf_counter() + if vsa_available: + vsa_sparsity = self.training_args.VSA_sparsity + vsa_decay_rate = self.training_args.VSA_decay_rate + vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps + current_decay_times = min(step // vsa_decay_interval_steps, + vsa_sparsity // vsa_decay_rate) + current_vsa_sparsity = current_decay_times * vsa_decay_rate + else: + current_vsa_sparsity = 0.0 + + training_batch = TrainingBatch() + training_batch.current_timestep = step + training_batch.current_vsa_sparsity = current_vsa_sparsity + training_batch = self.train_one_step(training_batch) + + loss = training_batch.total_loss + grad_norm = training_batch.grad_norm + + step_time = time.perf_counter() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + progress_bar.set_postfix({ + "loss": f"{loss:.4f}", + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm, + }) + progress_bar.update(1) + if self.global_rank == 0: + wandb.log( + { + "train_loss": loss, + "learning_rate": self.lr_scheduler.get_last_lr()[0], + "step_time": step_time, + "avg_step_time": avg_step_time, + "grad_norm": grad_norm, + "vsa_sparsity": current_vsa_sparsity, + }, + step=step, + ) + + if step % self.training_args.checkpointing_steps == 0: + save_checkpoint(self.transformer, self.global_rank, + self.training_args.output_dir, step, + self.optimizer, self.train_dataloader, + self.lr_scheduler, self.noise_random_generator) + self.transformer.train() + self.sp_group.barrier() + + wandb.finish() + save_checkpoint(self.transformer, self.global_rank, + self.training_args.output_dir, + self.training_args.max_train_steps, self.optimizer, + self.train_dataloader, self.lr_scheduler, + self.noise_random_generator) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + def _log_training_info(self) -> None: + total_batch_size = (self.world_size * + self.training_args.gradient_accumulation_steps / + self.training_args.sp_size * + self.training_args.train_sp_batch_size) + logger.info("***** Running training *****") + logger.info(" Num examples = %s", len(self.train_dataset)) + logger.info(" Dataloader size = %s", len(self.train_dataloader)) + logger.info(" Num Epochs = %s", self.num_train_epochs) + logger.info(" Resume training from step %s", + self.init_steps) # type: ignore + logger.info(" Instantaneous batch size per device = %s", + self.training_args.train_batch_size) + logger.info( + " Total train batch size (w. data & sequence parallel, accumulation) = %s", + total_batch_size) + logger.info(" Gradient Accumulation steps = %s", + self.training_args.gradient_accumulation_steps) + logger.info(" Total optimization steps = %s", + self.training_args.max_train_steps) + logger.info(" Total training parameters per FSDP shard = %s B", + round(_get_trainable_params(self.transformer) / 1e9, 3)) + # print dtype + logger.info(" Master weight dtype: %s", + self.transformer.parameters().__next__().dtype) + + gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + logger.info("GPU memory usage before train_one_step: %s MB", + gpu_memory_usage) + logger.info("VSA validation sparsity: %s", + self.training_args.VSA_sparsity) + + def _prepare_validation_batch(self, sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: dict[str, Any], + num_inference_steps: int) -> ForwardBatch: + sampling_param.prompt = validation_batch['prompt'] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + assert self.seed is not None + sampling_param.seed = self.seed + + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 + sampling_param.num_frames = num_frames + batch = ForwardBatch( + **shallow_asdict(sampling_param), + latents=None, + generator=self.validation_random_generator, + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + + return batch + + @torch.no_grad() + def _log_validation(self, transformer, training_args, global_step) -> None: + """ + Generate a validation video and log it to wandb to check the quality during training. + """ + training_args.inference_mode = True + training_args.dit_cpu_offload = True + if not training_args.log_validation: + return + if self.validation_pipeline is None: + raise ValueError("Validation pipeline is not set") + + logger.info("Starting validation") + + # Create sampling parameters if not provided + sampling_param = SamplingParam.from_pretrained(training_args.model_path) + + # Prepare validation prompts + logger.info('rank: %s: trainer_args.validation_dataset_file: %s', + self.global_rank, + training_args.validation_dataset_file, + local_main_process_only=False) + validation_dataset = ValidationDataset( + training_args.validation_dataset_file) + validation_dataloader = DataLoader(validation_dataset, + batch_size=None, + num_workers=0) + + transformer.eval() + + validation_steps = training_args.validation_sampling_steps.split(",") + validation_steps = [int(step) for step in validation_steps] + validation_steps = [step for step in validation_steps if step > 0] + # Log validation results for this step + world_group = get_world_group() + num_sp_groups = world_group.world_size // self.sp_group.world_size + + # Process each validation prompt for each validation step + for num_inference_steps in validation_steps: + logger.info("rank: %s: num_inference_steps: %s", + self.global_rank, + num_inference_steps, + local_main_process_only=False) + step_videos: list[np.ndarray] = [] + step_captions: list[str] = [] + + for validation_batch in validation_dataloader: + batch = self._prepare_validation_batch(sampling_param, + training_args, + validation_batch, + num_inference_steps) + logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s", + self.global_rank, + self.rank_in_sp_group, + batch.prompt, + local_main_process_only=False) + + assert batch.prompt is not None and isinstance( + batch.prompt, str) + step_captions.append(batch.prompt) + + # Run validation inference + output_batch = self.validation_pipeline.forward( + batch, training_args) + samples = output_batch.output + + if self.rank_in_sp_group != 0: + continue + + # Process outputs + video = rearrange(samples, "b c t h w -> t b c h w") + frames = [] + for x in video: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + step_videos.append(frames) + + # Only sp_group leaders (rank_in_sp_group == 0) need to send their + # results to global rank 0 + if self.rank_in_sp_group == 0: + if self.global_rank == 0: + # Global rank 0 collects results from all sp_group leaders + all_videos = step_videos # Start with own results + all_captions = step_captions + + # Receive from other sp_group leaders + for sp_group_idx in range(1, num_sp_groups): + src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders + recv_videos = world_group.recv_object(src=src_rank) + recv_captions = world_group.recv_object(src=src_rank) + all_videos.extend(recv_videos) + all_captions.extend(recv_captions) + + video_filenames = [] + for i, (video, caption) in enumerate( + zip(all_videos, all_captions, strict=True)): + os.makedirs(training_args.output_dir, exist_ok=True) + filename = os.path.join( + training_args.output_dir, + f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" + ) + imageio.mimsave(filename, video, fps=sampling_param.fps) + video_filenames.append(filename) + + logs = { + f"validation_videos_{num_inference_steps}_steps": [ + wandb.Video(filename, caption=caption) + for filename, caption in zip( + video_filenames, all_captions, strict=True) + ] + } + wandb.log(logs, step=global_step) + else: + # Other sp_group leaders send their results to global rank 0 + world_group.send_object(step_videos, dst=0) + world_group.send_object(step_captions, dst=0) + + # Re-enable gradients for training + training_args.inference_mode = False + transformer.train() diff --git a/trainer/training/ar_hunyuan_w_mem_training_pipeline.py b/trainer/training/ar_hunyuan_w_mem_training_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..e68eddec44f2257a217e4fa4843281eb4e472258 --- /dev/null +++ b/trainer/training/ar_hunyuan_w_mem_training_pipeline.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +from copy import deepcopy +import os +import sys +sys.path.append(os.path.abspath('.')) + +from trainer.trainer_args import TrainerArgs, TrainingArgs +from trainer.logger import init_logger +from trainer.training.ar_hunyuan_mem_training_pipeline import TrainingPipeline +from trainer.utils import is_vsa_available + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class HunyuanTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Hunyuan. + """ + _required_config_modules = ["transformer"] + + def initialize_pipeline(self, trainer_args: TrainerArgs): + pass + # self.modules["scheduler"] = FlowUniPCMultistepScheduler( + # shift=trainer_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + pass + + +def main(args) -> None: + logger.info("Starting training pipeline...") + + pipeline = HunyuanTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + import torch.multiprocessing as mp + mp.set_start_method('spawn', force=True) + + argv = sys.argv + from trainer.trainer_args import TrainingArgs + from trainer.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = TrainerArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/trainer/training/checkpointing_utils.py b/trainer/training/checkpointing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bc6aeed55aba4ce7541f16d70600172485686c48 --- /dev/null +++ b/trainer/training/checkpointing_utils.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +import random +from typing import Any + +import numpy as np +import torch +import torch.distributed.checkpoint.stateful +from torch.distributed.checkpoint.state_dict import (StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict) + + +class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful): + + def __init__(self, model: torch.nn.Module) -> None: + self.model = model + + def state_dict(self) -> dict[str, Any]: + state_dict = get_model_state_dict( + self.model) # type: ignore[no-any-return] + # filter out non-trainable parameters + param_requires_grad = set([ + k for k, v in dict(self.model.named_parameters()).items() + if v.requires_grad + ]) + state_dict = { + k: v + for k, v in state_dict.items() if k in param_requires_grad + } + return state_dict # type: ignore + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + set_model_state_dict( + self.model, + model_state_dict=state_dict, + options=StateDictOptions(strict=False), + ) + + +class OptimizerWrapper(torch.distributed.checkpoint.stateful.Stateful): + + def __init__(self, model: torch.nn.Module, + optimizer: torch.optim.Optimizer) -> None: + self.model = model + self.optimizer = optimizer + + def state_dict(self) -> dict[str, Any]: + return get_optimizer_state_dict( # type: ignore[no-any-return] + self.model, + self.optimizer, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + set_optimizer_state_dict( + self.model, + self.optimizer, + optim_state_dict=state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + + +class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful): + + def __init__(self, scheduler) -> None: + self.scheduler = scheduler + + def state_dict(self) -> dict[str, Any]: + return {"scheduler": self.scheduler.state_dict()} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.scheduler.load_state_dict(state_dict["scheduler"]) + + +class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful): + + def __init__(self, noise_generator: torch.Generator | None = None) -> None: + self.noise_generator = noise_generator + + def state_dict(self) -> dict[str, Any]: + state = { + "torch_rng_state": torch.get_rng_state(), + "numpy_rng_state": np.random.get_state(), + "python_rng_state": random.getstate(), + } + + if torch.cuda.is_available(): + state["cuda_rng_state"] = torch.cuda.get_rng_state() + if torch.cuda.device_count() > 1: + state["cuda_rng_state_all"] = torch.cuda.get_rng_state_all() + + if self.noise_generator is not None: + state["noise_generator_state"] = self.noise_generator.get_state() + + return state + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if "torch_rng_state" in state_dict: + torch.set_rng_state(state_dict["torch_rng_state"]) + + if "numpy_rng_state" in state_dict: + np.random.set_state(state_dict["numpy_rng_state"]) + + if "python_rng_state" in state_dict: + random.setstate(state_dict["python_rng_state"]) + + # Restore CUDA random state + if torch.cuda.is_available(): + if "cuda_rng_state" in state_dict: + torch.cuda.set_rng_state(state_dict["cuda_rng_state"]) + if "cuda_rng_state_all" in state_dict: + torch.cuda.set_rng_state_all(state_dict["cuda_rng_state_all"]) + + # Restore noise generator state + if "noise_generator_state" in state_dict and self.noise_generator is not None: + self.noise_generator.set_state(state_dict["noise_generator_state"]) diff --git a/trainer/training/muon.py b/trainer/training/muon.py new file mode 100644 index 0000000000000000000000000000000000000000..eff6b95918c979e183a832233db29412bc333cea --- /dev/null +++ b/trainer/training/muon.py @@ -0,0 +1,234 @@ +import math +import torch +try: + from torch.distributed.tensor import DTensor + from torch.distributed.tensor.placement_types import ( + Partial, + Placement, + Replicate, + Shard, + ) +except ImportError: + # handle old pytorch versions + Dtensor = None + +# This code is modified from the GitHub repository of KellerJordan: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +def zeropower_via_newtonschulz5(G, steps = 5): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + if isinstance(G, DTensor): + device_mesh = G.device_mesh + G = G.full_tensor() + else: + device_mesh = None + + assert len(G.shape) >= 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G + if G.size(-2) > G.size(-1): + X = X.mT + # Ensure spectral norm is at most 1 + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng + X = a * X + B @ X + + if G.size(-2) > G.size(-1): + X = X.mT + + if device_mesh != None: + return DTensor.from_local(X, device_mesh) + else: + return X + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + lr=1e-3, + wd=0.1, + muon_params=None, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_params=None, + adamw_betas=(0.95, 0.95), + adamw_eps=1e-8, + ): + + defaults = dict( + lr=lr, + wd=wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) + + params = list(muon_params) + adamw_params = list(adamw_params) if adamw_params is not None else [] + params.extend(adamw_params) + super().__init__(params, defaults) + # Sort parameters into those for which we will use Muon, and those for which we will not + for p in muon_params: + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim >= 2, p.ndim + self.state[p]["use_muon"] = True + for p in adamw_params: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + # generate weight updates in distributed fashion + for p in params: + # sanity check + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + g = g.bfloat16() + u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u.view(p.shape), alpha=-adjusted_lr) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group['lr'] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss + +# help function to create the Muon optimizer +def get_muon_optimizer(model, + lr=1e-3, + weight_decay=0.1, + momentum=0.95, + adamw_betas=(0.95, 0.95), + adamw_eps=1e-8): + muon_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and p.ndim >= 2 + ] + adamw_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and not (p.ndim >= 2) + ] + + return Muon( + lr=lr, + wd=weight_decay, + muon_params=muon_params, + momentum=momentum, + adamw_params=adamw_params, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) \ No newline at end of file diff --git a/trainer/training/training_utils.py b/trainer/training/training_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..53d3d54caa76b6a5057396ede69acc258f36ac9b --- /dev/null +++ b/trainer/training/training_utils.py @@ -0,0 +1,1082 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import math +import os +import time +from collections.abc import Callable, Iterator +from enum import Enum +from typing import Any +import shutil + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from einops import rearrange +from safetensors.torch import save_file +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LambdaLR + +from trainer.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) +from trainer.logger import init_logger +from trainer.training.checkpointing_utils import (ModelWrapper, + OptimizerWrapper, + RandomStateWrapper, + SchedulerWrapper) + +logger = init_logger(__name__) + +_HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False + + +def gather_state_dict_on_cpu_rank0( + model, + device: torch.device | None = None, +) -> dict[str, Any]: + rank = dist.get_rank() + cpu_state_dict = {} + sharded_sd = model.state_dict() + param_requires_grad = set([ + k.replace("._checkpoint_wrapped_module.", ".") + for k, v in dict(model.named_parameters()).items() if v.requires_grad + ]) + for param_name, param in sharded_sd.items(): + if param_name not in param_requires_grad: + continue + if hasattr(param, "_local_tensor"): + # DTensor case + if param.is_cpu: + # Gather directly on CPU + param = param.full_tensor() + else: + if device is not None: + param = param.to(device) + param = param.full_tensor() + else: + # Regular tensor case + if param.is_cpu: + pass + else: + if device is not None: + param = param.to(device) + + if rank == 0: + cpu_state_dict[param_name] = param.cpu() + + return cpu_state_dict + +def compute_density_for_timestep_sampling( + weighting_scheme: str, + batch_size: int, + generator, + logit_mean: float | None = None, + logit_std: float | None = None, + mode_scale: float | None = None, +): + """ + Compute the density for sampling the timesteps when doing SD3 training. + + Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. + + SD3 paper reference: https://arxiv.org/abs/2403.03206v1. + """ + if weighting_scheme == "logit_normal": + # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). + u = torch.normal( + mean=logit_mean, + std=logit_std, + size=(batch_size, ), + device="cpu", + generator=generator, + ) + u = torch.nn.functional.sigmoid(u) + elif weighting_scheme == "mode": + u = torch.rand(size=(batch_size, ), device="cpu", generator=generator) + u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2)**2 - 1 + u) + else: + u = torch.rand(size=(batch_size, ), device="cpu", generator=generator) + return u + + +def get_sigmas(noise_scheduler, + device, + timesteps, + n_dim=4, + dtype=torch.float32) -> torch.Tensor: + sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(device) + timesteps = timesteps.to(device) + step_indices = [(schedule_timesteps == t).nonzero().item() + for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + +def remove_previous_state(output_dir, rank): + checkpoint_dir_list = os.listdir(output_dir) + for checkpoint_dir in checkpoint_dir_list: + save_dir = os.path.join(output_dir, checkpoint_dir) + dcp_dir = os.path.join(save_dir, "distributed_checkpoint") + if os.path.exists(dcp_dir) and rank == 0: + shutil.rmtree(dcp_dir) + +def save_checkpoint(transformer, + rank, + output_dir, + step, + optimizer=None, + dataloader=None, + scheduler=None, + noise_generator=None) -> None: + """ + Save checkpoint following finetrainer's distributed checkpoint approach. + Saves both distributed checkpoint and consolidated model weights. + """ + if os.path.exists(output_dir): + remove_previous_state(output_dir, rank) + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + + states = { + "model": ModelWrapper(transformer), + "random_state": RandomStateWrapper(noise_generator), + } + + if optimizer is not None: + states["optimizer"] = OptimizerWrapper(transformer, optimizer) + + if dataloader is not None: + states["dataloader"] = dataloader + + if scheduler is not None: + states["scheduler"] = SchedulerWrapper(scheduler) + dcp_dir = os.path.join(save_dir, "distributed_checkpoint") + logger.info("rank: %s, saving distributed checkpoint to %s", + rank, + dcp_dir, + local_main_process_only=False) + + begin_time = time.perf_counter() + dcp.save(states, checkpoint_id=dcp_dir) + end_time = time.perf_counter() + + logger.info("rank: %s, distributed checkpoint saved in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False) + + cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None) + if rank == 0: + # Save model weights (consolidated) + transformer_save_dir = os.path.join(save_dir, "transformer") + os.makedirs(transformer_save_dir, exist_ok=True) + weight_path = os.path.join(transformer_save_dir, + "diffusion_pytorch_model.safetensors") + logger.info("rank: %s, saving consolidated checkpoint to %s", + rank, + weight_path, + local_main_process_only=False) + + # Convert training format to diffusers format and save + # diffusers_state_dict = custom_to_hf_state_dict( + # cpu_state, transformer.reverse_param_names_mapping) + save_file(cpu_state, weight_path) + + logger.info("rank: %s, consolidated checkpoint saved to %s", + rank, + weight_path, + local_main_process_only=False) + + # Save model config + config_dict = transformer.config + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(transformer_save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + logger.info("--> checkpoint saved at step %s to %s", step, weight_path) + +def load_checkpoint(transformer, + rank, + checkpoint_path, + optimizer=None, + dataloader=None, + scheduler=None, + noise_generator=None) -> int: + """ + Load checkpoint following finetrainer's distributed checkpoint approach. + Returns the step number from which training should resume. + """ + if not os.path.exists(checkpoint_path): + logger.warning("Checkpoint path %s does not exist", checkpoint_path) + return 0 + + # Extract step number from checkpoint path + step = int(os.path.basename(checkpoint_path).split('-')[-1]) + + if rank == 0: + logger.info("Loading checkpoint from step %s", step) + + dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint") + + if not os.path.exists(dcp_dir): + logger.warning("Distributed checkpoint directory %s does not exist", + dcp_dir) + return 0 + + states = { + "model": ModelWrapper(transformer), + "random_state": RandomStateWrapper(noise_generator), + } + + if optimizer is not None: + states["optimizer"] = OptimizerWrapper(transformer, optimizer) + + if dataloader is not None: + states["dataloader"] = dataloader + + if scheduler is not None: + states["scheduler"] = SchedulerWrapper(scheduler) + + logger.info("rank: %s, loading distributed checkpoint from %s", + rank, + dcp_dir, + local_main_process_only=False) + + begin_time = time.perf_counter() + dcp.load(states, checkpoint_id=dcp_dir) + end_time = time.perf_counter() + + logger.info("rank: %s, distributed checkpoint loaded in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False) + logger.info("--> checkpoint loaded from step %s", step) + + return step + + +def normalize_dit_input(model_type, latents, vae) -> torch.Tensor: + if model_type == "hunyuan_hf" or model_type == "hunyuan": + return latents * 0.476986 + elif model_type == "wan": + latents_mean = torch.tensor(vae.latents_mean) + latents_std = 1.0 / torch.tensor(vae.latents_std) + + latents_mean = latents_mean.view(1, -1, 1, 1, + 1).to(device=latents.device) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device) + latents = ((latents.float() - latents_mean) * latents_std).to(latents) + return latents + else: + raise NotImplementedError(f"model_type {model_type} not supported") + + +def shard_latents_across_sp(latents: torch.Tensor, + num_latent_t: int) -> torch.Tensor: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + latents = latents[:, :, :num_latent_t] + if sp_world_size > 1: + latents = rearrange(latents, + "b c (n s) h w -> b c n s h w", + n=sp_world_size).contiguous() + latents = latents[:, :, rank_in_sp_group, :, :, :] + return latents + +def shard_latents_dim_across_sp(latents: torch.Tensor, total_dim) -> torch.Tensor: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + # latents = latents[:, :, :num_latent_t] + if sp_world_size > 1: + if total_dim == 5: + latents = rearrange(latents, + "b c (n s) h w -> b c n s h w", + n=sp_world_size).contiguous() + latents = latents[:, :, rank_in_sp_group, :, :, :] + elif total_dim == 4: + latents = rearrange(latents, + "b (n s) h w -> b n s h w", + n=sp_world_size).contiguous() + latents = latents[:, rank_in_sp_group, :, :, :] + elif total_dim == 3: + latents = rearrange(latents, + "b (n s) c -> b n s c", + n=sp_world_size).contiguous() + latents = latents[:, rank_in_sp_group, :, :] + elif total_dim == 2: + latents = rearrange(latents, + "b (n s)-> b n s", + n=sp_world_size).contiguous() + latents = latents[:, rank_in_sp_group, :] + elif total_dim == 1: + latents = rearrange(latents, + "(n s)-> n s", + n=sp_world_size).contiguous() + latents = latents[rank_in_sp_group, :] + else: + raise NotImplementedError(f"total_dim {total_dim} not supported") + return latents + + +def clip_grad_norm_while_handling_failing_dtensor_cases( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, + pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, +) -> torch.Tensor | None: + global _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES + + if not _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES: + try: + return clip_grad_norm_(parameters, max_norm, norm_type, + error_if_nonfinite, foreach, pp_mesh) + except NotImplementedError as e: + if "DTensor does not support cross-mesh operation" in str(e): + # https://github.com/pytorch/pytorch/issues/134212 + logger.warning( + "DTensor does not support cross-mesh operation. If you haven't fully tensor-parallelized your " + "model, while combining other parallelisms such as FSDP, it could be the reason for this error. " + "Gradient clipping will be skipped and gradient norm will not be logged." + ) + except Exception as e: + logger.warning( + "An error occurred while clipping gradients: %s. Gradient clipping will be skipped and gradient " + "norm will not be logged.", e) + _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = True + return None + + +# Copied from https://github.com/pytorch/torchtitan/blob/4a169701555ab9bd6ca3769f9650ae3386b84c6e/torchtitan/utils.py#L362 +@torch.no_grad() +def clip_grad_norm_( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, + pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, +) -> torch.Tensor: + r""" + Clip the gradient norm of parameters. + + Gradient norm clipping requires computing the gradient norm over the entire model. + `torch.nn.utils.clip_grad_norm_` only computes gradient norm along DP/FSDP/TP dimensions. + We need to manually reduce the gradient norm across PP stages. + See https://github.com/pytorch/torchtitan/issues/596 for details. + + Args: + parameters (`torch.Tensor` or `List[torch.Tensor]`): + Tensors that will have gradients normalized. + max_norm (`float`): + Maximum norm of the gradients after clipping. + norm_type (`float`, defaults to `2.0`): + Type of p-norm to use. Can be `inf` for infinity norm. + error_if_nonfinite (`bool`, defaults to `False`): + If `True`, an error is thrown if the total norm of the gradients + from `parameters` is `nan`, `inf`, or `-inf`. + foreach (`bool`, defaults to `None`): + Use the faster foreach-based implementation. If `None`, + use the foreach implementation for CUDA and CPU native tensors + and silently fall back to the slow implementation for other device types. + pp_mesh (`torch.distributed.device_mesh.DeviceMesh`, defaults to `None`): + Pipeline parallel device mesh. If not `None`, will reduce gradient norm across PP stages. + + Returns: + `torch.Tensor`: + Total norm of the gradients + """ + grads = [p.grad for p in parameters if p.grad is not None] + + # TODO(aryan): Wait for next Pytorch release to use `torch.nn.utils.get_total_norm` + # total_norm = torch.nn.utils.get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + + # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. + # We can simply reduce the DTensor to get the total norm in this tensor's process group + # and then convert it to a local tensor. + # It has two purposes: + # 1. to make sure the total norm is computed correctly when PP is used (see below) + # 2. to return a reduced total_norm tensor whose .item() would return the correct value + if isinstance(total_norm, torch.distributed.tensor.DTensor): + # Will reach here if any non-PP parallelism is used. + # If only using PP, total_norm will be a local tensor. + total_norm = total_norm.full_tensor() + + if pp_mesh is not None: + raise NotImplementedError("Pipeline parallel is not supported") + if math.isinf(norm_type): + dist.all_reduce(total_norm, + op=dist.ReduceOp.MAX, + group=pp_mesh.get_group()) + else: + total_norm **= norm_type + dist.all_reduce(total_norm, + op=dist.ReduceOp.SUM, + group=pp_mesh.get_group()) + total_norm **= 1.0 / norm_type + + _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) + return total_norm + + +@torch.no_grad() +def _clip_grads_with_norm_( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + total_norm: torch.Tensor, + foreach: bool | None = None, +) -> None: + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + grads = [p.grad for p in parameters if p.grad is not None] + max_norm = float(max_norm) + if len(grads) == 0: + return + grouped_grads: dict[tuple[torch.device, torch.dtype], + tuple[list[list[torch.Tensor]], + list[int]]] = (_group_tensors_by_device_and_dtype( + [grads])) # type: ignore[assignment] + + clip_coef = max_norm / (total_norm + 1e-6) + + # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so + # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization + # when the gradients do not reside in CPU memory. + clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + for (device, _), ([device_grads], _) in grouped_grads.items(): + if (foreach is None and _has_foreach_support(device_grads, device)) or ( + foreach and _device_has_foreach_support(device)): + torch._foreach_mul_(device_grads, clip_coef_clamped.to(device)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + clip_coef_clamped_device = clip_coef_clamped.to(device) + for g in device_grads: + g.mul_(clip_coef_clamped_device) + + +def _get_total_norm( + tensors: torch.Tensor | list[torch.Tensor], + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, +) -> torch.Tensor: + tensors = [tensors] if isinstance(tensors, torch.Tensor) else list(tensors) + norm_type = float(norm_type) + if len(tensors) == 0: + return torch.tensor(0.0) + first_device = tensors[0].device + grouped_tensors: dict[tuple[torch.device, torch.dtype], + tuple[list[list[torch.Tensor]], list[int]]] = ( + _group_tensors_by_device_and_dtype( + [tensors] # type: ignore[list-item] + )) # type: ignore[assignment] + + norms: list[torch.Tensor] = [] + for (device, _), ([device_tensors], _) in grouped_tensors.items(): + local_tensors = [ + t.to_local() + if isinstance(t, torch.distributed.tensor.DTensor) else t + for t in device_tensors + ] + if (foreach is None and _has_foreach_support(local_tensors, device) + ) or (foreach and _device_has_foreach_support(device)): + norms.extend(torch._foreach_norm(local_tensors, norm_type)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + norms.extend( + [torch.linalg.vector_norm(g, norm_type) for g in local_tensors]) + + total_norm = torch.linalg.vector_norm( + torch.stack([norm.to(first_device) for norm in norms]), norm_type) + + if error_if_nonfinite and torch.logical_or(total_norm.isnan(), + total_norm.isinf()): + raise RuntimeError( + f"The total norm of order {norm_type} for gradients from " + "`parameters` is non-finite, so it cannot be clipped. To disable " + "this error and scale the gradients by the non-finite norm anyway, " + "set `error_if_nonfinite=False`") + return total_norm + + +def _get_foreach_kernels_supported_devices() -> list[str]: + r"""Return the device type list that supports foreach kernels.""" + return ["cuda", "xpu", torch._C._get_privateuse1_backend_name()] + + +@torch.no_grad() +def _group_tensors_by_device_and_dtype( + tensorlistlist: list[list[torch.Tensor | None]], + with_indices: bool = False, +) -> dict[tuple[torch.device, torch.dtype], tuple[ + list[list[torch.Tensor | None]], list[int]]]: + return torch._C._group_tensors_by_device_and_dtype( # type: ignore[no-any-return] + tensorlistlist, with_indices) + + +def _device_has_foreach_support(device: torch.device) -> bool: + return device.type in (_get_foreach_kernels_supported_devices() + + ["cpu"]) and not torch.jit.is_scripting() + + +def _has_foreach_support(tensors: list[torch.Tensor], + device: torch.device) -> bool: + return _device_has_foreach_support(device) and all( + t is None or type(t) in [torch.Tensor] for t in tensors) + + +def custom_to_hf_state_dict( + state_dict: dict[str, Any] | Iterator[tuple[str, torch.Tensor]], + reverse_param_names_mapping: dict[str, tuple[str, int, + int]]) -> dict[str, Any]: + """ + Convert trainer's custom model format to diffusers format using reverse_param_names_mapping. + + Args: + state_dict: State dict in trainer's custom format + reverse_param_names_mapping: Reverse mapping from trainer's custom format to diffusers format + + Returns: + State dict in diffusers format + """ + assert len( + reverse_param_names_mapping) > 0, "reverse_param_names_mapping is empty" + if isinstance(state_dict, Iterator): + state_dict = dict(state_dict) + new_state_dict = {} + # Group parameters that need to be split (merged parameters) + merge_groups: dict[str, list[tuple[str, int, int]]] = {} + + # First pass: collect all merge groups + for training_key, ( + diffusers_key, merge_index, + num_params_to_merge) in reverse_param_names_mapping.items(): + if merge_index is not None: + # This is a merged parameter that needs to be split + if training_key not in merge_groups: + merge_groups[training_key] = [] + merge_groups[training_key].append( + (diffusers_key, merge_index, num_params_to_merge)) + + # Second pass: handle merged parameters by splitting them + used_keys = set() + for training_key, splits in merge_groups.items(): + if training_key in state_dict: + v = state_dict[training_key] + # Sort by merge_index to ensure correct order + splits.sort(key=lambda x: x[1]) + total = splits[0][2] + split_size = v.shape[0] // total + split_tensors = torch.split(v, split_size, dim=0) + + for diffusers_key, split_index, _ in splits: + new_state_dict[diffusers_key] = split_tensors[split_index] + used_keys.add(training_key) + + # Third pass: handle regular parameters (direct mappings) + for training_key, v in state_dict.items(): + if training_key in used_keys: + continue + + if training_key in reverse_param_names_mapping: + diffusers_key, merge_index, _ = reverse_param_names_mapping[ + training_key] + if merge_index is None: + # Direct mapping + new_state_dict[diffusers_key] = v + else: + # No mapping found, keep as is + new_state_dict[training_key] = v + + return new_state_dict + + +def pred_noise_to_pred_video(pred_noise: torch.Tensor, + noise_input_latent: torch.Tensor, + timestep: torch.Tensor, + scheduler: Any) -> torch.Tensor: + """ + Convert predicted noise to clean latent. + """ + timestep = timestep.expand(noise_input_latent.shape[0]) + dtype = pred_noise.dtype + device = pred_noise.device + pred_noise = pred_noise.float().to(device) + noise_input_latent = noise_input_latent.float().to(device) + # sigmas = scheduler.sigmas.float().to(device) + # timesteps = scheduler.timesteps.float().to(device) + # timestep_id = torch.argmin( + # (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + # sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + sigma_t = timestep / 1000.0 + sigma_t = sigma_t.reshape(-1, 1, 1, 1) + # if get_sp_parallel_rank() == 0: + # print('sigma:', sigma_t) + pred_video = noise_input_latent.to(torch.float32) - sigma_t.to(torch.float32) * pred_noise.to(torch.float32) + return pred_video.to(dtype) + + +def shift_timestep(timestep: torch.Tensor, shift: float, + num_train_timestep: float) -> torch.Tensor: + if shift == 1: + return timestep + t = timestep / num_train_timestep + denominator = 1 + (shift - 1) * t + return num_train_timestep * (shift * t / denominator) + + +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Modified from diffusers/src/diffusers/optimization.py +"""PyTorch optimization for diffusion models.""" + + +class SchedulerType(Enum): + LINEAR = "linear" + COSINE = "cosine" + COSINE_WITH_RESTARTS = "cosine_with_restarts" + COSINE_WITH_MIN_LR = "cosine_with_min_lr" + POLYNOMIAL = "polynomial" + CONSTANT = "constant" + CONSTANT_WITH_WARMUP = "constant_with_warmup" + PIECEWISE_CONSTANT = "piecewise_constant" + + +def get_constant_schedule(optimizer: Optimizer, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch) + + +def get_constant_schedule_with_warmup(optimizer: Optimizer, + num_warmup_steps: int, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate + increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1.0, num_warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def get_piecewise_constant_schedule(optimizer: Optimizer, + step_rules: str, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + step_rules (`string`): + The rules for the learning rate. ex: rule_steps="1:10,0.1:20,0.01:30,0.005" it means that the learning rate + if multiple 1 for the first 10 steps, multiple 0.1 for the next 20 steps, multiple 0.01 for the next 30 + steps and multiple 0.005 for the other steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + rules_dict = {} + rule_list = step_rules.split(",") + for rule_str in rule_list[:-1]: + value_str, steps_str = rule_str.split(":") + steps = int(steps_str) + value = float(value_str) + rules_dict[steps] = value + last_lr_multiple = float(rule_list[-1]) + + def create_rules_function( + rules_dict: dict, + last_lr_multiple: float) -> Callable[[int], float]: + + def rule_func(steps: int) -> float: + for step_threshold, lr_multiple in sorted(rules_dict.items()): + if steps < step_threshold: + return lr_multiple + return last_lr_multiple + + return rule_func + + rules_func = create_rules_function(rules_dict, last_lr_multiple) + + return LambdaLR(optimizer, rules_func, last_epoch=last_epoch) + + +def get_linear_schedule_with_warmup(optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after + a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return max( + 0.0, + float(num_training_steps - current_step) / + float(max(1, num_training_steps - num_warmup_steps))) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_schedule_with_warmup(optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float = 0.5, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_periods (`float`, *optional*, defaults to 0.5): + The number of periods of the cosine function in a schedule (the default is to just decrease from the max + value to 0 following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps)) + return max( + 0.0, 0.5 * + (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_schedule_with_min_lr(optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + min_lr_ratio: float = 0.1, + num_cycles: float = 0.5, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to a minimum lr (min_lr_ratio * initial_lr), after a warmup period during which + it increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + min_lr_ratio (`float`, *optional*, defaults to 0.1): + The ratio of minimum learning rate to initial learning rate. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of periods of the cosine function in a schedule. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps)) + # Cosine decay from 1.0 to min_lr_ratio over num_cycles periods + # Use the same formula as standard cosine but ensure minimum is min_lr_ratio instead of 0 + cosine_value = 0.5 * ( + 1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) + # Ensure the value doesn't go below min_lr_ratio + return max(min_lr_ratio, cosine_value) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_with_hard_restarts_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: int = 1, + last_epoch: int = -1) -> LambdaLR: + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases + linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`int`, *optional*, defaults to 1): + The number of hard restarts to use. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps)) + if progress >= 1.0: + return 0.0 + return max( + 0.0, 0.5 * (1.0 + math.cos(math.pi * + ((float(num_cycles) * progress) % 1.0)))) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_polynomial_decay_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + lr_end: float = 1e-7, + power: float = 1.0, + last_epoch: int = -1, +) -> LambdaLR: + """ + Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the + optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + lr_end (`float`, *optional*, defaults to 1e-7): + The end LR. + power (`float`, *optional*, defaults to 1.0): + Power factor. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT + implementation at + https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37 + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + + """ + + lr_init = optimizer.defaults["lr"] + if not (lr_init > lr_end): + raise ValueError( + f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})") + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + elif current_step > num_training_steps: + return lr_end / lr_init # as LambdaLR multiplies by lr_init + else: + lr_range = lr_init - lr_end + decay_steps = num_training_steps - num_warmup_steps + pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps + decay = lr_range * pct_remaining**power + lr_end + return decay / lr_init # as LambdaLR multiplies by lr_init + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +# Type alias for scheduler functions +SchedulerFunction = Callable[..., LambdaLR] + +TYPE_TO_SCHEDULER_FUNCTION: dict[SchedulerType, SchedulerFunction] = { + SchedulerType.LINEAR: get_linear_schedule_with_warmup, + SchedulerType.COSINE: get_cosine_schedule_with_warmup, + SchedulerType.COSINE_WITH_RESTARTS: + get_cosine_with_hard_restarts_schedule_with_warmup, + SchedulerType.COSINE_WITH_MIN_LR: get_cosine_schedule_with_min_lr, + SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup, + SchedulerType.CONSTANT: get_constant_schedule, + SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup, + SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule, +} + + +def get_scheduler( + name: str | SchedulerType, + optimizer: Optimizer, + step_rules: str | None = None, + num_warmup_steps: int | None = None, + num_training_steps: int | None = None, + num_cycles: int = 1, + power: float = 1.0, + min_lr_ratio: float = 0.1, + last_epoch: int = -1, +) -> LambdaLR: + """ + Unified API to get any scheduler from its name. + + Args: + name (`str` or `SchedulerType`): + The name of the scheduler to use. + optimizer (`torch.optim.Optimizer`): + The optimizer that will be used during training. + step_rules (`str`, *optional*): + A string representing the step rules to use. This is only used by the `PIECEWISE_CONSTANT` scheduler. + num_warmup_steps (`int`, *optional*): + The number of warmup steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_training_steps (`int``, *optional*): + The number of training steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_cycles (`int`, *optional*): + The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler. + power (`float`, *optional*, defaults to 1.0): + Power factor. See `POLYNOMIAL` scheduler + min_lr_ratio (`float`, *optional*, defaults to 0.1): + The ratio of minimum learning rate to initial learning rate. Used in `COSINE_WITH_MIN_LR` scheduler. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + """ + name = SchedulerType(name) + schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] + if name == SchedulerType.CONSTANT: + return schedule_func(optimizer, last_epoch=last_epoch) + + if name == SchedulerType.PIECEWISE_CONSTANT: + return schedule_func(optimizer, + step_rules=step_rules, + last_epoch=last_epoch) + + # All other schedulers require `num_warmup_steps` + if num_warmup_steps is None: + raise ValueError( + f"{name} requires `num_warmup_steps`, please provide that argument." + ) + + if name == SchedulerType.CONSTANT_WITH_WARMUP: + return schedule_func(optimizer, + num_warmup_steps=num_warmup_steps, + last_epoch=last_epoch) + + # All other schedulers require `num_training_steps` + if num_training_steps is None: + raise ValueError( + f"{name} requires `num_training_steps`, please provide that argument." + ) + + if name == SchedulerType.COSINE_WITH_RESTARTS: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + last_epoch=last_epoch, + ) + + if name == SchedulerType.POLYNOMIAL: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + power=power, + last_epoch=last_epoch, + ) + + if name == SchedulerType.COSINE_WITH_MIN_LR: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + min_lr_ratio=min_lr_ratio, + last_epoch=last_epoch, + ) + + return schedule_func(optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + last_epoch=last_epoch) diff --git a/trainer/utils.py b/trainer/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c55543ce538c350426a537bbabdd5ccc3ae1db4d --- /dev/null +++ b/trainer/utils.py @@ -0,0 +1,814 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py + +import argparse +import ctypes +import hashlib +import importlib +import importlib.util +import inspect +import json +import math +import os +import signal +import socket +import sys +import tempfile +import threading +import traceback +from collections.abc import Callable +from dataclasses import dataclass, fields, is_dataclass +from functools import lru_cache, partial, wraps +from typing import Any, TypeVar, cast + +import cloudpickle +import filelock +import torch +import yaml +from diffusers.loaders.lora_base import ( + _best_guess_weight_name) # watch out for potetential removal from diffusers +from huggingface_hub import snapshot_download +from remote_pdb import RemotePdb +from torch.distributed.fsdp import MixedPrecisionPolicy + +import trainer.envs as envs +from trainer.logger import init_logger + +logger = init_logger(__name__) + +T = TypeVar("T") + +# TODO(will): used to convert trainer_args.precision to torch.dtype. Find a +# cleaner way to do this. +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +STR_BACKEND_ENV_VAR: str = "TRAINER_ATTENTION_BACKEND" +STR_ATTN_CONFIG_ENV_VAR: str = "TRAINER_ATTENTION_CONFIG" + + +def find_nccl_library() -> str: + """ + We either use the library file specified by the `VLLM_NCCL_SO_PATH` + environment variable, or we find the library file brought by PyTorch. + After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be + found by `ctypes` automatically. + """ + so_file = envs.TRAINER_NCCL_SO_PATH + + # manually load the nccl library + if so_file: + logger.info( + "Found nccl from environment variable TRAINER_NCCL_SO_PATH=%s", + so_file) + else: + if torch.version.cuda is not None: + so_file = "libnccl.so.2" + elif torch.version.hip is not None: + so_file = "librccl.so.1" + else: + raise ValueError("NCCL only supports CUDA and ROCm backends.") + logger.info("Found nccl from library %s", so_file) + return str(so_file) + + +prev_set_stream = torch.cuda.set_stream + +_current_stream = None + + +def _patched_set_stream(stream: torch.cuda.Stream | None) -> None: + global _current_stream + _current_stream = stream + if stream is not None: + prev_set_stream(stream) + + +torch.cuda.set_stream = _patched_set_stream + + +def current_stream() -> torch.cuda.Stream | None: + """ + replace `torch.cuda.current_stream()` with `trainer.utils.current_stream()`. + it turns out that `torch.cuda.current_stream()` is quite expensive, + as it will construct a new stream object at each call. + here we patch `torch.cuda.set_stream` to keep track of the current stream + directly, so that we can avoid calling `torch.cuda.current_stream()`. + + the underlying hypothesis is that we do not call `torch._C._cuda_setStream` + from C/C++ code. + """ + from trainer.platforms import current_platform + + # For non-CUDA platforms, return None + if not current_platform.is_cuda_alike(): + return None + + global _current_stream + if _current_stream is None: + # when this function is called before any stream is set, + # we return the default stream. + # On ROCm using the default 0 stream in combination with RCCL + # is hurting performance. Therefore creating a dedicated stream + # per process + _current_stream = torch.cuda.Stream() if current_platform.is_rocm( + ) else torch.cuda.current_stream() + return _current_stream + + +class StoreBoolean(argparse.Action): + + def __init__(self, + option_strings, + dest, + default=False, + required=False, + help=None): + super().__init__(option_strings=option_strings, + dest=dest, + nargs='?', + const=True, + default=default, + required=required, + help=help) + + def __call__(self, parser, namespace, values, option_string=None): + if values is None: + setattr(namespace, self.dest, True) + elif isinstance(values, str): + if values.lower() == "true": + setattr(namespace, self.dest, True) + elif values.lower() == "false": + setattr(namespace, self.dest, False) + else: + raise ValueError(f"Invalid boolean value: {values}. " + "Expected 'true' or 'false'.") + else: + setattr(namespace, self.dest, bool(values)) + + +class SortedHelpFormatter(argparse.HelpFormatter): + """SortedHelpFormatter that sorts arguments by their option strings.""" + + def add_arguments(self, actions): + actions = sorted(actions, key=lambda x: x.option_strings) + super().add_arguments(actions) + + +class FlexibleArgumentParser(argparse.ArgumentParser): + """ArgumentParser that allows both underscore and dash in names.""" + + def __init__(self, *args, **kwargs) -> None: + # Set the default 'formatter_class' to SortedHelpFormatter + if 'formatter_class' not in kwargs: + kwargs['formatter_class'] = SortedHelpFormatter + super().__init__(*args, **kwargs) + + def parse_args( # type: ignore[override] + self, args=None, namespace=None) -> argparse.Namespace: + if args is None: + args = sys.argv[1:] + + if '--config' in args: + args = self._pull_args_from_config(args) + + # Convert underscores to dashes and vice versa in argument names + processed_args = [] + for arg in args: + if arg.startswith('--'): + if '=' in arg: + key, value = arg.split('=', 1) + key = '--' + key[len('--'):].replace('_', '-') + processed_args.append(f'{key}={value}') + else: + processed_args.append('--' + + arg[len('--'):].replace('_', '-')) + elif arg.startswith('-O') and arg != '-O' and len(arg) == 2: + # allow -O flag to be used without space, e.g. -O3 + processed_args.append('-O') + processed_args.append(arg[2:]) + else: + processed_args.append(arg) + + namespace = super().parse_args(processed_args, namespace) + + # Track which arguments were explicitly provided + namespace._provided = set() + + i = 0 + while i < len(args): + arg = args[i] + if arg.startswith('--'): + # Handle --key=value format + if '=' in arg: + key = arg.split('=')[0][2:].replace('-', '_') + namespace._provided.add(key) + i += 1 + # Handle --key value format + else: + key = arg[2:].replace('-', '_') + namespace._provided.add(key) + # Skip the value if there is one + if i + 1 < len(args) and not args[i + 1].startswith('-'): + i += 2 + else: + i += 1 + else: + i += 1 + + return namespace # type: ignore[no-any-return] + + def _pull_args_from_config(self, args: list[str]) -> list[str]: + """Method to pull arguments specified in the config file + into the command-line args variable. + + The arguments in config file will be inserted between + the argument list. + + example: + ```yaml + port: 12323 + tensor-parallel-size: 4 + ``` + ```python + $: vllm {serve,chat,complete} "facebook/opt-12B" \ + --config config.yaml -tp 2 + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--config', 'config.yaml', + '-tp', '2' + ] + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--port', '12323', + '--tp-size', '4', + '-tp', '2' + ] + ``` + + Please note how the config args are inserted after the sub command. + this way the order of priorities is maintained when these are args + parsed by super(). + """ + assert args.count( + '--config') <= 1, "More than one config file specified!" + + index = args.index('--config') + if index == len(args) - 1: + raise ValueError("No config file specified! \ + Please check your command-line arguments.") + + file_path = args[index + 1] + + config_args = self._load_config_file(file_path) + + # 0th index is for {serve,chat,complete} + # followed by model_tag (only for serve) + # followed by config args + # followed by rest of cli args. + # maintaining this order will enforce the precedence + # of cli > config > defaults + if args[0] == "serve": + if index == 1: + raise ValueError( + "No model_tag specified! Please check your command-line" + " arguments.") + args = [args[0]] + [ + args[1] + ] + config_args + args[2:index] + args[index + 2:] + else: + args = [args[0]] + config_args + args[1:index] + args[index + 2:] + + return args + + def _load_config_file(self, file_path: str) -> list[str]: + """Loads a yaml file and returns the key value pairs as a + flattened list with argparse like pattern + ```yaml + port: 12323 + tensor-parallel-size: 4 + vae_config: + load_encoder: false + load_decoder: true + ``` + returns: + processed_args: list[str] = [ + '--port': '12323', + '--tp-size': '4', + '--vae-config.load-encoder': 'false', + '--vae-config.load-decoder': 'true' + ] + """ + + extension: str = file_path.split('.')[-1] + if extension not in ('yaml', 'yml', 'json'): + raise ValueError( + "Config file must be of a yaml/yml/json type.\ + %s supplied", extension) + + processed_args: list[str] = [] + + config: dict[str, Any] = {} + try: + with open(file_path) as config_file: + config = yaml.safe_load(config_file) + except Exception as ex: + logger.error( + "Unable to read the config file at %s. \ + Make sure path is correct", file_path) + raise ex + + store_boolean_arguments = [ + action.dest for action in self._actions + if isinstance(action, StoreBoolean) + ] + + def process_dict(prefix: str, d: dict[str, Any]): + for key, value in d.items(): + full_key = f"{prefix}.{key}" if prefix else key + + if isinstance(value, + bool) and full_key not in store_boolean_arguments: + if value: + processed_args.append('--' + full_key) + else: + processed_args.append('--' + full_key) + processed_args.append('false') + elif isinstance(value, list): + processed_args.append('--' + full_key) + for item in value: + processed_args.append(str(item)) + elif isinstance(value, dict): + process_dict(full_key, value) + else: + processed_args.append('--' + full_key) + processed_args.append(str(value)) + + process_dict("", config) + + return processed_args + + +def get_lock(model_name_or_path: str): + lock_dir = tempfile.gettempdir() + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +def warn_for_unimplemented_methods(cls: type[T]) -> type[T]: + """ + A replacement for `abc.ABC`. + When we use `abc.ABC`, subclasses will fail to instantiate + if they do not implement all abstract methods. + Here, we only require `raise NotImplementedError` in the + base class, and log a warning if the method is not implemented + in the subclass. + """ + + original_init = cls.__init__ + + def find_unimplemented_methods(self: object): + unimplemented_methods = [] + for attr_name in dir(self): + # bypass inner method + if attr_name.startswith('_'): + continue + + try: + attr = getattr(self, attr_name) + # get the func of callable method + if callable(attr): + attr_func = attr.__func__ + except AttributeError: + continue + src = inspect.getsource(attr_func) + if "NotImplementedError" in src: + unimplemented_methods.append(attr_name) + if unimplemented_methods: + method_names = ','.join(unimplemented_methods) + msg = (f"Methods {method_names} not implemented in {self}") + logger.warning(msg) + + @wraps(original_init) + def wrapped_init(self, *args, **kwargs) -> None: + original_init(self, *args, **kwargs) + find_unimplemented_methods(self) + + type.__setattr__(cls, '__init__', wrapped_init) + return cls + + +def align_to(value: int, alignment: int) -> int: + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) + + +def resolve_obj_by_qualname(qualname: str) -> Any: + """ + Resolve an object by its fully qualified name. + """ + module_name, obj_name = qualname.rsplit(".", 1) + module = importlib.import_module(module_name) + return getattr(module, obj_name) + + +# From vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py +def import_pynvml(): + """ + Historical comments: + + libnvml.so is the library behind nvidia-smi, and + pynvml is a Python wrapper around it. We use it to get GPU + status without initializing CUDA context in the current process. + Historically, there are two packages that provide pynvml: + - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official + wrapper. It is a dependency of Trainer, and is installed when users + install Trainer. It provides a Python module named `pynvml`. + - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper. + Prior to version 12.0, it also provides a Python module `pynvml`, + and therefore conflicts with the official one which is a standalone Python file. + This causes errors when both of them are installed. + Starting from version 12.0, it migrates to a new module + named `pynvml_utils` to avoid the conflict. + It is so confusing that many packages in the community use the + unofficial one by mistake, and we have to handle this case. + For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial + one, and it will cause errors, see the issue + https://github.com/vllm-project/vllm/issues/12847 for example. + After all the troubles, we decide to copy the official `pynvml` + module to our codebase, and use it directly. + """ + import trainer.third_party.pynvml as pynvml + return pynvml + + +def maybe_download_model(model_name_or_path: str, + local_dir: str | None = None, + download: bool = True) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + # If the path exists locally, return it + if os.path.exists(model_name_or_path): + logger.info("Model already exists locally at %s", model_name_or_path) + return model_name_or_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info("Downloading model snapshot from HF Hub for %s...", + model_name_or_path) + with get_lock(model_name_or_path): + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir) + logger.info("Downloaded model to %s", local_path) + return str(local_path) + except Exception as e: + raise ValueError( + f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" + ) from e + + +def maybe_download_lora(model_name_or_path: str, + local_dir: str | None = None, + download: bool = True) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + local_path = maybe_download_model(model_name_or_path, local_dir, download) + weight_name = _best_guess_weight_name(model_name_or_path, + file_extension=".safetensors") + return os.path.join(local_path, weight_name) + + +def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: + """ + Verify that the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary + """ + + # Check for model_index.json which is required for diffusers models + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + raise ValueError( + f"Model directory {model_path} does not contain model_index.json. " + "Only Hugging Face diffusers format is supported.") + + # Check for transformer and vae directories + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + + if not os.path.exists(transformer_dir): + raise ValueError( + f"Model directory {model_path} does not contain a transformer/ directory." + ) + + if not os.path.exists(vae_dir): + raise ValueError( + f"Model directory {model_path} does not contain a vae/ directory.") + + # Load the config + with open(config_path) as f: + config = json.load(f) + + # Verify diffusers version exists + if "_diffusers_version" not in config: + raise ValueError("model_index.json does not contain _diffusers_version") + + logger.info("Diffusers version: %s", config["_diffusers_version"]) + return cast(dict[str, Any], config) + + +def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: + """ + Download and extract just the model_index.json for a Hugging Face model. + + Args: + model_name_or_path: Path or HF Hub model ID + + Returns: + The parsed model_index.json as a dictionary + """ + import tempfile + + from huggingface_hub import hf_hub_download + + # If it's a local path, verify it directly + if os.path.exists(model_name_or_path): + return verify_model_config_and_directory(model_name_or_path) + + # For remote models, download just the model_index.json + try: + with tempfile.TemporaryDirectory() as tmp_dir: + # Download just the model_index.json file + model_index_path = hf_hub_download(repo_id=model_name_or_path, + filename="model_index.json", + local_dir=tmp_dir) + + # Load the model_index.json + with open(model_index_path) as f: + config: dict[str, Any] = json.load(f) + + # Verify it has the required fields + if "_class_name" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _class_name field" + ) + + if "_diffusers_version" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _diffusers_version field" + ) + + # Add the pipeline name for downstream use + config["pipeline_name"] = config["_class_name"] + + logger.info("Downloaded model_index.json for %s, pipeline: %s", + model_name_or_path, config["_class_name"]) + return config + + except Exception as e: + raise ValueError( + f"Failed to download or parse model_index.json for {model_name_or_path}: {e}" + ) from e + + +def update_environment_variables(envs: dict[str, str]): + for k, v in envs.items(): + if k in os.environ and os.environ[k] != v: + logger.warning( + "Overwriting environment variable %s " + "from '%s' to '%s'", k, os.environ[k], v) + os.environ[k] = v + + +def run_method(obj: Any, method: str | bytes | Callable, args: tuple[Any], + kwargs: dict[str, Any]) -> Any: + """ + Run a method of an object with the given arguments and keyword arguments. + If the method is string, it will be converted to a method using getattr. + If the method is serialized bytes and will be deserialized using + cloudpickle. + If the method is a callable, it will be called directly. + """ + if isinstance(method, bytes): + func = partial(cloudpickle.loads(method), obj) + elif isinstance(method, str): + try: + func = getattr(obj, method) + except AttributeError: + raise NotImplementedError(f"Method {method!r} is not" + " implemented.") from None + else: + func = partial(method, obj) # type: ignore + return func(*args, **kwargs) + + +def shallow_asdict(obj) -> dict[str, Any]: + if not is_dataclass(obj): + raise TypeError("Expected dataclass instance") + return {f.name: getattr(obj, f.name) for f in fields(obj)} + + +# TODO: validate that this is fine +def kill_itself_when_parent_died() -> None: + # if sys.platform == "linux": + # sigkill this process when parent worker manager dies + PR_SET_PDEATHSIG = 1 + import platform + if platform.system() == "Linux": + libc = ctypes.CDLL("libc.so.6") + libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) + # elif platform.system() == "Darwin": + # libc = ctypes.CDLL("libc.dylib") + # logger.warning("kill_itself_when_parent_died is only supported in linux.") + else: + logger.warning( + "kill_itself_when_parent_died is only supported in linux.") + + +def get_exception_traceback() -> str: + etype, value, tb = sys.exc_info() + err_str = "".join(traceback.format_exception(etype, value, tb)) + return err_str + + +class TypeBasedDispatcher: + + def __init__(self, mapping: list[tuple[type, Callable]]): + self._mapping = mapping + + def __call__(self, obj: Any): + for ty, fn in self._mapping: + if isinstance(obj, ty): + return fn(obj) + raise ValueError(f"Invalid object: {obj}") + + +# For non-torch.distributed debugging +def remote_breakpoint() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("localhost", 0)) # Let the OS pick an ephemeral port. + port = s.getsockname()[1] + RemotePdb(host="localhost", port=port).set_trace() + + +@dataclass +class MixedPrecisionState: + param_dtype: torch.dtype | None = None + reduce_dtype: torch.dtype | None = None + output_dtype: torch.dtype | None = None + compute_dtype: torch.dtype | None = None + mp_policy: MixedPrecisionPolicy | None = None + + +# Thread-local storage for mixed precision state +_mixed_precision_state = threading.local() + + +def get_mixed_precision_state() -> MixedPrecisionState: + """Get the current mixed precision state.""" + if not hasattr(_mixed_precision_state, 'state'): + raise ValueError("Mixed precision state not set") + return cast(MixedPrecisionState, _mixed_precision_state.state) + + +def set_mixed_precision_policy( + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + output_dtype: torch.dtype | None = None, + mp_policy: MixedPrecisionPolicy | None = None, +): + """Set mixed precision policy globally. + + Args: + param_dtype: Parameter dtype used for training + reduce_dtype: Reduction dtype used for gradients + output_dtype: Optional output dtype + """ + state = MixedPrecisionState( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + _mixed_precision_state.state = state + + +def get_compute_dtype() -> torch.dtype: + """Get the current compute dtype from mixed precision policy. + + Returns: + torch.dtype: The compute dtype to use, defaults to get_default_dtype() if no policy set + """ + if not hasattr(_mixed_precision_state, 'state'): + return torch.get_default_dtype() + else: + state = get_mixed_precision_state() + return state.param_dtype + + +def dict_to_3d_list( + mask_strategy: dict[str, Any] | None = None, + t_max: int | None = None, + l_max: int | None = None, + h_max: int | None = None, +) -> list[list[list[torch.Tensor | None]]]: + """ + Convert a dictionary of mask indices to a 3D list of tensors. + Args: + mask_strategy: keys are "t_l_h", values are torch.Tensor masks. + t_max, l_max, h_max: if provided (all three), force the output shape to (t_max, l_max, h_max). + If all three are None, infer shape from the data. + """ + # Case 1: no data, but fixed shape requested + if mask_strategy is None: + assert t_max is not None and l_max is not None and h_max is not None, ( + "If mask_strategy is None, you must provide t_max, l_max, and h_max" + ) + return [[[None for _ in range(h_max)] for _ in range(l_max)] + for _ in range(t_max)] + + # Parse all keys into integer tuples + indices = [tuple(map(int, key.split("_"))) for key in mask_strategy] + + # Decide on dimensions + if t_max is None and l_max is None and h_max is None: + # fully dynamic: infer from data + max_timesteps_idx = max(t for t, _, _ in indices) + 1 + max_layer_idx = max(l for _, l, _ in indices) + 1 # noqa: E741 + max_head_idx = max(h for _, _, h in indices) + 1 + else: + # require all three to be provided + assert t_max is not None and l_max is not None and h_max is not None, ( + "Either supply none of (t_max, l_max, h_max) to infer dimensions, " + "or supply all three to fix the shape.") + max_timesteps_idx = t_max + max_layer_idx = l_max + max_head_idx = h_max + + # Preallocate + result = [[[None for _ in range(max_head_idx)] + for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)] + + # Fill in, skipping any out-of-bounds entries + for key, value in mask_strategy.items(): + t, l, h = map(int, key.split("_")) # noqa: E741 + if 0 <= t < max_timesteps_idx and 0 <= l < max_layer_idx and 0 <= h < max_head_idx: + result[t][l][h] = value + # else: silently ignore any key that doesn't fit + + return result + + +def set_random_seed(seed: int) -> None: + from trainer.platforms import current_platform + current_platform.seed_everything(seed) + + +@lru_cache(maxsize=1) +def is_vsa_available() -> bool: + return importlib.util.find_spec("vsa") is not None diff --git a/trainer/version.py b/trainer/version.py new file mode 100644 index 0000000000000000000000000000000000000000..1276d0254ffc2b5e4fcf48cd868123238d4ad06f --- /dev/null +++ b/trainer/version.py @@ -0,0 +1 @@ +__version__ = "0.1.5" diff --git a/trainer/worker/__init__.py b/trainer/worker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..85413d3ca5d3f378b4fa729b510c799a4e4a91b9 --- /dev/null +++ b/trainer/worker/__init__.py @@ -0,0 +1,5 @@ +from .executor import Executor +from .gpu_worker import run_worker_process +from .multiproc_executor import MultiprocExecutor + +__all__ = ["Executor", "run_worker_process", "MultiprocExecutor"] \ No newline at end of file diff --git a/trainer/worker/executor.py b/trainer/worker/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..6ddb0eb1888d87e08189689b4094971c5d375a32 --- /dev/null +++ b/trainer/worker/executor.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, TypeVar, cast + +from trainer.trainer_args import TrainerArgs +from trainer.pipelines import ForwardBatch +from trainer.utils import init_logger + +logger = init_logger(__name__) + +_R = TypeVar("_R") + + +class Executor(ABC): + + def __init__(self, trainer_args: TrainerArgs): + self.trainer_args = trainer_args + + self._init_executor() + + @abstractmethod + def _init_executor(self) -> None: + raise NotImplementedError + + @classmethod + def get_class(cls, trainer_args: TrainerArgs) -> type["Executor"]: + if trainer_args.distributed_executor_backend == "mp": + from trainer.worker.multiproc_executor import MultiprocExecutor + return cast(type["Executor"], MultiprocExecutor) + else: + raise ValueError( + f"Unsupported distributed executor backend: {trainer_args.distributed_executor_backend}" + ) + + def execute_forward( + self, + forward_batch: ForwardBatch, + trainer_args: TrainerArgs, + ) -> ForwardBatch: + outputs: list[dict[str, + Any]] = self.collective_rpc("execute_forward", + kwargs={ + "forward_batch": + forward_batch, + "trainer_args": + trainer_args + }) + return cast(ForwardBatch, outputs[0]["output_batch"]) + + @abstractmethod + def set_lora_adapter(self, + lora_nickname: str, + lora_path: str | None = None) -> None: + """ + Set the LoRA adapter for the workers. + """ + raise NotImplementedError + + @abstractmethod + def collective_rpc(self, + method: str | Callable[..., _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None) -> list[_R]: + """ + Execute an RPC call on all workers. + + Args: + method: Name of the worker method to execute, or a callable that + is serialized and sent to all workers to execute. + + If the method is a callable, it should accept an additional + `self` argument, in addition to the arguments passed in `args` + and `kwargs`. The `self` argument will be the worker object. + timeout: Maximum time in seconds to wait for execution. Raises a + :exc:`TimeoutError` on timeout. `None` means wait indefinitely. + args: Positional arguments to pass to the worker method. + kwargs: Keyword arguments to pass to the worker method. + + Returns: + A list containing the results from each worker. + + Note: + It is recommended to use this API to only pass control messages, + and set up data-plane communication to pass data. + """ + raise NotImplementedError + + @abstractmethod + def shutdown(self) -> None: + """ + Shutdown the executor. + """ + raise NotImplementedError diff --git a/trainer/worker/gpu_worker.py b/trainer/worker/gpu_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..3fa92a00409d146184921fc22932a11072d1e96e --- /dev/null +++ b/trainer/worker/gpu_worker.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +import contextlib +import faulthandler +import multiprocessing as mp +import os +import signal +import sys +from multiprocessing.connection import Connection +from typing import Any, TextIO, cast + +import psutil +import torch + +import trainer.envs as envs +from trainer.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel) +from trainer.distributed.parallel_state import get_local_torch_device +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines import ForwardBatch, build_pipeline +from trainer.platforms import current_platform +from trainer.utils import (get_exception_traceback, + kill_itself_when_parent_died) + +logger = init_logger(__name__) + +# ANSI color codes +CYAN = '\033[1;36m' +RESET = '\033[0;0m' + + +class Worker: + + def __init__(self, trainer_args: TrainerArgs, local_rank: int, + rank: int, pipe: Connection, master_port: int): + self.trainer_args = trainer_args + self.local_rank = local_rank + self.rank = rank + # TODO(will): don't hardcode this + self.distributed_init_method = "env://" + self.pipe = pipe + self.master_port = master_port + self.init_device() + + # Init request dispatcher + # TODO(will): add request dispatcher: use TypeBasedDispatcher from + # utils.py + # self._request_dispatcher = TypeBasedDispatcher( + # [ + # (RpcReqInput, self.handle_rpc_request), + # (GenerateRequest, self.handle_generate_request), + # (ExpertDistributionReq, self.expert_distribution_handle), + # ] + # ) + + def init_device(self) -> None: + """Initialize the device for the worker.""" + + # torch.distributed.all_reduce does not free the input tensor until + # the synchronization point. This causes the memory usage to grow + # as the number of all_reduce calls increases. This env var disables + # this behavior. + # Related issue: + # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + # This env var set by Ray causes exceptions with graph building. + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + + # Platform-agnostic device initialization + self.device = get_local_torch_device() + + # _check_if_gpu_supports_dtype(self.model_config.dtype) + if current_platform.is_cuda_alike(): + self.init_gpu_memory = torch.cuda.mem_get_info()[0] + else: + # For MPS, we can't get memory info the same way + self.init_gpu_memory = 0 + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(self.master_port) + os.environ["LOCAL_RANK"] = str(self.local_rank) + os.environ["RANK"] = str(self.rank) + os.environ["WORLD_SIZE"] = str(self.trainer_args.num_gpus) + + # Initialize the distributed environment. + maybe_init_distributed_environment_and_model_parallel( + self.trainer_args.tp_size, self.trainer_args.sp_size) + + self.pipeline = build_pipeline(self.trainer_args) + + def execute_forward(self, forward_batch: ForwardBatch, + trainer_args: TrainerArgs) -> ForwardBatch: + output_batch = self.pipeline.forward(forward_batch, self.trainer_args) + return cast(ForwardBatch, output_batch) + + def set_lora_adapter(self, + lora_nickname: str, + lora_path: str | None = None) -> None: + self.pipeline.set_lora_adapter(lora_nickname, lora_path) + + def shutdown(self) -> dict[str, Any]: + """Gracefully shut down the worker process""" + logger.info("Worker %d shutting down...", + self.rank, + local_main_process_only=False) + # Clean up resources + if hasattr(self, 'pipeline') and self.pipeline is not None: + # Clean up pipeline resources if needed + pass + + # Destroy the distributed environment + cleanup_dist_env_and_memory(shutdown_ray=False) + + logger.info("Worker %d shutdown complete", + self.rank, + local_main_process_only=False) + return {"status": "shutdown_complete"} + + def event_loop(self) -> None: + """Event loop for the worker.""" + logger.info("Worker %d starting event loop...", + self.rank, + local_main_process_only=False) + while True: + try: + recv_rpc = self.pipe.recv() + method_name = recv_rpc.get('method') + + # Handle shutdown request + if method_name == 'shutdown': + response = self.shutdown() + with contextlib.suppress(Exception): + self.pipe.send(response) + break # Exit the loop + + # Handle regular RPC calls + if method_name == 'execute_forward': + forward_batch = recv_rpc['kwargs']['forward_batch'] + trainer_args = recv_rpc['kwargs']['trainer_args'] + output_batch = self.execute_forward(forward_batch, + trainer_args) + logging_info = None + if envs.TRAINER_STAGE_LOGGING: + logging_info = output_batch.logging_info + self.pipe.send({ + "output_batch": output_batch.output.cpu(), + "logging_info": logging_info + }) + elif method_name == 'set_lora_adapter': + lora_nickname = recv_rpc['kwargs']['lora_nickname'] + lora_path = recv_rpc['kwargs']['lora_path'] + self.set_lora_adapter(lora_nickname, lora_path) + logger.info("Worker %d set LoRA adapter %s with path %s", + self.rank, lora_nickname, lora_path) + self.pipe.send({"status": "lora_adapter_set"}) + else: + # Handle other methods dynamically if needed + args = recv_rpc.get('args', ()) + kwargs = recv_rpc.get('kwargs', {}) + if hasattr(self, method_name): + method = getattr(self, method_name) + result = method(*args, **kwargs) + self.pipe.send(result) + else: + self.pipe.send( + {"error": f"Unknown method: {method_name}"}) + except KeyboardInterrupt: + logger.error( + "Worker %d in loop received KeyboardInterrupt, aborting forward pass", + self.rank) + try: + self.pipe.send( + {"error": "Operation aborted by KeyboardInterrupt"}) + logger.info("Worker %d sent error response after interrupt", + self.rank) + except Exception as e: + logger.error("Worker %d failed to send error response: %s", + self.rank, str(e)) + continue + + +def run_worker_process(trainer_args: TrainerArgs, local_rank: int, + rank: int, pipe: Connection, master_port: int): + # Add process-specific prefix to stdout and stderr + process_name = mp.current_process().name + pid = os.getpid() + _add_prefix(sys.stdout, process_name, pid) + _add_prefix(sys.stderr, process_name, pid) + + # Config the process + kill_itself_when_parent_died() + faulthandler.enable() + parent_process = psutil.Process().parent() + + logger.info("Worker %d initializing...", + rank, + local_main_process_only=False) + + try: + worker = Worker(trainer_args, local_rank, rank, pipe, master_port) + logger.info("Worker %d sending ready", rank) + pipe.send({ + "status": "ready", + "local_rank": local_rank, + }) + worker.event_loop() + except Exception: + traceback = get_exception_traceback() + logger.error("Worker %d hit an exception: %s", rank, traceback) + parent_process.send_signal(signal.SIGQUIT) + + +def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None: + """Prepend each output line with process-specific prefix""" + + prefix = f"{CYAN}({worker_name} pid={pid}){RESET} " + file_write = file.write + + def write_with_prefix(s: str): + if not s: + return + if file.start_new_line: # type: ignore[attr-defined] + file_write(prefix) + idx = 0 + while (next_idx := s.find('\n', idx)) != -1: + next_idx += 1 + file_write(s[idx:next_idx]) + if next_idx == len(s): + file.start_new_line = True # type: ignore[attr-defined] + return + file_write(prefix) + idx = next_idx + file_write(s[idx:]) + file.start_new_line = False # type: ignore[attr-defined] + + file.start_new_line = True # type: ignore[attr-defined] + file.write = write_with_prefix # type: ignore[method-assign] diff --git a/trainer/worker/multiproc_executor.py b/trainer/worker/multiproc_executor.py new file mode 100644 index 0000000000000000000000000000000000000000..ffd784c9001590a373dd581b7c97bebf068bda8d --- /dev/null +++ b/trainer/worker/multiproc_executor.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +import atexit +import contextlib +import multiprocessing as mp +import os +import signal +import socket +import time +from collections.abc import Callable +from multiprocessing.process import BaseProcess +from typing import Any + +import trainer.envs as envs +from trainer.trainer_args import TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines.pipeline_batch_info import ForwardBatch +from trainer.worker.executor import Executor +from trainer.worker.gpu_worker import run_worker_process + +logger = init_logger(__name__) + + +class MultiprocExecutor(Executor): + + def _init_executor(self) -> None: + self.world_size = self.trainer_args.num_gpus + self.shutting_down = False + + # this will force the use of the `spawn` multiprocessing start if cuda + # is initialized + mp.set_start_method("spawn", force=True) + + self.workers: list[BaseProcess] = [] + self.worker_pipes = [] + + # Check if master_port is provided in trainer_args + if hasattr( + self.trainer_args, + 'master_port') and self.trainer_args.master_port is not None: + self.master_port = self.trainer_args.master_port + logger.info("Using provided master port: %s", self.master_port) + else: + # Auto-find available port + import random + for port in range(29503 + random.randint(0, 10000), 65535): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + if s.connect_ex(('localhost', port)) != 0: + self.master_port = port + break + else: + raise ValueError("No unused port found to use as master port") + logger.info("Auto-selected master port: %s", self.master_port) + + # Create pipes and start workers + for rank in range(self.world_size): + executor_pipe, worker_pipe = mp.Pipe(duplex=True) + self.worker_pipes.append(executor_pipe) + worker = mp.Process(target=run_worker_process, + name=f"FVWorkerProc-{rank}", + kwargs=dict(trainer_args=self.trainer_args, + local_rank=rank, + rank=rank, + pipe=worker_pipe, + master_port=self.master_port)) + worker.start() + self.workers.append(worker) + + # Wait for all workers to be ready + for idx, pipe in enumerate(self.worker_pipes): + data = pipe.recv() + if data["status"] != "ready" or data["local_rank"] != idx: + raise RuntimeError(f"Worker {idx} failed to start") + logger.info("%d workers ready", self.world_size) + + # Register shutdown on exit + atexit.register(self.shutdown) + + def execute_forward(self, forward_batch: ForwardBatch, + trainer_args: TrainerArgs) -> ForwardBatch: + responses = self.collective_rpc("execute_forward", + kwargs={ + "forward_batch": forward_batch, + "trainer_args": trainer_args + }) + output = responses[0]["output_batch"] + + logging_info = None + if envs.TRAINER_STAGE_LOGGING: + logging_info = responses[0]["logging_info"] + else: + logging_info = None + + result_batch = ForwardBatch(data_type=forward_batch.data_type, + output=output, + logging_info=logging_info) + + return result_batch + + def set_lora_adapter(self, + lora_nickname: str, + lora_path: str | None = None) -> None: + responses = self.collective_rpc("set_lora_adapter", + kwargs={ + "lora_nickname": lora_nickname, + "lora_path": lora_path + }) + for i, response in enumerate(responses): + if response["status"] != "lora_adapter_set": + raise RuntimeError( + f"Worker {i} failed to set LoRA adapter to {lora_path}") + + def collective_rpc(self, + method: str | Callable, + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None) -> list[Any]: + kwargs = kwargs or {} + + try: + for pipe in self.worker_pipes: + pipe.send({"method": method, "args": args, "kwargs": kwargs}) + + responses = [] + for pipe in self.worker_pipes: + response = pipe.recv() + responses.append(response) + return responses + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e + except KeyboardInterrupt as e: + # if we catch a KeyboardInterrupt, user wants to stop the execution. + # we need to send a signal to all workers to stop. + logger.info( + "Received KeyboardInterrupt, sending SIGINT to all workers") + for worker in self.workers: + if worker.pid is not None: + os.kill(worker.pid, signal.SIGINT) + raise e + except Exception as e: + raise e + + def shutdown(self) -> None: + """Properly shut down the executor and its workers""" + if hasattr(self, 'shutting_down') and self.shutting_down: + return # Prevent multiple shutdown calls + + logger.info("Shutting down MultiprocExecutor...") + self.shutting_down = True + + # First try gentle termination + try: + # Send termination message to all workers + for pipe in self.worker_pipes: + with contextlib.suppress(Exception): + pipe.send({"method": "shutdown", "args": (), "kwargs": {}}) + + # Give workers some time to exit gracefully + start_time = time.perf_counter() + while time.perf_counter() - start_time < 5.0: # 5 seconds timeout + if all(not worker.is_alive() for worker in self.workers): + break + time.sleep(0.1) + + # Force terminate any remaining workers + for worker in self.workers: + if worker.is_alive(): + worker.terminate() + + # Final timeout for terminate + start_time = time.perf_counter() + while time.perf_counter() - start_time < 2.0: # 2 seconds timeout + if all(not worker.is_alive() for worker in self.workers): + break + time.sleep(0.1) + + # Kill if still alive + for worker in self.workers: + if worker.is_alive(): + worker.kill() + worker.join(timeout=1.0) + + except Exception as e: + logger.error("Error during shutdown: %s", e) + # Last resort, try to kill all workers + for worker in self.workers: + with contextlib.suppress(Exception): + if worker.is_alive(): + worker.kill() + + # Clean up pipes + for pipe in self.worker_pipes: + with contextlib.suppress(Exception): + pipe.close() + + self.workers = [] + self.worker_pipes = [] + logger.info("MultiprocExecutor shutdown complete") + + def __del__(self): + """Ensure cleanup on garbage collection""" + self.shutdown() + + def __enter__(self): + """Support for context manager protocol""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Ensure cleanup when exiting context""" + self.shutdown() diff --git a/trainer/workflow/__init__.py b/trainer/workflow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/workflow/preprocess/__init__.py b/trainer/workflow/preprocess/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/trainer/workflow/preprocess/components.py b/trainer/workflow/preprocess/components.py new file mode 100644 index 0000000000000000000000000000000000000000..25aeff5a2e24e08a696663ab558f46169e367bc7 --- /dev/null +++ b/trainer/workflow/preprocess/components.py @@ -0,0 +1,397 @@ +import dataclasses +import gc +import multiprocessing +import os +import random +from collections.abc import Callable +from concurrent.futures import ProcessPoolExecutor +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from trainer.logger import init_logger +from trainer.pipelines.pipeline_batch_info import PreprocessBatch + +logger = init_logger(__name__) + + +class PreprocessingDataValidator: + + def __init__(self, + max_height: int = 1024, + max_width: int = 1024, + max_h_div_w_ratio: float = 17 / 16, + min_h_div_w_ratio: float = 8 / 16, + num_frames: int = 16, + train_fps: int = 24, + speed_factor: float = 1.0, + video_length_tolerance_range: float = 5.0, + drop_short_ratio: float = 0.0, + hw_aspect_threshold: float = 1.5): + self.max_height = max_height + self.max_width = max_width + self.max_h_div_w_ratio = max_h_div_w_ratio + self.min_h_div_w_ratio = min_h_div_w_ratio + self.num_frames = num_frames + self.train_fps = train_fps + self.speed_factor = speed_factor + self.video_length_tolerance_range = video_length_tolerance_range + self.drop_short_ratio = drop_short_ratio + self.hw_aspect_threshold = hw_aspect_threshold + self.validators: dict[str, Callable[[dict[str, Any]], bool]] = {} + self.filter_counts: dict[str, int] = {} + + self.num_items_before_filtering = 0 + self.num_items_after_filtering = 0 + + self.register_validators() + + def register_validators(self) -> None: + self.add_validator("data_type_validator", self._validate_data_type) + self.add_validator("resolution_validator", self._validate_resolution) + self.add_validator("frame_sampling_validator", + self._validate_frame_sampling) + + def add_validator(self, name: str, validator: Callable[[dict[str, Any]], + bool]) -> None: + self.validators[name] = validator + self.filter_counts[name] = 0 + + def __call__(self, batch: dict[str, Any]) -> bool: + """ + Validate whether the preprocessing data batch is valid. + """ + self.num_items_before_filtering += 1 + + for name, validator in self.validators.items(): + if not validator(batch): + self.filter_counts[name] += 1 + return False + + self.num_items_after_filtering += 1 + return True + + def _validate_data_type(self, batch: dict[str, Any]) -> bool: + """Validate basic validity of data items""" + return not (batch["caption"] is None or batch["caption"] == "" + or batch["fps"] is None or batch["fps"] <= 0 + or batch["num_frames"] is None or batch["num_frames"] <= 0) + + def _validate_resolution(self, batch: dict[str, Any]) -> bool: + """Validate resolution constraints""" + + aspect = self.max_height / self.max_width + if batch["resolution"] is not None: + height = batch["resolution"].get("height", None) + width = batch["resolution"].get("width", None) + + if height is None or width is None: + return False + + return self._filter_resolution( + height, + width, + max_h_div_w_ratio=self.hw_aspect_threshold * aspect, + min_h_div_w_ratio=1 / self.hw_aspect_threshold * aspect, + ) + + def _filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float, + min_h_div_w_ratio: float) -> bool: + """Filter based on aspect ratio""" + return (min_h_div_w_ratio <= h / w <= max_h_div_w_ratio) and ( + self.min_h_div_w_ratio <= h / w <= self.max_h_div_w_ratio) + + def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool: + """Validate frame sampling constraints""" + + if (batch["num_frames"] / batch["fps"] + > self.video_length_tolerance_range * + (self.num_frames / self.train_fps * self.speed_factor)): + return False + + frame_interval = batch["fps"] / self.train_fps + start_frame_idx = 0 + frame_indices = np.arange(start_frame_idx, batch["num_frames"], + frame_interval).astype(int) + return not (len(frame_indices) < self.num_frames + and random.random() < self.drop_short_ratio) + + def log_validation_stats(self): + info = "" + for name, count in self.filter_counts.items(): + info += f"failed in {name}: {count}, " + info += f"number of items before filtering: {self.num_items_before_filtering}, " + info += f"number of items after filtering: {self.num_items_after_filtering}" + + logger.info(info) + + +class VideoForwardBatchBuilder: + + def __init__(self, seed: int): + self.seed = seed + + def __call__(self, batch: list) -> PreprocessBatch: + forward_batch = PreprocessBatch( + video_loader=[item["video"] for item in batch], + video_file_name=[item["name"] for item in batch], + height=[item["resolution"]["height"] for item in batch], + width=[item["resolution"]["width"] for item in batch], + fps=[item["fps"] for item in batch], + num_frames=[item["num_frames"] for item in batch], + prompt=[item["caption"] for item in batch], + prompt_attention_mask=[], + data_type="video", + generator=torch.Generator("cpu").manual_seed(self.seed), + ) + return forward_batch + + +class ParquetDatasetSaver: + """Component for saving and writing Parquet datasets""" + + def __init__(self, + flush_frequency: int, + samples_per_file: int, + schema_fields: list[str], + record_creator: Callable[..., list[dict[str, Any]]], + file_writer_fn: Callable | None = None): + """ + Initialize ParquetDatasetSaver + + Args: + schema_fields: schema fields list + record_creator: Function for creating records + file_writer_fn: Function for writing records to files, uses default implementation if None + """ + self.flush_frequency = flush_frequency + self.samples_per_file = samples_per_file + self.schema_fields = schema_fields + self.create_records_from_batch = record_creator + self.file_writer_fn: Callable[ + [tuple], int] = file_writer_fn or self._default_file_writer_fn + self.all_tables: list[pa.Table] = [] + self.num_processed_samples: int = 0 + self.num_saved_files: int = 0 + + def save_and_write_parquet_batch( + self, + batch: PreprocessBatch, + output_dir: str, + extra_features: dict[str, Any] | None = None) -> None: + """ + Save and write Parquet dataset batch + + Args: + batch: PreprocessBatch containing video and metadata information + output_dir: Output directory + extra_features: Extra features + + Returns: + Number of processed samples + """ + assert isinstance(batch.latents, torch.Tensor) + assert isinstance(batch.prompt_embeds, list) + assert isinstance(batch.prompt_attention_mask, list) + + # Process non-padded embeddings (if needed) + if batch.prompt_attention_mask is not None: + batch.prompt_embeds = self._process_non_padded_embeddings( + batch.prompt_embeds[0], batch.prompt_attention_mask[0]) + else: + raise ValueError("prompt_attention_mask is None") + + # Prepare batch data for Parquet dataset + batch_data: list[dict[str, Any]] = [] + + for key in dataclasses.fields(batch): + value = getattr(batch, key.name) + if isinstance(value, list): + for idx in range(len(value)): + if isinstance(value[idx], torch.Tensor): + value[idx] = value[idx].cpu().numpy() + elif isinstance(value, torch.Tensor): + value = value.cpu().numpy() + setattr(batch, key.name, value) + + # Create record for Parquet dataset + records = self.create_records_from_batch(batch) + batch_data.extend(records) + + if batch_data: + self.num_processed_samples += len(batch_data) + # Convert batch data to PyArrow arrays + table = self._convert_batch_to_pyarrow_table(batch_data) + + # Store the table in a list for later processing + self.all_tables.append(table) + logger.debug("Collected batch with %s samples", len(table)) + + # If flush is needed + if self.num_processed_samples >= self.flush_frequency: + self.flush_tables(output_dir) + + def _process_non_padded_embeddings( + self, prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor) -> list[torch.Tensor]: + """Process non-padded embeddings""" + assert isinstance(prompt_embeds, torch.Tensor) + assert isinstance(prompt_attention_mask, torch.Tensor) + assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0] + + # Get sequence lengths from attention masks (number of 1s) + seq_lens = prompt_attention_mask.sum(dim=1) + + non_padded_embeds = [] + + # Process each item in the batch + for i in range(prompt_embeds.size(0)): + seq_len = seq_lens[i].item() + # Slice the embeddings and masks to keep only non-padding parts + non_padded_embeds.append(prompt_embeds[i, :seq_len]) + + return non_padded_embeds + + def _convert_batch_to_pyarrow_table(self, + batch_data: list[dict]) -> pa.Table: + """Convert batch data to PyArrow table""" + arrays = [] + + for field in self.schema_fields: + if field.endswith('_bytes'): + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.binary())) + elif field.endswith('_shape'): + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.list_(pa.int32()))) + elif field in ['width', 'height', 'num_frames']: + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.int32())) + elif field in ['duration_sec', 'fps']: + arrays.append( + pa.array([record[field] for record in batch_data], + type=pa.float32())) + else: + arrays.append(pa.array([record[field] + for record in batch_data])) + + return pa.Table.from_arrays(arrays, names=self.schema_fields) + + def flush_tables(self, output_dir: str): + """Flush collected tables to disk""" + if not hasattr(self, 'all_tables') or not self.all_tables: + return + + logger.debug("Combining %d batches...", len(self.all_tables)) + combined_table = pa.concat_tables(self.all_tables) + assert len(combined_table) == self.num_processed_samples + logger.debug("Total samples collected: %d", len(combined_table)) + + # Calculate total number of chunks needed, putting remainder into self.all_tables + total_files = max(self.num_processed_samples // self.samples_per_file, + 1) + + logger.debug("Fixed samples per parquet file: %d", + self.samples_per_file) + logger.debug("Total number of parquet files: %d", total_files) + logger.debug( + "Total samples to be processed: %d (putting %d samples into self.all_tables)", + total_files * self.samples_per_file, + self.num_processed_samples % self.samples_per_file) + + # Split work among processes + num_workers = int(min(multiprocessing.cpu_count(), total_files)) + files_per_worker = (total_files + num_workers - 1) // num_workers + + logger.debug("Using %d workers to process %d files", num_workers, + total_files) + logger.debug("Files per worker: %s", files_per_worker) + + # Prepare work ranges + work_ranges = [] + for i in range(num_workers): + start_idx = i * files_per_worker + end_idx = min((i + 1) * files_per_worker, total_files) + if start_idx < total_files: + work_ranges.append((start_idx, end_idx, combined_table, i, + output_dir, self.samples_per_file)) + + total_written = 0 + failed_ranges = [] + with ProcessPoolExecutor(max_workers=num_workers) as executor: + futures = { + executor.submit(self.file_writer_fn, work_range): work_range + for work_range in work_ranges + } + for future in futures: + try: + written = future.result() + total_written += written + logger.info("Processed file with %s samples", written) + except Exception as e: + work_range = futures[future] + failed_ranges.append(work_range) + logger.error("Failed to process range %s-%s: %s", + work_range[0], work_range[1], str(e)) + + # Retry failed ranges sequentially + if failed_ranges: + logger.warning("Retrying %s failed ranges sequentially", + len(failed_ranges)) + for work_range in failed_ranges: + try: + total_written += self.file_writer_fn(work_range) + except Exception as e: + logger.error( + "Failed to process range %s-%s after retry: %s", + work_range[0], work_range[1], str(e)) + + self.num_saved_files += total_files + + # Clear tables list + self.all_tables = [] + if self.num_processed_samples > self.samples_per_file: + saved_samples = total_files * self.samples_per_file + self.all_tables.append(combined_table.slice(saved_samples)) + self.num_processed_samples -= saved_samples + else: + self.num_processed_samples = 0 + + del combined_table + gc.collect() + + def clean_up(self) -> None: + """Clean up all tables""" + self.all_tables = [] + self.num_processed_samples = 0 + self.num_saved_files = 0 + gc.collect() + + def _default_file_writer_fn(self, args_tuple: tuple) -> int: + """Default chunk processing implementation""" + start_idx, end_idx, combined_table, worker_id, output_dir, samples_per_file = args_tuple + + written_count = 0 + for file_idx in range(start_idx, end_idx): + start_row = file_idx * samples_per_file + end_row = min(start_row + samples_per_file, len(combined_table)) + + if start_row >= len(combined_table): + break + + chunk_table = combined_table.slice(start_row, end_row - start_row) + + # Write to file + output_file = os.path.join( + output_dir, + f"chunk_{file_idx + self.num_saved_files:06d}.parquet") + pq.write_table(chunk_table, output_file) + written_count += len(chunk_table) + + return written_count diff --git a/trainer/workflow/preprocess/preprocess_workflow.py b/trainer/workflow/preprocess/preprocess_workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..5a33d34b7821f25fe4fee5d271e3104f04730d7e --- /dev/null +++ b/trainer/workflow/preprocess/preprocess_workflow.py @@ -0,0 +1,127 @@ +import os +from typing import cast + +from datasets import load_dataset +from torch.utils.data import DataLoader + +from trainer.configs.configs import PreprocessConfig +from trainer.dataset.dataloader.schema import (pyarrow_schema_i2v, + pyarrow_schema_t2v) +from trainer.trainer_args import TrainerArgs, WorkloadType +from trainer.logger import init_logger +from trainer.pipelines.pipeline_registry import PipelineType +from trainer.workflow.preprocess.components import ( + ParquetDatasetSaver, PreprocessingDataValidator, VideoForwardBatchBuilder) +from trainer.workflow.preprocess.record_schema import ( + basic_t2v_record_creator, i2v_record_creator) +from trainer.workflow.workflow_base import WorkflowBase + +logger = init_logger(__name__) + + +class PreprocessWorkflow(WorkflowBase): + + def register_pipelines(self) -> None: + self.add_pipeline_config("preprocess_pipeline", + (PipelineType.PREPROCESS, self.trainer_args)) + + def register_components(self) -> None: + assert self.trainer_args.preprocess_config is not None + preprocess_config: PreprocessConfig = self.trainer_args.preprocess_config + + # raw data validator + raw_data_validator = PreprocessingDataValidator( + max_height=preprocess_config.max_height, + max_width=preprocess_config.max_width, + num_frames=preprocess_config.num_frames, + train_fps=preprocess_config.train_fps, + speed_factor=preprocess_config.speed_factor, + video_length_tolerance_range=preprocess_config. + video_length_tolerance_range, + drop_short_ratio=preprocess_config.drop_short_ratio, + ) + self.add_component("raw_data_validator", raw_data_validator) + + # training dataset + training_dataset = load_dataset(preprocess_config.dataset_path, + split="train") + # set load_from_cache_file to False to check filter stats + training_dataset = training_dataset.filter(raw_data_validator) + # we do not use collate_fn here because we use iterable-style Dataset + # and want to keep the original type of the dataset + training_dataloader = DataLoader( + training_dataset, + batch_size=preprocess_config.preprocess_video_batch_size, + num_workers=preprocess_config.dataloader_num_workers, + collate_fn=lambda x: x, + ) + self.add_component("training_dataloader", training_dataloader) + + # try to load validation dataset if it exists + try: + validation_dataset = load_dataset(preprocess_config.dataset_path, + split="validation") + validation_dataset = validation_dataset.filter(raw_data_validator) + validation_dataloader = DataLoader( + validation_dataset, + batch_size=preprocess_config.preprocess_video_batch_size, + num_workers=preprocess_config.dataloader_num_workers, + collate_fn=lambda x: x, + ) + except ValueError: + logger.warning( + "Validation dataset not found, skipping validation dataset preprocessing." + ) + validation_dataloader = None + + self.add_component("validation_dataloader", validation_dataloader) + + # forward batch builder + video_forward_batch_builder = VideoForwardBatchBuilder( + seed=self.trainer_args.preprocess_config.seed) + self.add_component("video_forward_batch_builder", + video_forward_batch_builder) + + # record creator + if self.trainer_args.workload_type == WorkloadType.I2V: + record_creator = i2v_record_creator + schema_fields = [f.name for f in pyarrow_schema_i2v] + else: + record_creator = basic_t2v_record_creator + schema_fields = [f.name for f in pyarrow_schema_t2v] + processed_dataset_saver = ParquetDatasetSaver( + flush_frequency=self.trainer_args.preprocess_config. + flush_frequency, + samples_per_file=self.trainer_args.preprocess_config. + samples_per_file, + schema_fields=schema_fields, + record_creator=record_creator, + ) + self.add_component("processed_dataset_saver", processed_dataset_saver) + + def prepare_system_environment(self) -> None: + assert self.trainer_args.preprocess_config is not None + dataset_output_dir = self.trainer_args.preprocess_config.dataset_output_dir + os.makedirs(dataset_output_dir, exist_ok=True) + + validation_dataset_output_dir = os.path.join(dataset_output_dir, + "validation_dataset") + os.makedirs(validation_dataset_output_dir, exist_ok=True) + self.validation_dataset_output_dir = validation_dataset_output_dir + + training_dataset_output_dir = os.path.join(dataset_output_dir, + "training_dataset") + os.makedirs(training_dataset_output_dir, exist_ok=True) + self.training_dataset_output_dir = training_dataset_output_dir + + @classmethod + def get_workflow_cls(cls, + trainer_args: TrainerArgs) -> "PreprocessWorkflow": + if trainer_args.workload_type == WorkloadType.T2V or trainer_args.workload_type == WorkloadType.I2V: + from trainer.workflow.preprocess.preprocess_workflow_t2v import ( + PreprocessWorkflowT2V) + return cast(PreprocessWorkflow, PreprocessWorkflowT2V) + else: + raise ValueError( + f"Workload type: {trainer_args.workload_type} is not supported in preprocessing workflow." + ) diff --git a/trainer/workflow/preprocess/preprocess_workflow_t2v.py b/trainer/workflow/preprocess/preprocess_workflow_t2v.py new file mode 100644 index 0000000000000000000000000000000000000000..e0925dc4859aba20efd116dd4f80024fd3de212c --- /dev/null +++ b/trainer/workflow/preprocess/preprocess_workflow_t2v.py @@ -0,0 +1,56 @@ +from typing import TYPE_CHECKING, Optional + +from tqdm import tqdm + +from trainer.pipelines.pipeline_batch_info import PreprocessBatch +from trainer.workflow.preprocess.components import ParquetDatasetSaver +from trainer.workflow.preprocess.preprocess_workflow import PreprocessWorkflow + +if TYPE_CHECKING: + from torch.utils.data import DataLoader + + from trainer.pipelines.composed_pipeline_base import ComposedPipelineBase + from trainer.workflow.preprocess.components import ( + VideoForwardBatchBuilder) + + +class PreprocessWorkflowT2V(PreprocessWorkflow): + training_dataloader: "DataLoader" + validation_dataloader: Optional["DataLoader"] + preprocess_pipeline: "ComposedPipelineBase" + processed_dataset_saver: "ParquetDatasetSaver" + video_forward_batch_builder: "VideoForwardBatchBuilder" + + def run(self) -> None: + # Training dataset preprocessing + for batch in tqdm(self.training_dataloader, + desc="Preprocessing training dataset", + unit="batch"): + forward_batch: PreprocessBatch = self.video_forward_batch_builder( + batch) + + forward_batch = self.preprocess_pipeline.forward( + forward_batch, self.trainer_args) + + self.processed_dataset_saver.save_and_write_parquet_batch( + forward_batch, self.training_dataset_output_dir) + + self.processed_dataset_saver.flush_tables( + self.training_dataset_output_dir) + self.processed_dataset_saver.clean_up() + + # Validation dataset preprocessing + if self.validation_dataloader is not None: + for batch in tqdm(self.validation_dataloader, + desc="Preprocessing validation dataset", + unit="batch"): + forward_batch = self.video_forward_batch_builder(batch) + + forward_batch = self.preprocess_pipeline.forward( + forward_batch, self.trainer_args) + + self.processed_dataset_saver.save_and_write_parquet_batch( + forward_batch, self.validation_dataset_output_dir) + self.processed_dataset_saver.flush_tables( + self.validation_dataset_output_dir) + self.processed_dataset_saver.clean_up() diff --git a/trainer/workflow/preprocess/record_schema.py b/trainer/workflow/preprocess/record_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..2fa37660a22fc3e166f0db10cf920c464dcf28ff --- /dev/null +++ b/trainer/workflow/preprocess/record_schema.py @@ -0,0 +1,122 @@ +from typing import Any + +from trainer.pipelines.pipeline_batch_info import PreprocessBatch + + +def basic_t2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: + """Create a record for the Parquet dataset from PreprocessBatch.""" + # For batch processing, we need to handle the case where some fields might be single values + # or lists depending on the batch size + + assert isinstance(batch.prompt, list) + assert isinstance(batch.width, list) + assert isinstance(batch.height, list) + assert isinstance(batch.fps, list) + assert isinstance(batch.num_frames, list) + + records = [] + for idx, video_name in enumerate(batch.video_file_name): + width = batch.width[idx] if batch.width is not None else 0 + height = batch.height[idx] if batch.height is not None else 0 + + # Get FPS - single value in PreprocessBatch + fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0 + + # For duration, we need to calculate it or use a default since it's not in PreprocessBatch + # duration = num_frames / fps if available + duration_val = 0.0 + if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0: + duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx]) + + record = { + "id": + video_name, + "vae_latent_bytes": + batch.latents[idx].tobytes(), + "vae_latent_shape": + list(batch.latents[idx].shape), + "vae_latent_dtype": + str(batch.latents[idx].dtype), + "text_embedding_bytes": + batch.prompt_embeds[idx].tobytes(), + "text_embedding_shape": + list(batch.prompt_embeds[idx].shape), + "text_embedding_dtype": + str(batch.prompt_embeds[idx].dtype), + "file_name": + video_name, + "caption": + batch.prompt[idx], + "media_type": + "video", + "width": + int(width), + "height": + int(height), + "num_frames": + batch.latents[idx].shape[1] + if len(batch.latents[idx].shape) > 1 else 0, + "duration_sec": + duration_val, + "fps": + fps_val, + } + records.append(record) + + return records + + +def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: + """Create a record for the Parquet dataset with CLIP features.""" + records = basic_t2v_record_creator(batch) + + assert len( + batch.image_embeds) == 1, "image embedding should be a single tensor" + image_embeds = batch.image_embeds[0] + image_latent = batch.image_latent + pil_image = batch.pil_image + + for idx, record in enumerate(records): + if image_embeds is not None: + record.update({ + "clip_feature_bytes": image_embeds[idx].tobytes(), + "clip_feature_shape": list(image_embeds[idx].shape), + "clip_feature_dtype": str(image_embeds[idx].dtype), + }) + else: + record.update({ + "clip_feature_bytes": b"", + "clip_feature_shape": [], + "clip_feature_dtype": "", + }) + + if image_latent is not None: + record.update({ + "first_frame_latent_bytes": + image_latent[idx].tobytes(), + "first_frame_latent_shape": + list(image_latent[idx].shape), + "first_frame_latent_dtype": + str(image_latent[idx].dtype), + }) + else: + record.update({ + "first_frame_latent_bytes": b"", + "first_frame_latent_shape": [], + "first_frame_latent_dtype": "", + }) + + if pil_image is not None: + record.update({ + "pil_image_bytes": pil_image[idx].tobytes(), + "pil_image_shape": list(pil_image[idx].shape), + "pil_image_dtype": str(pil_image[idx].dtype), + }) + else: + record.update({ + "pil_image_bytes": b"", + "pil_image_shape": [], + "pil_image_dtype": "", + }) + + return records diff --git a/trainer/workflow/workflow_base.py b/trainer/workflow/workflow_base.py new file mode 100644 index 0000000000000000000000000000000000000000..396922ee3f5c23bfaa0585429daa01ef3a309a17 --- /dev/null +++ b/trainer/workflow/workflow_base.py @@ -0,0 +1,186 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional + +from trainer.trainer_args import ExecutionMode, TrainerArgs +from trainer.logger import init_logger +from trainer.pipelines import ComposedPipelineBase, build_pipeline +from trainer.pipelines.pipeline_registry import PipelineType + +logger = init_logger(__name__) + + +class WorkflowBase(ABC): + """ + Abstract base class for defining video processing workflows. + + A workflow serves as the top-level orchestrator that coordinates multiple pipelines + and components to accomplish a specific video processing task. The workflow pattern + provides several key benefits: + + 1. **Separation of Concerns**: Workflows separate high-level orchestration logic + from low-level processing implementations in pipelines. + + 2. **Modularity**: Different workflows can be created for different execution modes + (preprocess, inference, etc.) while sharing common pipeline components. + + 3. **Configuration Management**: Workflows manage the configuration and initialization + of multiple related pipelines and components in a centralized manner. + + 4. **Environment Setup**: Workflows handle system-level setup and resource + allocation before pipeline execution begins. + + 5. **Lifecycle Management**: Workflows control the complete lifecycle from + initialization through execution to cleanup. + + The workflow acts as a factory and coordinator, creating the appropriate pipelines + based on configuration, setting up the execution environment, and orchestrating + the overall processing flow. + """ + + def __init__(self, trainer_args: TrainerArgs): + """ + Initialize the workflow with configuration arguments. + + Args: + trainer_args: Configuration object containing all parameters + needed for workflow and pipeline setup. + """ + self.trainer_args = trainer_args + + # TODO: pipeline_config should be: dict[str, PipelineConfig] + # pipeline_type should be included in the PipelineConfig + # pipeline_config[pipeline_name] = (pipeline_type, trainer_args) + self._pipeline_configs: dict[str, tuple[PipelineType, + TrainerArgs]] = {} + self._pipelines: dict[str, ComposedPipelineBase] = {} + self._components: dict[str, Any] = {} + self.register_pipelines() + self.register_components() + + self.prepare_system_environment() + self.load_pipelines() + + def load_pipelines(self) -> None: + """ + Create and initialize all registered pipelines. + + This method instantiates pipeline objects from their configurations + and makes them available as both dictionary entries and instance + attributes for convenient access. + """ + for pipeline_name, pipeline_config in self._pipeline_configs.items(): + pipeline_type, trainer_args = pipeline_config + pipeline = build_pipeline(trainer_args, pipeline_type) + self._pipelines[pipeline_name] = pipeline + setattr(self, pipeline_name, pipeline) + + def add_pipeline_config( + self, pipeline_name: str, + pipeline_config: tuple[PipelineType, TrainerArgs]) -> None: + """ + Register a pipeline configuration for later instantiation. + + Args: + pipeline_name: Unique identifier for the pipeline. + pipeline_config: Tuple containing the pipeline type and + configuration arguments. + """ + self._pipeline_configs[pipeline_name] = pipeline_config + + def add_component(self, component_name: str, component: Any) -> None: + """ + Register a component instance with the workflow. + + Components are auxiliary objects that may be shared across pipelines + or used for workflow-level functionality (e.g., databases, caches, + external services). + + Args: + component_name: Unique identifier for the component. + component: The component instance to register. + """ + self._components[component_name] = component + setattr(self, component_name, component) + + def get_component(self, component_name: str) -> Any: + """ + Retrieve a registered component by name. + + Args: + component_name: The name of the component to retrieve. + + Returns: + The component instance. + """ + return self._components[component_name] + + @abstractmethod + def register_components(self) -> None: + """ + Register workflow-specific components. + + Subclasses must implement this method to register any components + needed for their specific workflow (e.g., databases, external APIs, + shared resources). + """ + pass + + @abstractmethod + def register_pipelines(self) -> None: + """ + Register workflow-specific pipelines. + + Subclasses must implement this method to define which pipelines + are needed for their specific workflow and how they should be + configured. + """ + pass + + @abstractmethod + def prepare_system_environment(self) -> None: + """ + Prepare the system environment for workflow execution. + + Subclasses must implement this method to handle any system-level + setup required before pipeline execution (e.g., GPU initialization, + temporary directories, resource allocation). + """ + pass + + @abstractmethod + def run(self): + """ + Execute the main workflow logic. + + Subclasses must implement this method to define the specific + execution flow for their workflow, coordinating the registered + pipelines and components to accomplish the desired task. + """ + pass + + @classmethod + def get_workflow_cls( + cls, trainer_args: TrainerArgs) -> Optional["WorkflowBase"]: + """ + Factory method to get the appropriate workflow class based on execution mode. + + This method acts as a workflow factory, returning the appropriate + workflow class implementation based on the specified execution mode + in the configuration arguments. + + Args: + trainer_args: Configuration object containing the execution mode + and other parameters. + + Returns: + The appropriate workflow class for the specified execution mode, + or None if no workflow is available for the given mode. + """ + if trainer_args.mode == ExecutionMode.PREPROCESS: + from trainer.workflow.preprocess.preprocess_workflow import ( + PreprocessWorkflow) + return PreprocessWorkflow.get_workflow_cls(trainer_args) + else: + raise ValueError( + f"Execution mode: {trainer_args.mode} is not supported in workflow." + ) diff --git a/wan/README.md b/wan/README.md new file mode 100644 index 0000000000000000000000000000000000000000..29890672208fa8048b207009cfbe7073f0f165a5 --- /dev/null +++ b/wan/README.md @@ -0,0 +1,130 @@ +# WAN Pipeline Inference + +The WAN pipeline provides a lightweight alternative with distributed inference support. It's optimized for multi-GPU setups and offers faster inference with lower memory footprint. + +## Prerequisites + +Make sure you have completed the base environment setup from the [main README](../README.md), including: +- Creating the `worldplay` conda environment +- Installing `requirements.txt` +- Installing **SageAttention** (required for WAN pipeline) + +**Important:** All WAN commands must be run from the **repository root directory**, and require `PYTHONPATH` to include both the repo root and the `wan/` subdirectory: + +```bash +cd /path/to/HY-WorldPlay +export PYTHONPATH=$(pwd):$(pwd)/wan:$PYTHONPATH +``` + +## Download WAN Models + +The WAN models are included in the `tencent/HY-WorldPlay` repository. You can download them with: + +```bash +# Download both wan_transformer and wan_distilled_model into a local directory +huggingface-cli download tencent/HY-WorldPlay wan_transformer wan_distilled_model --local-dir /path/to/models +``` + +This creates the following structure: +``` +/path/to/models/ +├── wan_transformer/ +│ ├── config.json +│ └── diffusion_pytorch_model.safetensors +└── wan_distilled_model/ + └── model.pt +``` + +Then use: +- `--ar_model_path /path/to/models/wan_transformer` +- `--ckpt_path /path/to/models/wan_distilled_model/model.pt` + +> **Tip:** If you already have these files in your HuggingFace cache, you can find the snapshot path with: +> ```bash +> python -c "from huggingface_hub import snapshot_download; print(snapshot_download('tencent/HY-WorldPlay', local_files_only=True))" +> ``` +> Then use `/wan_transformer` and `/wan_distilled_model/model.pt` directly. + +## Configuration Parameters + +The WAN pipeline accepts the following parameters: + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `--input` | Text prompt or path to txt file with prompts | Required | +| `--image_path` | Input image path for I2V generation | None (T2V mode) | +| `--num_chunk` | Number of chunks to generate (each chunk = 4 latents) | 4 | +| `--pose` | Camera trajectory (pose string or JSON file) | `w-96` | +| `--ar_model_path` | Path to WAN transformer model directory | Required | +| `--ckpt_path` | Path to trained checkpoint file (model.pt) | Required | +| `--out` | Output directory for generated videos | `outputs` | + +> **Note:** The total pose duration in latents should match `num_chunk * 4`. For example, with `num_chunk=1`, use `--pose "w-4"` (4 latents). + +## Single-GPU Inference + +```bash +export PYTHONPATH=$(pwd):$(pwd)/wan:$PYTHONPATH + +torchrun --nproc_per_node=1 wan/generate.py \ + --input "First-person view walking around ancient Athens, with Greek architecture and marble structures" \ + --num_chunk 1 \ + --pose "w-4" \ + --ar_model_path /path/to/models/wan_transformer \ + --ckpt_path /path/to/models/wan_distilled_model/model.pt \ + --out outputs +``` + +## Multi-GPU Distributed Inference + +For multi-GPU inference with better performance: + +```bash +export PYTHONPATH=$(pwd):$(pwd)/wan:$PYTHONPATH + +# Using 4 GPUs +torchrun --nproc_per_node=4 wan/generate.py \ + --input "First-person view walking around ancient Athens, with Greek architecture and marble structures" \ + --num_chunk 4 \ + --pose "w-16" \ + --ar_model_path /path/to/models/wan_transformer \ + --ckpt_path /path/to/models/wan_distilled_model/model.pt \ + --out outputs +``` + +## Batch Processing with Text Files + +You can process multiple prompts by providing a text file: + +```bash +export PYTHONPATH=$(pwd):$(pwd)/wan:$PYTHONPATH + +# Create a prompts.txt file with one prompt per line +echo "First-person view of a medieval castle" > prompts.txt +echo "Walking through a cyberpunk city at night" >> prompts.txt +echo "Exploring an underwater coral reef" >> prompts.txt + +# Run inference on all prompts +torchrun --nproc_per_node=1 wan/generate.py \ + --input prompts.txt \ + --ar_model_path /path/to/models/wan_transformer \ + --ckpt_path /path/to/models/wan_distilled_model/model.pt +``` + +## Camera Control with WAN + +WAN uses the same camera control system as the HunyuanVideo pipeline: + +**Pose String Format:** +```bash +# Forward movement for 4 latents (with num_chunk=1) +--pose "w-4" + +# Complex trajectory +--pose "w-20, right-10, d-30, up-36" +``` + +**Supported Actions:** +- **Movement**: `w` (forward), `s` (backward), `a` (left), `d` (right) +- **Rotation**: `up` (pitch up), `down` (pitch down), `left` (yaw left), `right` (yaw right) +- **Format**: `action-duration` where duration specifies the number of latents corresponding to the given action. diff --git a/wan/README_zh.md b/wan/README_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..59ce218f8bba318182169ffb0c378033b6f8ea8c --- /dev/null +++ b/wan/README_zh.md @@ -0,0 +1,80 @@ +# WAN 模型推理 + +WAN Pipeline 提供了支持分布式推理的轻量级替代方案。它针对多 GPU 设置进行了优化，提供更快的推理速度和更低的内存占用。 + +## 下载 WAN 模型 + +从 HuggingFace 下载所需的 WAN 模型： + +```bash +# 下载 WAN transformer 模型 +huggingface-cli download tencent/HY-WorldPlay wan_transformer --local-dir /path/to/models/wan_transformer + +# 下载 WAN 蒸馏模型检查点 +huggingface-cli download tencent/HY-WorldPlay wan_distilled_model/model.pt --local-dir /path/to/models/wan_distilled_model +``` + +## 配置参数 + +WAN 管道接受以下参数： + +| 参数 | 描述 | 默认值 | +|------|------|--------| +| `--input` | 文本提示或包含提示的 txt 文件路径 | 必需 | +| `--image_path` | I2V 生成的输入图像路径 | None | +| `--num_chunk` | 块数（每块 = 16 帧） | 4 | +| `--pose` | 相机轨迹（姿态字符串或 JSON 文件） | `w-96` | +| `--ar_model_path` | WAN transformer 模型路径 | 必需 | +| `--ckpt_path` | 训练检查点路径 (model.pt) | 必需 | +| `--out` | 生成视频的输出目录 | `outputs` | + +## 多 GPU 分布式推理 + +多 GPU 推理以获得更好的性能： + +```bash +# 使用 4 个 GPU +torchrun --nproc_per_node=4 wan/generate.py \ + --input "First-person view walking around ancient Athens, with Greek architecture and marble structures" \ + --image_path /path/to/input/image.jpg \ + --num_chunk 4 \ + --pose "w-96" \ + --ar_model_path /path/to/models/wan_transformer \ + --ckpt_path /path/to/models/wan_distilled_model/model.pt \ + --out outputs +``` + +## 批量处理文本文件 + +您可以通过提供文本文件来处理多个提示： + +```bash +# 创建一个 prompts.txt 文件，每行一个提示 +echo "第一人称视角的中世纪城堡" > prompts.txt +echo "夜晚在赛博朋克城市中漫步" >> prompts.txt +echo "探索水下珊瑚礁" >> prompts.txt + +# 对所有提示运行推理 +python wan/generate.py \ + --input prompts.txt \ + --ar_model_path /path/to/models/wan_transformer \ + --ckpt_path /path/to/models/wan_distilled_model/model.pt +``` + +## WAN 的相机控制 + +WAN 使用与 HunyuanVideo 管道相同的相机控制系统： + +**姿态字符串格式：** +```bash +# 前进 96 个 latents（使用 num_chunk=4 生成 961 帧） +--pose "w-96" + +# 复杂轨迹 +--pose "w-20, right-10, d-30, up-36" +``` + +**支持的动作：** +- **移动**: `w` (前进), `s` (后退), `a` (左移), `d` (右移) +- **旋转**: `up` (抬头), `down` (低头), `left` (左转), `right` (右转) +- **格式**: `动作-时长`，时长代表动作对应的latents数量。每个latent对应4帧。 diff --git a/wan/__init__.py b/wan/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/wan/distributed/__init__.py b/wan/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f61330e6ce77014de1aeb4430078b99052c87cae --- /dev/null +++ b/wan/distributed/__init__.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 + +from distributed.communication_op import * +from distributed.parallel_state import ( + cleanup_dist_env_and_memory, + get_dp_group, + get_dp_rank, + get_dp_world_size, + get_local_torch_device, + get_sp_group, + get_sp_parallel_rank, + get_sp_world_size, + get_tp_group, + get_tp_rank, + get_tp_world_size, + get_world_group, + get_world_rank, + get_world_size, + init_distributed_environment, + initialize_model_parallel, + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from distributed.utils import * + +__all__ = [ + # Initialization + "init_distributed_environment", + "initialize_model_parallel", + "cleanup_dist_env_and_memory", + "model_parallel_is_initialized", + "maybe_init_distributed_environment_and_model_parallel", + # World group + "get_world_group", + "get_world_rank", + "get_world_size", + # Data parallel group + "get_dp_group", + "get_dp_rank", + "get_dp_world_size", + # Sequence parallel group + "get_sp_group", + "get_sp_parallel_rank", + "get_sp_world_size", + # Tensor parallel group + "get_tp_group", + "get_tp_rank", + "get_tp_world_size", + # Get torch device + "get_local_torch_device", +] diff --git a/wan/distributed/communication_op.py b/wan/distributed/communication_op.py new file mode 100644 index 0000000000000000000000000000000000000000..21d2518dee696c8e77031f2ff517c709befcc4ec --- /dev/null +++ b/wan/distributed/communication_op.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py +# High-level distributed communication operations for model parallelism + +import torch +import torch.distributed + +from distributed.parallel_state import get_sp_group, get_tp_group + + +def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + # Sum tensors across tensor parallel ranks + return get_tp_group().all_reduce(input_) + + +def tensor_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + # Gather and concatenate tensors from all tensor parallel ranks + return get_tp_group().all_gather(input_, dim) + + +# TODO: remove model, make it sequence_parallel +def all_to_all_sp( + input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 +) -> torch.Tensor: + """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group.""" + # Scatter along one dimension and gather along another for sequence parallelism + return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim) + + +def sequence_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + # Gather sequence chunks from all sequence parallel ranks + return get_sp_group().all_gather(input_, dim) diff --git a/wan/distributed/device_communicators/__init__.py b/wan/distributed/device_communicators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/wan/distributed/device_communicators/base_device_communicator.py b/wan/distributed/device_communicators/base_device_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..dadab9d49f393cd29a52a15c3d3cea6d77af0356 --- /dev/null +++ b/wan/distributed/device_communicators/base_device_communicator.py @@ -0,0 +1,292 @@ +from typing import Any + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.distributed import ProcessGroup, ReduceOp + + +class DistributedAutograd: + """Collection of autograd functions for distributed operations. + + This class provides custom autograd functions for distributed operations like all_reduce, + all_gather, and all_to_all. Each operation is implemented as a static inner class with + proper forward and backward implementations. + """ + + class AllReduce(torch.autograd.Function): + """Differentiable all_reduce operation. + + The gradient of all_reduce is another all_reduce operation since the operation + combines values from all ranks equally. + """ + + @staticmethod + def forward( + ctx: Any, + group: ProcessGroup, + input_: Tensor, + op: dist.ReduceOp | None = None, + ) -> Tensor: + ctx.group = group + ctx.op = op + output = input_.clone() + dist.all_reduce(output, group=group, op=op) + return output + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None]: + grad_output = grad_output.clone() + dist.all_reduce(grad_output, group=ctx.group, op=ctx.op) + return None, grad_output, None + + class AllGather(torch.autograd.Function): + """Differentiable all_gather operation. + + The operation gathers tensors from all ranks and concatenates them along a specified dimension. + The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks. + """ + + @staticmethod + def forward( + ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, dim: int + ) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.dim = dim + ctx.input_shape = input_.shape + + input_size = input_.size() + output_size = (input_size[0] * world_size,) + input_size[1:] + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + + dist.all_gather_into_tensor(output_tensor, input_, group=group) + + output_tensor = output_tensor.reshape((world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + return output_tensor + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None, None]: + # Split the gradient tensor along the gathered dimension + dim_size = grad_output.size(ctx.dim) // ctx.world_size + grad_chunks = grad_output.reshape( + grad_output.shape[: ctx.dim] + + (ctx.world_size, dim_size) + + grad_output.shape[ctx.dim + 1 :] + ) + grad_chunks = grad_chunks.movedim(ctx.dim, 0) + + # Each rank only needs its corresponding gradient + grad_input = torch.empty( + ctx.input_shape, dtype=grad_output.dtype, device=grad_output.device + ) + dist.reduce_scatter_tensor( + grad_input, grad_chunks.contiguous(), group=ctx.group + ) + + return None, grad_input, None, None + + class AllToAll4D(torch.autograd.Function): + """Differentiable all_to_all operation specialized for 4D tensors. + + This operation is particularly useful for attention operations where we need to + redistribute data across ranks for efficient parallel processing. + + The operation supports two modes: + 1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads + 2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions + """ + + @staticmethod + def forward( + ctx: Any, + group: ProcessGroup, + input_: Tensor, + world_size: int, + scatter_dim: int, + gather_dim: int, + ) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + + if world_size == 1: + return input_ + + assert ( + input_.dim() == 4 + ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + + if scatter_dim == 2 and gather_dim == 1: + bs, shard_seqlen, hn, hd = input_.shape + seqlen = shard_seqlen * world_size + shard_hn = hn // world_size + + input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd + output = torch.empty_like(input_) + + dist.all_to_all_single( + output, input_, group=group + ) # hn, shard_seqlen, bs, hd + + output = torch.cat( + output.split(shard_hn), dim=1 + ) # sharded hn, seqlen, bs, hd + + output = output.transpose( + 0, 2 + ).contiguous() # bs, seqlen, sharded_hn, hd + + return output + elif scatter_dim == 1 and gather_dim == 2: + bs, seqlen, shard_hn, hd = input_.shape + hn = shard_hn * world_size + shard_seqlen = seqlen // world_size + + input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd + + input_ = ( + input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd) + .transpose(0, 1) + .reshape(shard_hn * world_size, shard_seqlen, bs, hd) + .contiguous() + ) + + output = torch.empty_like(input_) + + dist.all_to_all_single(output, input_, group=group) + + output = output.transpose( + 0, 2 + ).contiguous() # bs, seqlen, sharded_hn, hd + + return output + else: + raise RuntimeError( + f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " + f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported." + ) + + @staticmethod + def backward( + ctx: Any, grad_output: Tensor + ) -> tuple[None, Tensor, None, None, None]: + if ctx.world_size == 1: + return None, grad_output, None, None, None + + # For backward pass, we swap scatter_dim and gather_dim + output = DistributedAutograd.AllToAll4D.apply( + ctx.group, grad_output, ctx.world_size, ctx.gather_dim, ctx.scatter_dim + ) + return None, output, None, None, None + + +class DeviceCommunicatorBase: + """ + Base class for device-specific communicator with autograd support. + It can use the `cpu_group` to initialize the communicator. + If the device has PyTorch integration (PyTorch can recognize its + communication backend), the `device_group` will also be given. + """ + + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + self.device = device or torch.device("cpu") + self.cpu_group = cpu_group + self.device_group = device_group + self.unique_name = unique_name + self.rank = dist.get_rank(cpu_group) + self.world_size = dist.get_world_size(cpu_group) + self.ranks = dist.get_process_group_ranks(cpu_group) + self.global_rank = dist.get_rank() + self.global_world_size = dist.get_world_size() + self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) + + def all_reduce( + self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + """Performs an all_reduce operation with gradient support.""" + return DistributedAutograd.AllReduce.apply(self.device_group, input_, op) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Performs an all_gather operation with gradient support.""" + if dim < 0: + dim += input_.dim() + return DistributedAutograd.AllGather.apply( + self.device_group, input_, self.world_size, dim + ) + + def all_to_all_4D( + self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 + ) -> torch.Tensor: + """Performs a 4D all-to-all operation with gradient support.""" + return DistributedAutograd.AllToAll4D.apply( + self.device_group, input_, self.world_size, scatter_dim, gather_dim + ) + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + # Gather. + torch.distributed.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + pass diff --git a/wan/distributed/device_communicators/cpu_communicator.py b/wan/distributed/device_communicators/cpu_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..6dc3d5b166f44ea7f26a1a48b2b34cb30aec9cee --- /dev/null +++ b/wan/distributed/device_communicators/cpu_communicator.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py +# CPU-based communicator with optional shared memory optimization + +import os + +import torch +from torch.distributed import ProcessGroup + +from platforms import current_platform +from platforms.interface import CpuArchEnum + +from .base_device_communicator import DeviceCommunicatorBase + + +class CpuCommunicator(DeviceCommunicatorBase): + + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + super().__init__(cpu_group, device, device_group, unique_name) + # Default to standard PyTorch distributed module + self.dist_module = torch.distributed + + # Use shared memory optimization for x86 CPUs with tensor parallel + if ( + (current_platform.get_cpu_architecture() == CpuArchEnum.X86) + and hasattr(torch.ops._C, "init_shm_manager") + and unique_name.startswith("tp") + ): + # Enable shared memory for faster communication + self.dist_module = _CPUSHMDistributed(self) + + def all_reduce( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM, + ) -> torch.Tensor: + # Perform all-reduce using distributed or shared memory module + self.dist_module.all_reduce(input_, group=self.device_group, op=op) + return input_ + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + # Normalize negative dimension to positive + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + # Destination rank prepares list to receive from all ranks + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + + # Gather. + self.dist_module.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + + if self.rank_in_group == dst: + # Concatenate gathered tensors along specified dimension + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + # Normalize negative dimension to positive + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + input_size = input_.size() + # NOTE: we have to use concat-style all-gather here, + # stack-style all-gather has compatibility issues with + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 + output_size = (input_size[0] * self.world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + # All-gather: collect tensors from all ranks + self.dist_module.all_gather_into_tensor( + output_tensor, input_, group=self.device_group + ) + + # Reshape: move world_size dimension to target dim position + output_tensor = output_tensor.reshape((self.world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (self.world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + return output_tensor + + +class _CPUSHMDistributed: + """Shared memory distributed operations for CPU tensor parallel""" + + def __init__(self, communicator: CpuCommunicator): + # Get unique instance identifier from environment + instance_identifier = os.environ["VLLM_DIST_IDENT"] + unique_name = communicator.unique_name + instance_identifier = f"{instance_identifier}-{unique_name}" + self.communicator = communicator + + # Create group identifier from rank list + group_ranks = [str(rank) for rank in self.communicator.ranks] + shm_group_identifier = f"[{'-'.join(group_ranks)}]" + self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm" + + # Initialize shared memory manager + self.handle = self._init_cpu_shm() + + def _init_cpu_shm(self) -> int: + handle = torch.ops._C.init_shm_manager( + self.group_name, + self.communicator.world_size, + self.communicator.rank, + ) + torch.distributed.barrier(self.communicator.device_group) + torch.ops._C.join_shm_manager( + handle, + self.group_name, + ) + torch.distributed.barrier(self.communicator.device_group) + + return int(handle) + + def all_reduce( + self, input: torch.Tensor, group: ProcessGroup | None = None + ) -> None: + torch.ops._C.shm_allreduce(self.handle, input) + + def gather( + self, + input: torch.Tensor, + gather_list: list[torch.Tensor] | None, + dst: int = -1, + group: ProcessGroup | None = None, + ) -> None: + # Note: different from the torch gather, here we use local dst rank. + torch.ops._C.shm_gather( + self.handle, + input, + gather_list, + torch.distributed.get_group_rank(group, dst), + ) + + def all_gather_into_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + group: ProcessGroup | None = None, + ) -> None: + torch.ops._C.shm_all_gather(self.handle, input, output) diff --git a/wan/distributed/device_communicators/cuda_communicator.py b/wan/distributed/device_communicators/cuda_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..37ee35b224aee51172b07ffe557c08aab010c84a --- /dev/null +++ b/wan/distributed/device_communicators/cuda_communicator.py @@ -0,0 +1,81 @@ +# CUDA-based device communicator using PyNccl for efficient GPU communication +import torch +from torch.distributed import ProcessGroup + +from distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, +) + + +class CudaCommunicator(DeviceCommunicatorBase): + + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + super().__init__(cpu_group, device, device_group, unique_name) + + from distributed.device_communicators.pynccl import PyNcclCommunicator + + self.pynccl_comm: PyNcclCommunicator | None = None + # Initialize NCCL communicator for multi-GPU communication + if self.world_size > 1: + self.pynccl_comm = PyNcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None): + pynccl_comm = self.pynccl_comm + assert pynccl_comm is not None + # Use PyNccl for optimized GPU all-reduce + out = pynccl_comm.all_reduce(input_, op=op) + if out is None: + # fall back to the default all-reduce using PyTorch. + # this usually happens during testing. + # when we run the model, allreduce only happens for the TP + # group, where we always have either custom allreduce or pynccl. + # Fallback to PyTorch distributed for compatibility + out = input_.clone() + torch.distributed.all_reduce(out, group=self.device_group, op=op) + return out + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + # Default to next rank in ring topology + dst = (self.rank_in_group + 1) % self.world_size + + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + # Use PyNccl send for better GPU performance + pynccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + # Default to previous rank in ring topology + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + # Use PyNccl recv for better GPU performance + pynccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + # Clean up PyNccl communicator resources + if self.pynccl_comm is not None: + self.pynccl_comm = None diff --git a/wan/distributed/device_communicators/pynccl.py b/wan/distributed/device_communicators/pynccl.py new file mode 100644 index 0000000000000000000000000000000000000000..1e58493ac7b15a26e19d7f396bb176f81b2659f0 --- /dev/null +++ b/wan/distributed/device_communicators/pynccl.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, + buffer_type, + cudaStream_t, + ncclComm_t, + ncclDataTypeEnum, + ncclRedOpTypeEnum, + ncclUniqueId, +) +from distributed.utils import StatelessProcessGroup +from distributed.logger import init_logger +from utils import current_stream + +logger = init_logger(__name__) + + +class PyNcclCommunicator: + + def __init__( + self, + group: ProcessGroup | StatelessProcessGroup, + device: int | str | torch.device, + library_path: str | None = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyNcclCommunicator to. If None, + it will be bind to f"cuda:{local_rank}". + library_path: the path to the NCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "PyNcclCommunicator should be attached to a non-NCCL group." + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + return + try: + self.nccl = NCCLLibrary(library_path) + except Exception: + # disable because of missing NCCL library + # e.g. in a non-GPU environment + self.available = False + self.disabled = True + return + + self.available = True + self.disabled = False + + logger.info("FastVideo is using nccl==%s", self.nccl.ncclGetVersion()) + + if self.rank == 0: + # get the unique id from NCCL + self.unique_id = self.nccl.ncclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = ncclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + # nccl communicator and stream will use this device + # `torch.cuda.device` is a context manager that changes the + # current cuda device to the specified one + with torch.cuda.device(device): + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank + ) + + stream = current_stream() + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + if stream is not None: + stream.synchronize() + del data + + def all_reduce( + self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None + ) -> torch.Tensor: + if self.disabled: + return None + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert in_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {in_tensor.device}" + ) + + out_tensor = torch.empty_like(in_tensor) + + if stream is None: + stream = current_stream() + self.nccl.ncclAllReduce( + buffer_type(in_tensor.data_ptr()), + buffer_type(out_tensor.data_ptr()), + in_tensor.numel(), + ncclDataTypeEnum.from_torch(in_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + return out_tensor + + def all_gather( + self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclAllGather( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + input_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def reduce_scatter( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None, + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclReduceScatter( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + output_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def send(self, tensor: torch.Tensor, dst: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclSend( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + dst, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def recv(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclRecv( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def broadcast(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + if src == self.rank: + sendbuff = buffer_type(tensor.data_ptr()) + # NCCL requires the sender also to have a receive buffer + recvbuff = buffer_type(tensor.data_ptr()) + else: + sendbuff = buffer_type() + recvbuff = buffer_type(tensor.data_ptr()) + self.nccl.ncclBroadcast( + sendbuff, + recvbuff, + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) diff --git a/wan/distributed/device_communicators/pynccl_wrapper.py b/wan/distributed/device_communicators/pynccl_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..f855520fc7241691ce64de60cf9356ccfc17406e --- /dev/null +++ b/wan/distributed/device_communicators/pynccl_wrapper.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py + +# This file is a pure Python wrapper for the NCCL library. +# The main purpose is to use NCCL combined with CUDA graph. +# Before writing this script, we tried the following approach: +# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself +# often gets stuck when initializing the NCCL communicator. +# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` +# contains many other potential cuda APIs, that are not allowed during +# capturing the CUDA graph. For further details, please check +# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . +# +# Another rejected idea is to write a C/C++ binding for NCCL. It is usually +# doable, but we often encounter issues related with nccl versions, and need +# to switch between different versions of NCCL. See +# https://github.com/NVIDIA/nccl/issues/1234 for more details. +# A C/C++ binding is not flexible enough to handle this. It requires +# recompilation of the code every time we want to switch between different +# versions. This current implementation, with a **pure** Python wrapper, is +# more flexible. We can easily switch between different versions of NCCL by +# changing the environment variable `FASTVIDEO_NCCL_SO_PATH`, or the `so_file` +# variable in the code. + +# TODO(will): support FASTVIDEO_NCCL_SO_PATH + +import ctypes +import platform +from dataclasses import dataclass +from typing import Any + +import torch +from torch.distributed import ReduceOp + +from distributed.logger import init_logger +from utils import find_nccl_library + +logger = init_logger(__name__) + +# === export types and functions from nccl to Python === +# for the original nccl definition, please check +# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in + +ncclResult_t = ctypes.c_int +ncclComm_t = ctypes.c_void_p + + +class ncclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +cudaStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +ncclDataType_t = ctypes.c_int + + +class ncclDataTypeEnum: + ncclInt8 = 0 + ncclChar = 0 + ncclUint8 = 1 + ncclInt32 = 2 + ncclInt = 2 + ncclUint32 = 3 + ncclInt64 = 4 + ncclUint64 = 5 + ncclFloat16 = 6 + ncclHalf = 6 + ncclFloat32 = 7 + ncclFloat = 7 + ncclFloat64 = 8 + ncclDouble = 8 + ncclBfloat16 = 9 + ncclNumTypes = 10 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.ncclInt8 + if dtype == torch.uint8: + return cls.ncclUint8 + if dtype == torch.int32: + return cls.ncclInt32 + if dtype == torch.int64: + return cls.ncclInt64 + if dtype == torch.float16: + return cls.ncclFloat16 + if dtype == torch.float32: + return cls.ncclFloat32 + if dtype == torch.float64: + return cls.ncclFloat64 + if dtype == torch.bfloat16: + return cls.ncclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +ncclRedOp_t = ctypes.c_int + + +class ncclRedOpTypeEnum: + ncclSum = 0 + ncclProd = 1 + ncclMax = 2 + ncclMin = 3 + ncclAvg = 4 + ncclNumOps = 5 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.ncclSum + if op == ReduceOp.PRODUCT: + return cls.ncclProd + if op == ReduceOp.MAX: + return cls.ncclMax + if op == ReduceOp.MIN: + return cls.ncclMin + if op == ReduceOp.AVG: + return cls.ncclAvg + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: list[Any] + + +class NCCLLibrary: + exported_functions = [ + # const char* ncclGetErrorString(ncclResult_t result) + Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), + # ncclResult_t ncclGetVersion(int *version); + Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]), + # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); + Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]), + # ncclResult_t ncclCommInitRank( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); + # note that ncclComm_t is a pointer type, so the first argument + # is a pointer to a pointer + Function( + "ncclCommInitRank", + ncclResult_t, + [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int], + ), + # ncclResult_t ncclAllReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllReduce", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclAllGather( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllGather", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclReduceScatter( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclReduceScatter", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclSend( + # const void* sendbuff, size_t count, ncclDataType_t datatype, + # int dest, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclSend", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclRecv( + # void* recvbuff, size_t count, ncclDataType_t datatype, + # int src, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclRecv", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclBroadcast( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, int root, ncclComm_t comm, + # cudaStream_t stream); + Function( + "ncclBroadcast", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # be cautious! this is a collective call, it will block until all + # processes in the communicator have called this function. + # because Python object destruction can happen in random order, + # it is better not to call it at all. + # ncclResult_t ncclCommDestroy(ncclComm_t comm); + Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: dict[str, dict[str, Any]] = {} + + def __init__(self, so_file: str | None = None): + + so_file = so_file or find_nccl_library() + + try: + if so_file not in NCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + NCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = NCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + logger.error( + "Failed to load NCCL library from %s ." + "It is expected if you are not running on NVIDIA/AMD GPUs." + "Otherwise, the nccl library might not exist, be corrupted " + "or it does not support the current platform %s." + "If you already have the library, please set the " + "environment variable FASTVIDEO_NCCL_SO_PATH" + " to point to the correct nccl library path.", + so_file, + platform.platform(), + ) + raise e + + if so_file not in NCCLLibrary.path_to_dict_mapping: + _funcs: dict[str, Any] = {} + for func in NCCLLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + NCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] + + def ncclGetErrorString(self, result: ncclResult_t) -> str: + return str(self._funcs["ncclGetErrorString"](result).decode("utf-8")) + + def NCCL_CHECK(self, result: ncclResult_t) -> None: + if result != 0: + error_str = self.ncclGetErrorString(result) + raise RuntimeError(f"NCCL error: {error_str}") + + def ncclGetVersion(self) -> str: + version = ctypes.c_int() + self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) + version_str = str(version.value) + # something like 21903 --> "2.19.3" + major = version_str[0].lstrip("0") + minor = version_str[1:3].lstrip("0") + patch = version_str[3:].lstrip("0") + return f"{major}.{minor}.{patch}" + + def ncclGetUniqueId(self) -> ncclUniqueId: + unique_id = ncclUniqueId() + self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) + return unique_id + + def ncclCommInitRank( + self, world_size: int, unique_id: ncclUniqueId, rank: int + ) -> ncclComm_t: + comm = ncclComm_t() + self.NCCL_CHECK( + self._funcs["ncclCommInitRank"]( + ctypes.byref(comm), world_size, unique_id, rank + ) + ) + return comm + + def ncclAllReduce( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllReduce"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclReduceScatter( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclReduceScatter"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclAllGather( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # which is an aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllGather"]( + sendbuff, recvbuff, count, datatype, comm, stream + ) + ) + + def ncclSend( + self, + sendbuff: buffer_type, + count: int, + datatype: int, + dest: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream) + ) + + def ncclRecv( + self, + recvbuff: buffer_type, + count: int, + datatype: int, + src: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream) + ) + + def ncclBroadcast( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + root: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclBroadcast"]( + sendbuff, recvbuff, count, datatype, root, comm, stream + ) + ) + + def ncclCommDestroy(self, comm: ncclComm_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) + + +__all__ = [ + "NCCLLibrary", + "ncclDataTypeEnum", + "ncclRedOpTypeEnum", + "ncclUniqueId", + "ncclComm_t", + "cudaStream_t", + "buffer_type", +] diff --git a/wan/distributed/envs.py b/wan/distributed/envs.py new file mode 100644 index 0000000000000000000000000000000000000000..b1b11095c965f947b41f670cc77b28f0a62859b0 --- /dev/null +++ b/wan/distributed/envs.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60 + FASTVIDEO_NCCL_SO_PATH: str | None = None + LD_LIBRARY_PATH: str | None = None + LOCAL_RANK: int = 0 + CUDA_VISIBLE_DEVICES: str | None = None + FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo") + FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo") + FASTVIDEO_CONFIGURE_LOGGING: int = 1 + FASTVIDEO_LOGGING_LEVEL: str = "INFO" + FASTVIDEO_LOGGING_PREFIX: str = "" + FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None + FASTVIDEO_TRACE_FUNCTION: int = 0 + FASTVIDEO_ATTENTION_BACKEND: str | None = None + FASTVIDEO_ATTENTION_CONFIG: str | None = None + FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork" + FASTVIDEO_TARGET_DEVICE: str = "cuda" + MAX_JOBS: str | None = None + NVCC_THREADS: str | None = None + CMAKE_BUILD_TYPE: str | None = None + VERBOSE: bool = False + FASTVIDEO_SERVER_DEV_MODE: bool = False + FASTVIDEO_STAGE_LOGGING: bool = False + + +def get_default_cache_root() -> str: + return os.getenv( + "XDG_CACHE_HOME", + os.path.join(os.path.expanduser("~"), ".cache"), + ) + + +def get_default_config_root() -> str: + return os.getenv( + "XDG_CONFIG_HOME", + os.path.join(os.path.expanduser("~"), ".config"), + ) + + +def maybe_convert_int(value: str | None) -> int | None: + if value is None: + return None + return int(value) + + +# The begin-* and end* here are used by the documentation generator +# to extract the used env vars. + +# begin-env-vars-definition + +environment_variables: dict[str, Callable[[], Any]] = { + # ================== Installation Time Env Vars ================== + # Target device of FastVideo, supporting [cuda (by default), + # rocm, neuron, cpu, openvino] + "FASTVIDEO_TARGET_DEVICE": lambda: os.getenv("FASTVIDEO_TARGET_DEVICE", "cuda"), + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + # Number of threads to use for nvcc + # By default this is 1. + # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. + "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), + # If set, fastvideo will use precompiled binaries (*.so) + "FASTVIDEO_USE_PRECOMPILED": lambda: bool( + os.environ.get("FASTVIDEO_USE_PRECOMPILED") + ) + or bool(os.environ.get("FASTVIDEO_PRECOMPILED_WHEEL_LOCATION")), + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"), + # If set, fastvideo will print verbose logs during installation + "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), + # Root directory for FASTVIDEO configuration files + # Defaults to `~/.config/fastvideo` unless `XDG_CONFIG_HOME` is set + # Note that this not only affects how fastvideo finds its configuration files + # during runtime, but also affects how fastvideo installs its configuration + # files during **installation**. + "FASTVIDEO_CONFIG_ROOT": lambda: os.path.expanduser( + os.getenv( + "FASTVIDEO_CONFIG_ROOT", + os.path.join(get_default_config_root(), "fastvideo"), + ) + ), + # ================== Runtime Env Vars ================== + # Root directory for FASTVIDEO cache files + # Defaults to `~/.cache/fastvideo` unless `XDG_CACHE_HOME` is set + "FASTVIDEO_CACHE_ROOT": lambda: os.path.expanduser( + os.getenv( + "FASTVIDEO_CACHE_ROOT", + os.path.join(get_default_cache_root(), "fastvideo"), + ) + ), + # Interval in seconds to log a warning message when the ring buffer is full + "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL": lambda: int( + os.environ.get("FASTVIDEO_RINGBUFFER_WARNING_INTERVAL", "60") + ), + # Path to the NCCL library file. It is needed because nccl>=2.19 brought + # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 + "FASTVIDEO_NCCL_SO_PATH": lambda: os.environ.get("FASTVIDEO_NCCL_SO_PATH", None), + # when `FASTVIDEO_NCCL_SO_PATH` is not set, fastvideo will try to find the nccl + # library file in the locations specified by `LD_LIBRARY_PATH` + "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None), + # Internal flag to enable Dynamo fullgraph capture + "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool( + os.environ.get("FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0" + ), + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + # timeout for each iteration in the engine + "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S": lambda: int( + os.environ.get("FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S", "60") + ), + # Logging configuration + # If set to 0, fastvideo will not configure logging + # If set to 1, fastvideo will configure logging using the default configuration + # or the configuration file specified by FASTVIDEO_LOGGING_CONFIG_PATH + "FASTVIDEO_CONFIGURE_LOGGING": lambda: int( + os.getenv("FASTVIDEO_CONFIGURE_LOGGING", "1") + ), + "FASTVIDEO_LOGGING_CONFIG_PATH": lambda: os.getenv("FASTVIDEO_LOGGING_CONFIG_PATH"), + # this is used for configuring the default logging level + "FASTVIDEO_LOGGING_LEVEL": lambda: os.getenv("FASTVIDEO_LOGGING_LEVEL", "INFO"), + # if set, FASTVIDEO_LOGGING_PREFIX will be prepended to all log messages + "FASTVIDEO_LOGGING_PREFIX": lambda: os.getenv("FASTVIDEO_LOGGING_PREFIX", ""), + # Trace function calls + # If set to 1, fastvideo will trace function calls + # Useful for debugging + "FASTVIDEO_TRACE_FUNCTION": lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")), + # Backend for attention computation + # Available options: + # - "TORCH_SDPA": use torch.nn.MultiheadAttention + # - "FLASH_ATTN": use FlashAttention + # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention + # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention + # - "SAGE_ATTN": use Sage Attention + "FASTVIDEO_ATTENTION_BACKEND": lambda: os.getenv( + "FASTVIDEO_ATTENTION_BACKEND", None + ), + # Path to the attention configuration file. Only used for sliding tile + # attention for now. + "FASTVIDEO_ATTENTION_CONFIG": lambda: ( + None + if os.getenv("FASTVIDEO_ATTENTION_CONFIG", None) is None + else os.path.expanduser(os.getenv("FASTVIDEO_ATTENTION_CONFIG", ".")) + ), + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "FASTVIDEO_WORKER_MULTIPROC_METHOD": lambda: os.getenv( + "FASTVIDEO_WORKER_MULTIPROC_METHOD", "fork" + ), + # Enables torch profiler if set. Path to the directory where torch profiler + # traces are saved. Note that it must be an absolute path. + "FASTVIDEO_TORCH_PROFILER_DIR": lambda: ( + None + if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None + else os.path.expanduser(os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", ".")) + ), + # If set, fastvideo will run in development mode, which will enable + # some additional endpoints for developing and debugging, + # e.g. `/reset_prefix_cache` + "FASTVIDEO_SERVER_DEV_MODE": lambda: bool( + int(os.getenv("FASTVIDEO_SERVER_DEV_MODE", "0")) + ), + # If set, fastvideo will enable stage logging, which will print the time + # taken for each stage + "FASTVIDEO_STAGE_LOGGING": lambda: bool( + int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0")) + ), +} + +# end-env-vars-definition + + +def __getattr__(name: str): + # lazy evaluation of environment variables + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return list(environment_variables.keys()) diff --git a/wan/distributed/logger.py b/wan/distributed/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..1d1c33b221c4b224ccb75a6eac00688141110552 --- /dev/null +++ b/wan/distributed/logger.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py +"""Logging configuration for fastvideo.""" +import datetime +import json +import logging +import os +import sys +from functools import lru_cache, partial +from logging import Logger +from logging.config import dictConfig +from os import path +from types import MethodType +from typing import Any, cast + +import distributed.envs as envs + +FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING +FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH +FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL +FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX + +RED = "\033[91m" +GREEN = "\033[92m" +RESET = "\033[0;0m" + +_warned_local_main_process = False +_warned_main_process = False + +_FORMAT = ( + f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s " + "[%(filename)s:%(lineno)d] %(message)s" +) +_DATE_FORMAT = "%m-%d %H:%M:%S" + +DEFAULT_LOGGING_CONFIG = { + "formatters": { + "fastvideo": { + "class": "logging_utils.NewLineFormatter", + "datefmt": _DATE_FORMAT, + "format": _FORMAT, + }, + }, + "handlers": { + "fastvideo": { + "class": "logging.StreamHandler", + "formatter": "fastvideo", + "level": FASTVIDEO_LOGGING_LEVEL, + "stream": "ext://sys.stdout", + }, + }, + "loggers": { + "fastvideo": { + "handlers": ["fastvideo"], + "level": "DEBUG", + "propagate": False, + }, + }, + "root": { + "handlers": ["fastvideo"], + "level": "DEBUG", + }, + "version": 1, + "disable_existing_loggers": False, +} + + +@lru_cache +def _print_info_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.info(msg, stacklevel=2) + + +@lru_cache +def _print_warning_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.warning(msg, stacklevel=2) + + +# TODO(will): add env variable to control this process-aware logging behavior +def _info( + logger: Logger, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, +) -> None: + """Process-aware INFO level logging function. + + This function controls logging behavior based on the process rank, allowing for + selective logging from specific processes in a distributed environment. + + Args: + logger: The logger instance to use for logging + msg: The message format string to log + *args: Format string arguments + main_process_only: If True, only log if this is the global main process (RANK=0) + local_main_process_only: If True, only log if this is the local main process (LOCAL_RANK=0) + **kwargs: Additional keyword arguments to pass to the logger.log method + - stacklevel: Defaults to 2 to show the original caller's location + + Note: + - When both main_process_only and local_main_process_only are True, + the message will be logged only if both conditions are met + - When both are False, the message will be logged from all processes + - By default, only logs from processes with LOCAL_RANK=0 + """ + try: + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + except Exception: + local_rank = 0 + rank = 0 + + is_main_process = rank == 0 + is_local_main_process = local_rank == 0 + + if (main_process_only and is_main_process) or ( + local_main_process_only and is_local_main_process + ): + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + global _warned_local_main_process, _warned_main_process + + if not _warned_local_main_process and local_main_process_only: + logger.warning( + "%s is_local_main_process is set to True, logging only from the local main process.%s", + GREEN, + RESET, + ) + _warned_local_main_process = True + if not _warned_main_process and main_process_only: + logger.warning( + "%s is_main_process_only is set to True, logging only from the main process.%s", + GREEN, + RESET, + ) + _warned_main_process = True + + if not main_process_only and not local_main_process_only: + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + +class _FastvideoLogger(Logger): + """ + Note: + This class is just to provide type information. + We actually patch the methods directly on the :class:`logging.Logger` + instance to avoid conflicting with other libraries such as + `intel_extension_for_pytorch.utils._logger`. + """ + + def info_once(self, msg: str) -> None: + """ + As :meth:`info`, but subsequent calls with the same message + are silently dropped. + """ + _print_info_once(self, msg) + + def warning_once(self, msg: str) -> None: + """ + As :meth:`warning`, but subsequent calls with the same message + are silently dropped. + """ + _print_warning_once(self, msg) + + def info( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, + ) -> None: + _info( + self, + msg, + *args, + main_process_only=main_process_only, + local_main_process_only=local_main_process_only, + **kwargs, + ) + + +def _configure_fastvideo_root_logger() -> None: + logging_config = dict[str, Any]() + + if not FASTVIDEO_CONFIGURE_LOGGING and FASTVIDEO_LOGGING_CONFIG_PATH: + raise RuntimeError( + "FASTVIDEO_CONFIGURE_LOGGING evaluated to false, but " + "FASTVIDEO_LOGGING_CONFIG_PATH was given. FASTVIDEO_LOGGING_CONFIG_PATH " + "implies FASTVIDEO_CONFIGURE_LOGGING. Please enable " + "FASTVIDEO_CONFIGURE_LOGGING or unset FASTVIDEO_LOGGING_CONFIG_PATH." + ) + + if FASTVIDEO_CONFIGURE_LOGGING: + logging_config = DEFAULT_LOGGING_CONFIG + + if FASTVIDEO_LOGGING_CONFIG_PATH: + if not path.exists(FASTVIDEO_LOGGING_CONFIG_PATH): + raise RuntimeError( + "Could not load logging config. File does not exist: %s", + FASTVIDEO_LOGGING_CONFIG_PATH, + ) + with open(FASTVIDEO_LOGGING_CONFIG_PATH, encoding="utf-8") as file: + custom_config = json.loads(file.read()) + + if not isinstance(custom_config, dict): + raise ValueError( + "Invalid logging config. Expected Dict, got %s.", + type(custom_config).__name__, + ) + logging_config = custom_config + + for formatter in logging_config.get("formatters", {}).values(): + # This provides backwards compatibility after #10134. + if formatter.get("class") == "logging.NewLineFormatter": + formatter["class"] = "logging_utils.NewLineFormatter" + + if logging_config: + dictConfig(logging_config) + + +def init_logger(name: str) -> _FastvideoLogger: + """The main purpose of this function is to ensure that loggers are + retrieved in such a way that we can be sure the root fastvideo logger has + already been configured.""" + + logger = logging.getLogger(name) + + methods_to_patch = { + "info_once": _print_info_once, + "warning_once": _print_warning_once, + "info": _info, + } + + for method_name, method in methods_to_patch.items(): + setattr( + logger, method_name, MethodType(method, logger) + ) # type: ignore[arg-type] + + return cast(_FastvideoLogger, logger) + + +# The root logger is initialized when the module is imported. +# This is thread-safe as the module is only imported once, +# guaranteed by the Python GIL. +_configure_fastvideo_root_logger() + +logger = init_logger(__name__) + + +def _trace_calls(log_path, root_dir, frame, event, arg=None): + if event in ["call", "return"]: + # Extract the filename, line number, function name, and the code object + filename = frame.f_code.co_filename + lineno = frame.f_lineno + func_name = frame.f_code.co_name + if not filename.startswith(root_dir): + # only log the functions in the fastvideo root_dir + return + # Log every function call or return + try: + last_frame = frame.f_back + if last_frame is not None: + last_filename = last_frame.f_code.co_filename + last_lineno = last_frame.f_lineno + last_func_name = last_frame.f_code.co_name + else: + # initial frame + last_filename = "" + last_lineno = 0 + last_func_name = "" + with open(log_path, "a") as f: + ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") + if event == "call": + f.write( + f"{ts} Call to" + f" {func_name} in {filename}:{lineno}" + f" from {last_func_name} in {last_filename}:" + f"{last_lineno}\n" + ) + else: + f.write( + f"{ts} Return from" + f" {func_name} in {filename}:{lineno}" + f" to {last_func_name} in {last_filename}:" + f"{last_lineno}\n" + ) + except NameError: + # modules are deleted during shutdown + pass + return partial(_trace_calls, log_path, root_dir) + + +def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): + """ + Enable tracing of every function call in code under `root_dir`. + This is useful for debugging hangs or crashes. + `log_file_path` is the path to the log file. + `root_dir` is the root directory of the code to trace. If None, it is the + fastvideo root directory. + + Note that this call is thread-level, any threads calling this function + will have the trace enabled. Other threads will not be affected. + """ + logger.warning( + "FASTVIDEO_TRACE_FUNCTION is enabled. It will record every" + " function executed by Python. This will slow down the code. It " + "is suggested to be used for debugging hang or crashes only." + ) + logger.info("Trace frame log is saved to %s", log_file_path) + if root_dir is None: + # by default, this is the fastvideo root directory + root_dir = os.path.dirname(os.path.dirname(__file__)) + sys.settrace(partial(_trace_calls, log_file_path, root_dir)) diff --git a/wan/distributed/parallel_state.py b/wan/distributed/parallel_state.py new file mode 100644 index 0000000000000000000000000000000000000000..b1b629fdf8f24e4e64538f37133bf938aa4d4d42 --- /dev/null +++ b/wan/distributed/parallel_state.py @@ -0,0 +1,1300 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Adapted from +"""FastVideo distributed state. +It takes over the control of the distributed environment from PyTorch. +The typical workflow is: + +- call `init_distributed_environment` to initialize the distributed environment. +- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to + initialize the model parallel groups. + +- any code dealing with the distributed stuff + +- call `destroy_model_parallel` to destroy the model parallel groups. +- call `destroy_distributed_environment` to destroy the distributed environment. + +If you only need to use the distributed environment without model parallelism, + you can skip the model parallel initialization and destruction steps. +""" +import contextlib +import os +import pickle +import weakref +from collections import namedtuple +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from multiprocessing import shared_memory +from typing import Any, Optional +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup, ReduceOp + +import distributed.envs as envs +from distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, +) +from distributed.device_communicators.cpu_communicator import CpuCommunicator +from distributed.utils import StatelessProcessGroup +from distributed.logger import init_logger +from platforms import current_platform + +logger = init_logger(__name__) + + +@dataclass +class GraphCaptureContext: + stream: torch.cuda.Stream | None + + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + + +def _split_tensor_dict( + tensor_dict: dict[str, torch.Tensor | Any], +) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]: + """Split the tensor dictionary into two parts: + 1. A list of (key, value) pairs. If the value is a tensor, it is replaced + by its metadata. + 2. A list of tensors. + """ + metadata_list: list[tuple[str, Any]] = [] + tensor_list: list[torch.Tensor] = [] + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor): + # Note: we cannot use `value.device` here, + # because it contains not only the device type but also the device + # index (e.g. "cuda:0"). We only need the device type. + # receiving side will set the device index. + device = value.device.type + metadata_list.append( + (key, TensorMetadata(device, value.dtype, value.size())) + ) + tensor_list.append(value) + else: + metadata_list.append((key, value)) + return metadata_list, tensor_list + + +_group_name_counter: dict[str, int] = {} + + +def _get_unique_name(name: str) -> str: + """Get a unique name for the group. + Example: + _get_unique_name("tp") -> "tp:0" + _get_unique_name("tp") -> "tp:1" + """ + if name not in _group_name_counter: + _group_name_counter[name] = 0 + newname = f"{name}:{_group_name_counter[name]}" + _group_name_counter[name] += 1 + return newname + + +_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._all_reduce_out_place(tensor) + + +def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + return torch.empty_like(tensor) + + +class GroupCoordinator: + """ + PyTorch ProcessGroup wrapper for a group of processes. + PyTorch ProcessGroup is bound to one specific communication backend, + e.g. NCCL, Gloo, MPI, etc. + GroupCoordinator takes charge of all the communication operations among + the processes in the group. It manages both CPU and device + communication. + """ + + # available attributes: + rank: int # global rank + ranks: list[int] # global ranks in the group + world_size: int # size of the group + # difference between `local_rank` and `rank_in_group`: + # if we have a group of size 4 across two nodes: + # Process | Node | Rank | Local Rank | Rank in Group + # 0 | 0 | 0 | 0 | 0 + # 1 | 0 | 1 | 1 | 1 + # 2 | 1 | 2 | 0 | 2 + # 3 | 1 | 3 | 1 | 3 + local_rank: int # local rank used to assign devices + rank_in_group: int # rank inside the group + cpu_group: ProcessGroup # group for CPU communication + device_group: ProcessGroup # group for device communication + use_device_communicator: bool # whether to use device communicator + device_communicator: DeviceCommunicatorBase # device communicator + mq_broadcaster: Any | None # shared memory broadcaster + + def __init__( + self, + group_ranks: list[list[int]], + local_rank: int, + torch_distributed_backend: str | Backend, + use_device_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, + ): + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + + for ranks in group_ranks: + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend + ) + # a group with `gloo` backend, to allow direct coordination between + # processes through the CPU. + cpu_group = torch.distributed.new_group(ranks, backend="gloo") + if self.rank in ranks: + self.ranks = ranks + self.world_size = len(ranks) + self.rank_in_group = ranks.index(self.rank) + self.device_group = device_group + self.cpu_group = cpu_group + try: + assert self.cpu_group is not None + assert self.device_group is not None + except Exception as e: + print(f"rank: {self.rank} group not found") + raise e + + from platforms import current_platform + + # TODO: fix it for other platforms + self.device = get_local_torch_device() + + self.use_device_communicator = use_device_communicator + + self.device_communicator: DeviceCommunicatorBase = None # type: ignore + if use_device_communicator and self.world_size > 1: + # Platform-aware device communicator selection + if current_platform.is_cuda_alike(): + from distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + + self.device_communicator = CudaCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + else: + # For MPS and CPU, use the CPU communicator + self.device_communicator = CpuCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + + self.mq_broadcaster = None + + from platforms import current_platform + + # TODO(will): check if this is needed + # self.use_custom_op_call = current_platform.is_cuda_alike() + self.use_custom_op_call = False + + @property + def first_rank(self): + """Return the global rank of the first process in the group""" + return self.ranks[0] + + @property + def last_rank(self): + """Return the global rank of the last process in the group""" + return self.ranks[-1] + + @property + def is_first_rank(self): + """Return whether the caller is the first process in the group""" + return self.rank == self.first_rank + + @property + def is_last_rank(self): + """Return whether the caller is the last process in the group""" + return self.rank == self.last_rank + + @property + def next_rank(self): + """Return the global rank of the process that follows the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group + 1) % world_size] + + @property + def prev_rank(self): + """Return the global rank of the process that precedes the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group - 1) % world_size] + + @contextmanager + def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None): + # Platform-aware graph capture + from platforms import current_platform + + if current_platform.is_cuda_alike(): + if graph_capture_context is None: + stream = torch.cuda.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = torch.cuda.current_stream() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with torch.cuda.stream(stream): + yield graph_capture_context + else: + # For non-CUDA platforms (MPS, CPU), just yield the context without stream management + if graph_capture_context is None: + # Create a dummy context for non-CUDA platforms + graph_capture_context = GraphCaptureContext(None) + yield graph_capture_context + + def all_reduce( + self, input_: torch.Tensor, op: torch.distributed.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + """ + User-facing all-reduce function before we actually call the + all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we always make the all-reduce operation + out-of-place. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + if self.use_custom_op_call: + return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name) + else: + return self._all_reduce_out_place(input_, op=op) + + def _all_reduce_out_place( + self, input_: torch.Tensor, op: torch.distributed.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + return self.device_communicator.all_reduce(input_, op=op) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + return self.device_communicator.all_gather(input_, dim) + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + return self.device_communicator.gather(input_, dst, dim) + + def all_to_all_4D( + self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 + ) -> torch.Tensor: + if self.world_size == 1: + return input_ + return self.device_communicator.all_to_all_4D(input_, scatter_dim, gather_dim) + + def broadcast(self, input_: torch.Tensor, src: int = 0): + """Broadcast the input tensor. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + # Broadcast. + torch.distributed.broadcast( + input_, src=self.ranks[src], group=self.device_group + ) + return input_ + + def broadcast_object(self, obj: Any | None = None, src: int = 0): + """Broadcast the input object. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj + if self.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_broadcaster.broadcast_object(obj) + if self.rank_in_group == src: + torch.distributed.broadcast_object_list( + [obj], src=self.ranks[src], group=self.cpu_group + ) + return obj + else: + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=self.ranks[src], group=self.cpu_group + ) + return recv[0] + + def broadcast_object_list( + self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None + ): + """Broadcast the input object list. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj_list + # Broadcast. + torch.distributed.broadcast_object_list( + obj_list, src=self.ranks[src], group=self.device_group + ) + return obj_list + + def send_object(self, obj: Any, dst: int) -> None: + """Send the input object list to the destination rank.""" + """NOTE: `dst` is the local rank of the destination rank.""" + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank." + ) + + # Serialize object to tensor and get the size as well + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) + + size_tensor = torch.tensor( + [object_tensor.numel()], dtype=torch.long, device="cpu" + ) + + # Send object size + + torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group) + + # Send object + torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group) + + return None + + def recv_object(self, src: int) -> Any: + """Receive the input object list from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + + assert src < self.world_size, f"Invalid src rank ({src})" + + assert ( + src != self.rank_in_group + ), "Invalid source rank. Source rank is the same as the current rank." + + size_tensor = torch.empty(1, dtype=torch.long, device="cpu") + + # Receive object size + rank_size = torch.distributed.recv( + size_tensor, src=self.ranks[src], group=self.cpu_group + ) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device="cpu", + ) + + rank_object = torch.distributed.recv( + object_tensor, src=self.ranks[src], group=self.cpu_group + ) + + assert ( + rank_object == rank_size + ), "Received object sender rank does not match the size sender rank." + + obj = pickle.loads(object_tensor.numpy().tobytes()) + + return obj + + def broadcast_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any] | None = None, + src: int = 0, + group: ProcessGroup | None = None, + metadata_group: ProcessGroup | None = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Broadcast the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + group = self.device_group + metadata_group = self.cpu_group + assert src < self.world_size, f"Invalid src rank ({src})" + + rank_in_group = self.rank_in_group + if rank_in_group == src: + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `broadcast_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.broadcast_object(metadata_list, src=src) + async_handles = [] + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=metadata_group, async_op=True + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + for async_handle in async_handles: + async_handle.wait() + + else: + metadata_list = self.broadcast_object(None, src=src) + tensor_dict = {} + async_handles = [] + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty( + value.size, dtype=value.dtype, device=value.device + ) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + for async_handle in async_handles: + async_handle.wait() + return tensor_dict + + def send_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any], + dst: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Send the input tensor dictionary. + NOTE: `dst` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `send_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.send_object(metadata_list, dst=dst) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if all_gather_group is not None and tensor.numel() % all_gather_size == 0: + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.send( + tensor, dst=self.ranks[dst], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.send(tensor, dst=self.ranks[dst], group=group) + return None + + def recv_tensor_dict( + self, + src: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Recv the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return None + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: dict[str, Any] = {} + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.recv( + tensor, src=self.ranks[src], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.recv(tensor, src=self.ranks[src], group=group) + if use_all_gather: + # do the allgather + tensor = all_gather_group.all_gather(tensor, dim=0) # type: ignore + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + return tensor_dict + + def barrier(self) -> None: + """Barrier synchronization among the group. + NOTE: don't use `device_group` here! `barrier` in NCCL is + terrible because it is internally a broadcast operation with + secretly created GPU tensors. It is easy to mess up the current + device. Use the CPU group instead. + """ + torch.distributed.barrier(group=self.cpu_group) + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + self.device_communicator.send(tensor, dst) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + return self.device_communicator.recv(size, dtype, src) + + def destroy(self) -> None: + if self.device_group is not None: + torch.distributed.destroy_process_group(self.device_group) + self.device_group = None + if self.cpu_group is not None: + torch.distributed.destroy_process_group(self.cpu_group) + self.cpu_group = None + if self.device_communicator is not None: + self.device_communicator.destroy() + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: GroupCoordinator | None = None +_NODE: GroupCoordinator | None = None + + +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, "world group is not initialized" + return _WORLD + + +def init_world_group( + ranks: list[int], local_rank: int, backend: str +) -> GroupCoordinator: + return GroupCoordinator( + group_ranks=[ranks], + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: list[list[int]], + local_rank: int, + backend: str, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, +) -> GroupCoordinator: + + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: GroupCoordinator | None = None + + +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +_ENABLE_CUSTOM_ALL_REDUCE = True + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def init_distributed_environment( + world_size: int = 1, + rank: int = 0, + distributed_init_method: str = "env://", + local_rank: int = 0, + backend: str = "nccl", + device_id: torch.device | None = None, +): + # Determine the appropriate backend based on the platform + from platforms import current_platform + + if backend == "nccl" and not current_platform.is_cuda_alike(): + # Use gloo backend for non-CUDA platforms (MPS, CPU) + backend = "gloo" + logger.info("Using gloo backend for %s platform", current_platform.device_name) + + logger.debug( + "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s", + world_size, + rank, + local_rank, + distributed_init_method, + backend, + ) + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment" + ) + + # For MPS, don't pass device_id as it doesn't support device indices + if current_platform.is_mps(): + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + ) + else: + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + device_id=device_id, + ) + # set the local rank + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = envs.LOCAL_RANK + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert ( + _WORLD.world_size == torch.distributed.get_world_size() + ), "world group already initialized with a different world size" + + +_SP: GroupCoordinator | None = None + + +def get_sp_group() -> GroupCoordinator: + assert _SP is not None, "sequence model parallel group is not initialized" + return _SP + + +_DP: GroupCoordinator | None = None + + +def get_dp_group() -> GroupCoordinator: + assert _DP is not None, "data parallel group is not initialized" + return _DP + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + sequence_model_parallel_size: int = 1, + data_parallel_size: int = 1, + backend: str | None = None, +) -> None: + """ + Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism (used for language encoder). + sequence_model_parallel_size: number of GPUs used for sequence model + parallelism (used for DiT). + """ + # Get world size and rank. Ensure some consistencies. + assert ( + _WORLD is not None + ), "world group is not initialized, please call init_distributed_environment first" + world_size: int = get_world_size() + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + assert ( + world_size >= tensor_model_parallel_size + ), f"world_size({world_size}) < tensor_model_parallel_size({tensor_model_parallel_size})" + num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size + global _TP + assert _TP is None, "tensor model parallel group is already initialized" + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + ) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="tp", + ) + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = world_size // sequence_model_parallel_size + global _SP + assert _SP is None, "sequence model parallel group is already initialized" + group_ranks = [] + + # Since SP is incompatible with TP and PP, we can use a simpler group creation logic + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range( + i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size + ) + ) + group_ranks.append(ranks) + + _SP = init_model_parallel_group( + group_ranks, get_world_group().local_rank, backend, group_name="sp" + ) + + # Build the data parallel groups. + num_data_parallel_groups: int = sequence_model_parallel_size + global _DP + assert _DP is None, "data parallel group is already initialized" + group_ranks = [] + + for i in range(num_data_parallel_groups): + ranks = list(range(i, world_size, num_data_parallel_groups)) + group_ranks.append(ranks) + + _DP = init_model_parallel_group( + group_ranks, get_world_group().local_rank, backend, group_name="dp" + ) + + +def get_sp_world_size() -> int: + """Return world size for the sequence model parallel group.""" + return get_sp_group().world_size + + +def get_sp_parallel_rank() -> int: + """Return my rank for the sequence model parallel group.""" + return get_sp_group().rank_in_group + + +def get_sp_parallel_local_rank() -> int: + """Return my local rank for the sequence model parallel group.""" + return get_sp_group().local_rank + + +def get_world_size() -> int: + """Return world size for the world group.""" + return get_world_group().world_size + + +def get_world_rank() -> int: + """Return my rank for the world group.""" + return get_world_group().rank + + +def get_dp_world_size() -> int: + """Return world size for the data parallel group.""" + return get_dp_group().world_size + + +def get_dp_rank() -> int: + """Return my rank for the data parallel group.""" + return get_dp_group().rank_in_group + + +def get_local_torch_device() -> torch.device: + """Return the torch device for the current rank.""" + return ( + torch.device(f"cuda:{envs.LOCAL_RANK}") + if current_platform.is_cuda_alike() + else torch.device("mps") + ) + + +def maybe_init_distributed_environment_and_model_parallel( + tp_size: int, sp_size: int, distributed_init_method: str = "env://" +): + if _WORLD is not None and model_parallel_is_initialized(): + # make sure the tp and sp sizes are correct + assert ( + get_tp_world_size() == tp_size + ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" + assert ( + get_sp_world_size() == sp_size + ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + return + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + rank = int(os.environ.get("RANK", 0)) + device = get_local_torch_device() + logger.info( + "Initializing distributed environment with world_size=%d, device=%s", + world_size, + device, + ) + + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method=distributed_init_method, + device_id=device, + ) + initialize_model_parallel( + tensor_model_parallel_size=tp_size, sequence_model_parallel_size=sp_size + ) + + # Only set CUDA device if we're on a CUDA platform + if current_platform.is_cuda_alike(): + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + + +def model_parallel_is_initialized() -> bool: + """Check if tensor, sequence parallel groups are initialized.""" + return _TP is not None and _SP is not None and _DP is not None + + +_TP_STATE_PATCHED = False + + +@contextmanager +def patch_tensor_parallel_group(tp_group: GroupCoordinator): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _TP_STATE_PATCHED + assert not _TP_STATE_PATCHED, "Should not call when it's already patched" + + _TP_STATE_PATCHED = True + old_tp_group = get_tp_group() + global _TP + _TP = tp_group + try: + yield + finally: + # restore the original state + _TP_STATE_PATCHED = False + _TP = old_tp_group + + +def get_tp_world_size() -> int: + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tp_rank() -> int: + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def destroy_model_parallel() -> None: + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _SP + if _SP: + _SP.destroy() + _SP = None + + global _DP + if _DP: + _DP.destroy() + _DP = None + + +def destroy_distributed_environment() -> None: + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def cleanup_dist_env_and_memory(shutdown_ray: bool = False): + destroy_model_parallel() + destroy_distributed_environment() + with contextlib.suppress(AssertionError): + torch.distributed.destroy_process_group() + if shutdown_ray: + import ray # Lazy import Ray + + ray.shutdown() + + +def is_the_same_node_as( + pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0 +) -> list[int]: + """ + This is a collective operation that returns if each rank is in the same node + as the source rank. It tests if processes are attached to the same + memory system (shared access to shared memory). + """ + if isinstance(pg, ProcessGroup): + assert ( + torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL + ), "in_the_same_node_as should be tested with a non-NCCL group." + # local rank inside the group + rank = torch.distributed.get_rank(group=pg) + world_size = torch.distributed.get_world_size(group=pg) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + else: + rank = pg.rank + world_size = pg.world_size + ranks = list(range(world_size)) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + magic_message = b"magic_message" + shm = None + + try: + with contextlib.suppress(OSError): + if rank == source_rank: + # create a shared memory segment + shm = shared_memory.SharedMemory(create=True, size=128) + shm.buf[: len(magic_message)] = magic_message + if isinstance(pg, ProcessGroup): + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg + ) + else: + pg.broadcast_obj(shm.name, src=source_rank) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + if isinstance(pg, ProcessGroup): + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg + ) + name = recv[0] + else: + name = pg.broadcast_obj(None, src=source_rank) + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[: len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as e: + logger.error("Error ignored in is_in_the_same_node: %s", e) + finally: + if shm: + shm.close() + + if isinstance(pg, ProcessGroup): + torch.distributed.barrier(group=pg) + else: + pg.barrier() + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + + if isinstance(pg, ProcessGroup): + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + aggregated_data = is_in_the_same_node + else: + aggregated_data = torch.zeros_like(is_in_the_same_node) + for i in range(world_size): + rank_data = pg.broadcast_obj(is_in_the_same_node, src=i) + aggregated_data += rank_data + + return [x == 1 for x in aggregated_data.tolist()] + + +def initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % tensor_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + ) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_model_parallel_group( + tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name, + ) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % sequence_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = world_size // sequence_model_parallel_size + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range( + i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size + ) + ) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_model_parallel_group( + sp_group_ranks, get_world_group().local_rank, backend, group_name=group_name + ) + + return sp_group diff --git a/wan/distributed/utils.py b/wan/distributed/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3a7af7216bbfe5bed09d6f8828aa51e8355de3e7 --- /dev/null +++ b/wan/distributed/utils.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +import dataclasses +import pickle +import time +from collections import deque +from collections.abc import Sequence +from typing import Any + +import torch +from torch.distributed import TCPStore + +from distributed.logger import init_logger + +logger = init_logger(__name__) + + +def ensure_divisibility(numerator, denominator) -> None: + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, "{} is not divisible by {}".format( + numerator, denominator + ) + + +def divide(numerator: int, denominator: int) -> int: + """Ensure that numerator is divisible by the denominator and return + the division value.""" + ensure_divisibility(numerator, denominator) + return numerator // denominator + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> Sequence[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = divide(tensor.size()[last_dim], num_partitions) + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # NOTE: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tuple(tensor_list) + + +@dataclasses.dataclass +class StatelessProcessGroup: + """A dataclass to hold a metadata store, and the rank, world_size of the + group. Only use it to communicate metadata between processes. + For data-plane communication, create NCCL-related objects. + """ + + rank: int + world_size: int + store: torch._C._distributed_c10d.Store + data_expiration_seconds: int = 3600 # 1 hour + + # dst rank -> counter + send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict) + # src rank -> counter + recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) + broadcast_send_counter: int = 0 + broadcast_recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) + + # A deque to store the data entries, with key and timestamp. + entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque) + + def __post_init__(self): + assert self.rank < self.world_size + self.send_dst_counter = {i: 0 for i in range(self.world_size)} + self.recv_src_counter = {i: 0 for i in range(self.world_size)} + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def send_obj(self, obj: Any, dst: int): + """Send an object to a destination rank.""" + self.expire_data() + key = f"send_to/{dst}/{self.send_dst_counter[dst]}" + self.store.set(key, pickle.dumps(obj)) + self.send_dst_counter[dst] += 1 + self.entries.append((key, time.perf_counter())) + + def expire_data(self) -> None: + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.perf_counter() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def recv_obj(self, src: int) -> Any: + """Receive an object from a source rank.""" + obj = pickle.loads( + self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}") + ) + self.recv_src_counter[src] += 1 + return obj + + def broadcast_obj(self, obj: Any | None, src: int) -> Any: + """Broadcast an object from a source rank to all other ranks. + It does not clean up after all ranks have received the object. + Use it for limited times, e.g., for initialization. + """ + if self.rank == src: + self.expire_data() + key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.perf_counter())) + return obj + else: + key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + recv_obj = pickle.loads(self.store.get(key)) + self.broadcast_recv_src_counter[src] += 1 + return recv_obj + + def all_gather_obj(self, obj: Any) -> list[Any]: + """All gather an object from all ranks.""" + gathered_objs = [] + for i in range(self.world_size): + if i == self.rank: + gathered_objs.append(obj) + self.broadcast_obj(obj, src=self.rank) + else: + recv_obj = self.broadcast_obj(None, src=i) + gathered_objs.append(recv_obj) + return gathered_objs + + def barrier(self): + """A barrier to synchronize all ranks.""" + for i in range(self.world_size): + if i == self.rank: + self.broadcast_obj(None, src=self.rank) + else: + self.broadcast_obj(None, src=i) + + @staticmethod + def create( + host: str, + port: int, + rank: int, + world_size: int, + data_expiration_seconds: int = 3600, + ) -> "StatelessProcessGroup": + """A replacement for `torch.distributed.init_process_group` that does not + pollute the global state. + + If we have process A and process B called `torch.distributed.init_process_group` + to form a group, and then we want to form another group with process A, B, C, + D, it is not possible in PyTorch, because process A and process B have already + formed a group, and process C and process D cannot join that group. This + function is a workaround for this issue. + + `torch.distributed.init_process_group` is a global call, while this function + is a stateless call. It will return a `StatelessProcessGroup` object that can be + used for exchanging metadata. With this function, process A and process B + can call `StatelessProcessGroup.create` to form a group, and then process A, B, + C, and D can call `StatelessProcessGroup.create` to form another group. + """ # noqa + store = TCPStore( + host_name=host, + port=port, + world_size=world_size, + is_master=(rank == 0), + ) + + return StatelessProcessGroup( + rank=rank, + world_size=world_size, + store=store, + data_expiration_seconds=data_expiration_seconds, + ) diff --git a/wan/generate.py b/wan/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..56acff502373d1d33f654b1a46fdc49d1b911836 --- /dev/null +++ b/wan/generate.py @@ -0,0 +1,471 @@ +# Video generation script with distributed inference support +import torch +import sys, os +import time +import argparse +import json +from typing import List, Dict + +sys.path.append(os.path.abspath(".")) +import numpy as np + +from wan.inference.helper import ( + MyVAE, + CHUNK_SIZE, +) +# Lazy import pose utilities to avoid hyvideo's parallel state init at module load +# (hyvideo/generate.py calls initialize_parallel_state() at import time which +# conflicts with WAN's own distributed init when running in single-GPU mode) +from wan.inference.pipeline_wan_w_mem_relative_rope import WanPipeline +from wan.models.dits.arwan_w_action_w_mem_relative_rope import WanTransformer3DModel +from wan.distributed import ( + get_world_group, + maybe_init_distributed_environment_and_model_parallel, +) + +from wan.models.par_vae.tools import DistController +from wan.models.par_vae.context_parallel.wrapper_vae import DistWrapper +from tqdm import tqdm + +from diffusers.utils import export_to_video +from safetensors.torch import load_file + + + + +def sanitize_filename(text: str, max_length: int = 50) -> str: + """ + Convert text to safe filename format. + + Args: + text: Input text (e.g., prompt) + max_length: Maximum length of output filename + + Returns: + Sanitized filename string + """ + import re + + # Remove or replace invalid filename characters + # Strip special characters that are unsafe for filenames + text = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "", text) + # Replace spaces and commas with underscores + text = re.sub(r"[\s,]+", "_", text) + # Remove multiple consecutive underscores + text = re.sub(r"_+", "_", text) + # Strip leading/trailing underscores + text = text.strip("_") + + # Truncate if too long + if len(text) > max_length: + text = text[:max_length].rstrip("_") + + return text + + +class WanRunner: + def __init__( + self, + model_id, + ckpt_path, + ar_model_path, + ): + # Get distributed environment info from torchrun + # Extract rank information from environment variables + self.rank = int(os.environ.get("RANK", 0)) + self.local_rank = int(os.environ.get("LOCAL_RANK", 0)) + self.world_size = int(os.environ.get("WORLD_SIZE", 1)) + + self.model_id = model_id + self.ckpt_path = ckpt_path + self.ar_model_path = ar_model_path + + # Set device + # Assign CUDA device based on local rank for multi-GPU setup + device = torch.device( + f"cuda:{self.local_rank}" if torch.cuda.is_available() else "cpu" + ) + self.device = device + torch.cuda.set_device(self.local_rank) + + # Initialize distributed environment - needed even for single GPU + # Setup distributed training for multi-process inference + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl") + + # Initialize model parallel with project's custom function + # Enable sequence parallelism (sp_size=1 for single GPU) + maybe_init_distributed_environment_and_model_parallel( + 1, + sp_size=self.world_size, + distributed_init_method="env://", + ) + + # Initialize models + self._init_models() + + def _init_models(self): + # Load VAE for encoding/decoding video frames + self.vae = ( + MyVAE.from_pretrained( + self.model_id, subfolder="vae", torch_dtype=torch.bfloat16 + ) + .eval() + .requires_grad_(False) + ) + + # Load main diffusion pipeline - load directly to GPU + self.pipe = WanPipeline.from_pretrained( + self.model_id, vae=self.vae, torch_dtype=torch.bfloat16 + ) + # Move pipeline to GPU immediately after loading each component + self.pipe.to(self.device) + + # Reuse the same VAE for distributed wrapper instead of loading another copy + dist_vae = self.vae + vae_chunk_dim = 4 # chunk width dim for vae input + # Setup VAE for context-parallel decoding across GPUs + dist_controller = DistController( + self.rank, self.local_rank, self.world_size, None + ) + dist_vae = DistWrapper(dist_vae, dist_controller, None, vae_chunk_dim) + self.pipe.dist_vae = dist_vae + + # Load autoregressive transformer with action conditioning + transformer_ar_action = WanTransformer3DModel.from_pretrained( + self.ar_model_path, + use_safetensors=True, + local_files_only=True, + torch_dtype=torch.bfloat16, + ) + # Add action parameters to enable camera control + transformer_ar_action.add_discrete_action_parameters() + + # Load trained checkpoint weights + # Load directly to GPU to avoid OOM on systems with limited RAM + state_dict = torch.load( + self.ckpt_path, map_location=self.device + ) + + state_dict = state_dict["generator"] + # Clean up state dict keys from training artifacts + state_dict = { + key.replace("model.", "") if key.startswith("model.") else key: value + for key, value in state_dict.items() + } + state_dict = { + ( + key.replace("_fsdp_wrapped_module.", "") + if key.startswith("_fsdp_wrapped_module.") + else key + ): value + for key, value in state_dict.items() + } + + transformer_ar_action.load_state_dict(state_dict, strict=True) + # Transformer already on GPU from device_map, just ensure dtype + self.pipe.transformer = transformer_ar_action.to(dtype=torch.bfloat16) + # Ensure entire pipeline is on GPU + self.pipe.to(self.device) + + def predict(self, input_dict): + prompt = input_dict.get("prompt", "") + negative_prompt = input_dict.get("negative_prompt", "") + num_frames = input_dict.get("num_frames", 189) + num_inference_steps = input_dict.get("num_inference_steps", 50) + height = input_dict.get("height", 704) + width = input_dict.get("width", 1280) + image_path = input_dict.get("image_path", None) + use_memory = input_dict.get("use_memory", True) + context_window_length = input_dict.get("context_window_length", 16) + pose = input_dict.get("pose", "w-96") # Default to forward movement + num_chunk = input_dict.get("num_chunk", 4) + + seed = input_dict.get("seed", 0) + torch.manual_seed(seed) + + all_video = [] + run_args = dict( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=1.0, + few_step=True, + first_chunk_size=CHUNK_SIZE, + return_dict=False, + image_path=image_path, + use_memory=use_memory, + context_window_length=context_window_length, + ) + + # Convert pose string to pose json + # Lazy import to avoid hyvideo's parallel state init at module load time + # (by this point, WAN's distributed init has already completed) + from hyvideo.generate import pose_string_to_json, pose_to_input + + pose_json = pose_string_to_json(pose) + all_viewmats, all_Ks, all_action = pose_to_input(pose_json, len(pose_json)) + + all_save = [] + for chunk_i in range(num_chunk): + torch.cuda.synchronize() + begin_time = time.time() + + start_idx = chunk_i * CHUNK_SIZE + end_idx = start_idx + CHUNK_SIZE + + curr_viewmats = all_viewmats[start_idx:end_idx] + curr_Ks = all_Ks[start_idx:end_idx] + curr_action = all_action[start_idx:end_idx] + + curr_viewmat_save = curr_viewmats[0].cpu().numpy().tolist() + curr_action_save = curr_action[0].item() + curr_save = {"viewmat": curr_viewmat_save, "action": curr_action_save} + all_save.append(curr_save) + + self.pipe( + **run_args, + chunk_i=chunk_i, + viewmats=curr_viewmats.unsqueeze(0), + Ks=curr_Ks.unsqueeze(0), + action=curr_action.unsqueeze(0), + output_type="latent", + ) + + if self.rank == 0: + print( + "Generate time for chunk ", chunk_i, "is", time.time() - begin_time + ) + + # decode + for i in range(4): + decode_start_time = time.time() + video = self.pipe.decode_next_latent(output_type="np") + decode_time = time.time() - decode_start_time + if self.rank == 0 and i == 0: + print(f"Decode latent {i}: {decode_time:.4f} seconds") + all_video.append(video) + + if self.rank == 0: + video = np.concatenate(all_video, axis=1) + return {"video": video, "pose": pose} + else: + return {"video": None, "pose": pose} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num_frames", type=int, default=961) + parser.add_argument("--num_inference_steps", type=int, default=50) + parser.add_argument( + "--input", + type=str, + default="First-person view walking around ancient Athens, with Greek architecture and marble structures", + help="Model input in format 'prompt@image_path' or txt inputs file", + ) + parser.add_argument( + "--negative_prompt", + type=str, + default=( + "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰," + "最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部," + "画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面," + "杂乱的背景,三条腿,背景人很多,倒着走" + ), + ) + parser.add_argument("--image_path", type=str, default=None) + parser.add_argument( + "--out", + type=str, + default="outputs", + ) + parser.add_argument( + "--ckpt_path", + type=str, + default="/root/model.pt", + ) + parser.add_argument( + "--pose", + type=str, + default="w-96", + help="Pose string (e.g. 'w-3, right-0.5, d-4') or path to pose JSON file", + ) + parser.add_argument("--ar_model_path", type=str, default=None) + parser.add_argument("--num_chunk", type=int, default=4, help="Number of chunks to generate") + parser.add_argument( + "--model_id", + type=str, + default="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + help="Model ID for the pretrained model", + ) + + args = parser.parse_args() + + # Get rank info from torchrun environment + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + + if rank == 0: + print(f"Running with world_size={world_size}") + + # Initialize runner (distributed initialization handled inside) + runner = WanRunner( + model_id=args.model_id, + ckpt_path=args.ckpt_path, + ar_model_path=args.ar_model_path, + ) + + # Process input arguments (only rank 0 needs to parse inputs) + if rank == 0: + model_inputs = [] + if os.path.isfile(args.input): + # Read from txt file + with open(args.input, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + model_inputs.append(line) + else: + # Use direct input + model_inputs = [args.input] + + # Parse model inputs + parsed_inputs = [] + for model_input in model_inputs: + if "@" in model_input: + prompt, image_path = model_input.split("@", 1) + parsed_inputs.append( + {"prompt": prompt, "image_path": image_path if image_path else None} + ) + else: + parsed_inputs.append({"prompt": model_input, "image_path": None}) + + # Prepare error log file + save_dir = args.out + os.makedirs(save_dir, exist_ok=True) + error_log_path = os.path.join(save_dir, "err.txt") + else: + parsed_inputs = [] + save_dir = None + error_log_path = None + + # Broadcast inputs to all ranks if using distributed + if world_size > 1: + import torch.distributed as dist + + # Prepare data for broadcast + if rank == 0: + broadcast_data = { + "parsed_inputs": parsed_inputs, + "pose": args.pose, + "save_dir": save_dir, + "error_log_path": error_log_path, + } + else: + broadcast_data = None + + # Broadcast using object list + object_list = [broadcast_data] + dist.broadcast_object_list(object_list, src=0) + broadcast_data = object_list[0] + + parsed_inputs = broadcast_data["parsed_inputs"] + pose = broadcast_data["pose"] + save_dir = broadcast_data["save_dir"] + error_log_path = broadcast_data["error_log_path"] + else: + pose = args.pose + + try: + # Run inference + for i, input_data in enumerate(parsed_inputs): + seed_idx = 0 + if rank == 0: + prompt_prefix = sanitize_filename(input_data["prompt"], max_length=50) + # Create filename based on pose string (sanitized) + pose_prefix = sanitize_filename( + pose if not pose.endswith(".json") else os.path.basename(pose), + max_length=30, + ) + output_filename = f"{pose_prefix}_{prompt_prefix}.mp4" + output_path = os.path.join(save_dir, output_filename) + + # Check if output already exists + if os.path.exists(output_path): + print(f"\n{'='*80}") + print(f"Skipping case {i}, pose: {pose}, seed: {seed_idx}") + print(f"Output already exists: {output_filename}") + print(f"{'='*80}\n") + continue + + try: + # Check if image path is provided and exists + image_path = input_data["image_path"] or args.image_path + if rank == 0 and image_path and not os.path.exists(image_path): + print(f"\n{'='*80}") + print(f"Skipping case {i}, pose: {pose}, seed: {seed_idx}") + print(f"Image path does not exist: {image_path}") + print(f"{'='*80}\n") + continue + + input_dict = { + "prompt": input_data["prompt"], + "negative_prompt": args.negative_prompt, + "num_frames": args.num_frames, + "num_inference_steps": args.num_inference_steps, + "guidance_scale": 1, + "height": 704, + "width": 1280, + "image_path": image_path, + "use_memory": True, + "context_window_length": 16, + "seed": seed_idx, + "pose": pose, + "num_chunk": args.num_chunk, + } + + if rank == 0: + print(f"\n{'='*80}") + print(f"Processing case {i}, pose: {pose}, seed: {seed_idx}") + print(f"{'='*80}\n") + + start_time = time.time() + result = runner.predict(input_dict) + + # Save results (only rank 0) + if rank == 0: + video = result["video"] + export_to_video(video[0], output_path, fps=16) + print(f"Video saved to {output_path}") + print("Generate time:", time.time() - start_time) + + except Exception as e: + import traceback + if rank == 0: + error_msg = ( + f"\n{'='*80}\n" + f"Error occurred:\n" + f" Case: {i}\n" + f" Pose: {pose}\n" + f" Seed: {seed_idx}\n" + f" Prompt: {input_data['prompt']}\n" + f" Error: {str(e)}\n" + f" Traceback:\n{traceback.format_exc()}\n" + f"{'='*80}\n" + ) + print(error_msg) + + # Append error to log file + with open(error_log_path, "a", encoding="utf-8") as f: + f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {error_msg}\n") + + # Continue with next case + continue + + finally: + # Cleanup distributed environment + if world_size > 1 and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/wan/inference/helper.py b/wan/inference/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..33159616efcb57b6e4e7bf3108d803f363122fb7 --- /dev/null +++ b/wan/inference/helper.py @@ -0,0 +1,70 @@ +from diffusers import AutoencoderKLWan +import numpy as np +import torch +from scipy.spatial.transform import Rotation +import diffusers.models.autoencoders.vae as diff_vae +import diffusers.models.autoencoders.autoencoder_kl_wan as diff_wan_vae + +# from inference.gen_traj import generate_camera_trajectory_local + +CHUNK_SIZE = 4 +KS = [ + [ + [0.5051, 0.0000, 0.5000], + [0.0000, 0.8979, 0.5000], + [0.0000, 0.0000, 1.0000], + ] +] + + +class MyVAE(AutoencoderKLWan): + def decode(self, z, return_dict=False, is_first_chunk=False): + # self.enable_tiling() + return (self._decode(z, is_first_chunk=is_first_chunk).sample,) + + def _decode(self, z: torch.Tensor, return_dict: bool = True, is_first_chunk=False): + _, _, num_frame, height, width = z.shape + tile_latent_min_height = ( + self.tile_sample_min_height // self.spatial_compression_ratio + ) + tile_latent_min_width = ( + self.tile_sample_min_width // self.spatial_compression_ratio + ) + + if self.use_tiling and ( + width > tile_latent_min_width or height > tile_latent_min_height + ): + return self.tiled_decode(z, return_dict=return_dict) + + # print('use my vae with is_first_chunk', is_first_chunk) + if is_first_chunk: + self.clear_cache() + + x = self.post_quant_conv(z) + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=is_first_chunk, + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_], 2) + + if self.config.patch_size is not None: + out = diff_wan_vae.unpatchify(out, patch_size=self.config.patch_size) + + out = torch.clamp(out, min=-1.0, max=1.0) + + # self.clear_cache() + if not return_dict: + return (out,) + + return diff_vae.DecoderOutput(sample=out) diff --git a/wan/inference/pipeline_wan_w_mem_relative_rope.py b/wan/inference/pipeline_wan_w_mem_relative_rope.py new file mode 100644 index 0000000000000000000000000000000000000000..bd7a4e1e12cd5fcceb16221916fd91a589bf1dc4 --- /dev/null +++ b/wan/inference/pipeline_wan_w_mem_relative_rope.py @@ -0,0 +1,1110 @@ +import html + +# import json +from typing import Any, Callable, Dict, List, Optional, Union + +import os +import time +import regex as re +import torch +from PIL import Image +from tqdm import tqdm +import numpy as np +from transformers import AutoTokenizer, UMT5EncoderModel + +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.loaders import WanLoraLoaderMixin +from diffusers.models import AutoencoderKLWan, WanTransformer3DModel +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import ( + is_ftfy_available, + is_torch_xla_available, + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.wan.pipeline_output import WanPipelineOutput + +from models.utils import ( + shard_latents_dim_across_sp, + select_mem_frames_wan, +) +from hyvideo.utils.retrieval_context import generate_points_in_sphere +from distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_parallel_local_rank, + get_sp_world_size, +) +from distributed.communication_op import sequence_model_parallel_all_gather +from models.par_vae.tools import distribute_data_to_gpus, gather_results +from inference.helper import CHUNK_SIZE + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_ftfy_available(): + import ftfy + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + text = whitespace_clean(basic_clean(text)) + return text + + +class WanPipeline(DiffusionPipeline, WanLoraLoaderMixin): + model_cpu_offload_seq = "text_encoder->transformer->transformer_2->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + _optional_components = ["transformer", "transformer_2"] + + def __init__( + self, + tokenizer: AutoTokenizer, + text_encoder: UMT5EncoderModel, + vae: AutoencoderKLWan, + scheduler: FlowMatchEulerDiscreteScheduler, + transformer: Optional[WanTransformer3DModel] = None, + transformer_2: Optional[WanTransformer3DModel] = None, + boundary_ratio: Optional[float] = None, + expand_timesteps: bool = False, # Wan2.2 ti2v + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + transformer_2=transformer_2, + ) + self.register_to_config(boundary_ratio=boundary_ratio) + self.register_to_config(expand_timesteps=expand_timesteps) + self.vae_scale_factor_temporal = ( + self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + ) + self.vae_scale_factor_spatial = ( + self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + ) + self.video_processor = VideoProcessor( + vae_scale_factor=self.vae_scale_factor_spatial + ) + if os.getenv("WORLD_SIZE", "0") != "0": + self.parallel_run = True + else: + self.parallel_run = False + self.par_vae_encode = False + self.par_vae_decode = self.parallel_run and True + + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 226, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [prompt_clean(u) for u in prompt] + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder( + text_input_ids.to(device), mask.to(device) + ).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [ + torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0, + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view( + batch_size * num_videos_per_prompt, seq_len, -1 + ) + + return prompt_embeds + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + max_sequence_length: int = 226, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = ( + batch_size * [negative_prompt] + if isinstance(negative_prompt, str) + else negative_prompt + ) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, negative_prompt_embeds + + def check_inputs( + self, + prompt, + negative_prompt, + height, + width, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + guidance_scale_2=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError( + f"`height` and `width` have to be divisible by 16 but are {height} and {width}." + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, " + f"but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. " + "Please make sure to only forward one of the two." + ) + elif negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and ( + not isinstance(prompt, str) and not isinstance(prompt, list) + ): + raise ValueError( + f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" + ) + elif negative_prompt is not None and ( + not isinstance(negative_prompt, str) + and not isinstance(negative_prompt, list) + ): + raise ValueError( + f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}" + ) + + if self.config.boundary_ratio is None and guidance_scale_2 is not None: + raise ValueError( + "`guidance_scale_2` is only supported when the pipeline's `boundary_ratio` is not None." + ) + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + num_frames: int = 81, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @property + def attention_kwargs(self): + return self._attention_kwargs + + def init_kv_cache(self): + self._kv_cache = [] + self._kv_cache_neg = [] + transformer_num_layers = len(self.transformer.blocks) + for i in range(transformer_num_layers): + self._kv_cache.append({"k": None, "v": None}) + self._kv_cache_neg.append({"k": None, "v": None}) + + self.use_kv_cache = True + + def resize_center_crop_to_tensor(self, img_path, target_size): + """ + Read image -> Resize -> Center crop -> Convert to torch.tensor -> Normalize to [-1, 1] + :param img_path: Path to the image + :param target_size: (width, height) + :return: torch.Tensor, shape [C, H, W], range [-1, 1] + """ + img = Image.open(img_path).convert("RGB") + target_w, target_h = target_size + orig_w, orig_h = img.size + + # --- Step 1: proportional resize --- + scale = max(target_w / orig_w, target_h / orig_h) + new_w = int(orig_w * scale) + new_h = int(orig_h * scale) + img = img.resize((new_w, new_h), Image.LANCZOS) + + # --- Step 2: center crop --- + left = (new_w - target_w) // 2 + top = (new_h - target_h) // 2 + right = left + target_w + bottom = top + target_h + img = img.crop((left, top, right, bottom)) + + # --- Step 3: to numpy，and then tensor --- + img_np = np.array(img).astype(np.float32) # [H, W, C], 0~255 + img_np = img_np / 127.5 - 1.0 # normalize to [-1, 1] + img_tensor = torch.from_numpy(img_np).permute(2, 0, 1) # [C, H, W] + + return img_tensor + + def init_decode_state(self, all_latents, chunk_i): + """ + Initialize the decoding state and set all latents required for decoding. + """ + self._decode_state["total_latents"] = all_latents.shape[2] # all latent numbers + self._decode_state["all_latents"] = all_latents + self._decode_state["current_latent_idx"] = 0 + self._decode_state["chunk_i"] = chunk_i + assert self._decode_state["par_vae_decode"] + + if self._decode_state["latents_mean"] is None: + latents_mean = ( + torch.tensor(self.dist_vae.pipe.config.latents_mean) + .view(1, self.dist_vae.pipe.config.z_dim, 1, 1, 1) + .to(all_latents.device, all_latents.dtype) + ) + latents_std = 1.0 / torch.tensor( + self.dist_vae.pipe.config.latents_std + ).view(1, self.dist_vae.pipe.config.z_dim, 1, 1, 1).to( + all_latents.device, all_latents.dtype + ) + + self._decode_state["latents_mean"] = latents_mean + self._decode_state["latents_std"] = latents_std + + def decode_next_latent(self, output_type="np"): + """ + Decode the next latents and return the decoded frames. + If all latents have been decoded, return None. + """ + + current_idx = self._decode_state["current_latent_idx"] + total_latents = self._decode_state["total_latents"] + + if current_idx >= total_latents: + return None # all latents have been decoded + + # Get the current latent to be decoded. + all_latents = self._decode_state["all_latents"] + chunk_i = self._decode_state["chunk_i"] + current_latent = all_latents[ + :, :, current_idx : current_idx + 1, :, : + ] # [B, C, 1, H, W] + + sp_world_size = self._decode_state["sp_world_size"] + rank_in_sp_group = self._decode_state["rank_in_sp_group"] + latents_mean = self._decode_state["latents_mean"] + latents_std = self._decode_state["latents_std"] + + # Multi-GPU parallel VAE decoding + current_latent = current_latent.to(self.dist_vae.pipe.dtype) + # Ensure latents_mean and latents_std have the same type as current_latent + latents_mean = latents_mean.to(current_latent.dtype) + latents_std = latents_std.to(current_latent.dtype) + current_latent = current_latent / latents_std + latents_mean + + local_rank_in_sp_group_vae = self.ctx["local_rank_in_sp_group"] + + decode_input = distribute_data_to_gpus( + current_latent, + dim=4, + rank=rank_in_sp_group, + local_rank=local_rank_in_sp_group_vae, + num_gpus=sp_world_size, + dtype=current_latent.dtype, + ) + video = self.dist_vae.pipe.decode( + decode_input, + return_dict=False, + is_first_chunk=(current_idx == 0 and chunk_i == 0), + ) + video = gather_results( + video[0], + dim=4, + world_size=sp_world_size, + ) + + video = self.video_processor.postprocess_video(video, output_type=output_type) + + torch.cuda.synchronize() + + # Update the index + self._decode_state["current_latent_idx"] += 1 + + return video + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + height: int = 480, + width: int = 832, + num_frames: int = 81, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + guidance_scale_2: Optional[float] = None, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[ + Union[ + Callable[[int, int, Dict], None], + PipelineCallback, + MultiPipelineCallbacks, + ] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + few_step=False, + first_chunk_size=1, + image_path=None, + viewmats=None, + Ks=None, + action=None, + use_memory=True, + context_window_length=None, # mean the context length for each chunk + chunk_i: int = 0, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` + instead. Ignored when not using guidance (`guidance_scale` < `1`). + height (`int`, defaults to `480`): + The height in pixels of the generated image. + width (`int`, defaults to `832`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `81`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `5.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + guidance_scale_2 (`float`, *optional*, defaults to `None`): + Guidance scale for the low-noise stage transformer (`transformer_2`). If `None` and the pipeline's + `boundary_ratio` is not None, uses the same value as `guidance_scale`. Only used when `transformer_2` + and the pipeline's `boundary_ratio` are not None. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + output_type (`str`, *optional*, defaults to `"np"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`WanPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, defaults to `512`): + The maximum sequence length of the text encoder. If the prompt is longer than this, it will be + truncated. If the prompt is shorter, it will be padded to this length. + + Examples: + + Returns: + [`~WanPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`WanPipelineOutput`] is returned, otherwise a `tuple` is returned where + the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + if chunk_i == 0: + # print("init run") + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + negative_prompt, + height, + width, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + guidance_scale_2, + ) + + if num_frames % self.vae_scale_factor_temporal != 1: + logger.warning( + "`num_frames - 1` has to be divisible by %d. " + "Rounding to the nearest number.", + self.vae_scale_factor_temporal, + ) + num_frames = ( + num_frames + // self.vae_scale_factor_temporal + * self.vae_scale_factor_temporal + + 1 + ) + num_frames = max(num_frames, 1) + + if self.config.boundary_ratio is not None and guidance_scale_2 is None: + guidance_scale_2 = guidance_scale + + self._guidance_scale = guidance_scale + self._guidance_scale_2 = guidance_scale_2 + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + device = self._execution_device + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + self.points_local = generate_points_in_sphere(50000, 8.0).to(device) + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + ) + + transformer_dtype = ( + self.transformer.dtype + if self.transformer is not None + else self.transformer_2.dtype + ) + prompt_embeds = prompt_embeds.to(transformer_dtype) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) + + # 4. Prepare timesteps + if not few_step: + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + zeros_t = torch.zeros(1).to(timesteps.device).to(timesteps.dtype) + timesteps = torch.cat([timesteps, zeros_t]) + else: + sigmas = self.scheduler.sigmas + timesteps = torch.tensor( + [1000.0000, 960.0000, 888.8889, 727.2728, 0.0] + ).to(sigmas.device) + sigmas = timesteps / 1000.0 + + # 5. Prepare latent variables + num_channels_latents = ( + self.transformer.config.in_channels + if self.transformer is not None + else self.transformer_2.config.in_channels + ) + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + torch.float32, + device, + generator, + latents, + ) + + self._num_timesteps = len(timesteps) + chunk_size = CHUNK_SIZE + n_latent = latents.shape[2] + + # set the noise level for the inference stage + stabilization_level = 15 + + sp_world_size = 1 + rank_in_sp_group = 0 + local_rank_in_sp_group = 0 + if self.parallel_run: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + local_rank_in_sp_group = get_sp_parallel_local_rank() + + first_image_condition = None + if image_path: + first_image = self.resize_center_crop_to_tensor(image_path, (width, height)) # (B C T H W) + first_image = first_image.to(device).to(self.vae.dtype) + first_image = first_image.unsqueeze(0).unsqueeze(2) # [B C T H W] + first_image_condition = self.vae.encode(first_image, return_dict=False)[0].sample() + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( + 1, self.vae.config.z_dim, 1, 1, 1 + ).to(latents.device, latents.dtype) + first_image_condition = (first_image_condition - latents_mean) * latents_std + + + self.use_kv_cache = False + self.init_kv_cache() + + self._decode_state = { + "current_latent_idx": 0, + "total_latents": None, + "device": device, + "transformer_dtype": transformer_dtype, + "vae": self.vae, + "video_processor": self.video_processor, + "parallel_run": self.parallel_run, + "par_vae_decode": self.par_vae_decode, + "sp_world_size": sp_world_size, + "rank_in_sp_group": rank_in_sp_group, + "latents_mean": None, + "latents_std": None, + "is_initialized": False, + } + + + self.ctx = { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "guidance_scale": guidance_scale, + "guidance_scale_2": guidance_scale_2, + "attention_kwargs": attention_kwargs, + "transformer_dtype": ( + self.transformer.dtype + if self.transformer is not None + else self.transformer_2.dtype + ), + # Denoising scheduler and state + "timesteps": timesteps, + "num_inference_steps": num_inference_steps, + "latents": latents, + "first_image_condition": first_image_condition, + "chunk_size": chunk_size, + "first_chunk_size": first_chunk_size, + "n_latent": n_latent, + "viewmats": torch.zeros( + 1, n_latent, 4, 4, device=latents.device, dtype=latents.dtype + ), + "Ks": torch.zeros( + 1, n_latent, 3, 3, device=latents.device, dtype=latents.dtype + ), + "action": torch.zeros( + 1, n_latent, device=latents.device, dtype=latents.dtype + ), + "use_memory": use_memory, + "context_window_length": context_window_length, + "kv_cache": self._kv_cache, + "kv_cache_neg": self._kv_cache_neg, + "use_kv_cache": self.use_kv_cache, + # Distributed related + "sp_world_size": sp_world_size, + "rank_in_sp_group": rank_in_sp_group, + "local_rank_in_sp_group": local_rank_in_sp_group, + # Other auxiliary parameters + "device": device, + "batch_size": batch_size, + "few_step": few_step, + "sigmas": sigmas if few_step else None, + "stabilization_level": stabilization_level, + } + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # ========================================================== + if self.ctx is None and chunk_i > 0: + raise ValueError( + "When chunk_i > 0, the context self.ctx cannot be empty. Please ensure the call starts from chunk_i = 0." + ) + prompt_embeds = self.ctx["prompt_embeds"] + negative_prompt_embeds = self.ctx["negative_prompt_embeds"] + guidance_scale = self.ctx["guidance_scale"] + guidance_scale_2 = self.ctx["guidance_scale_2"] + attention_kwargs = self.ctx["attention_kwargs"] + transformer_dtype = self.ctx["transformer_dtype"] + + timesteps = self.ctx["timesteps"] + num_inference_steps = self.ctx["num_inference_steps"] + latents = self.ctx["latents"] + first_image_condition = self.ctx["first_image_condition"] + + chunk_size = self.ctx["chunk_size"] + first_chunk_size = self.ctx["first_chunk_size"] + n_latent = self.ctx["n_latent"] + stabilization_level = self.ctx["stabilization_level"] + + # Fill the poses of the current chunk into the corresponding positions. + start_idx = chunk_i * chunk_size + end_idx = start_idx + chunk_size + + self.ctx["viewmats"][:, start_idx:end_idx, ...] = viewmats + self.ctx["Ks"][:, start_idx:end_idx, ...] = Ks + self.ctx["action"][:, start_idx:end_idx, ...] = action + + viewmats = self.ctx["viewmats"] + Ks = self.ctx["Ks"] + action = self.ctx["action"] + + use_memory = self.ctx["use_memory"] + context_window_length = self.ctx["context_window_length"] + self._kv_cache = self.ctx["kv_cache"] + self._kv_cache_neg = self.ctx["kv_cache_neg"] + self.use_kv_cache = self.ctx["use_kv_cache"] + + sp_world_size_vae = self.ctx["sp_world_size"] + rank_in_sp_group_vae = self.ctx["rank_in_sp_group"] + local_rank_in_sp_group_vae = self.ctx["local_rank_in_sp_group"] + + device = self.ctx["device"] + batch_size = self.ctx["batch_size"] + few_step = self.ctx["few_step"] + sigmas = self.ctx["sigmas"] + + + # Restore the attributes of the pipeline + self._guidance_scale = guidance_scale + self._guidance_scale_2 = guidance_scale_2 + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + self._num_timesteps = len(timesteps) + + selected_frame_indices = None + if chunk_i == 0: + already_generate_num = first_chunk_size + generate_latent_num = first_chunk_size + if first_image_condition is not None: + latents[:, :, :1] = first_image_condition + latents_curr = latents[:, :, :already_generate_num].to(device=device, dtype=prompt_embeds.dtype) + else: + if not few_step: + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + zeros_t = torch.zeros(1).to(timesteps.device).to(timesteps.dtype) + timesteps = torch.cat([timesteps, zeros_t]) + + already_generate_num = chunk_i * chunk_size + first_chunk_size + latents_curr = latents[:, :, :already_generate_num].to(device=device, dtype=prompt_embeds.dtype) + generate_latent_num = chunk_size + + # add memory selection here + current_frame_idx = chunk_i * chunk_size # the current frame index to generate + if ( + use_memory + and (current_frame_idx >= context_window_length) + and current_frame_idx < n_latent + ): + selected_frame_indices = select_mem_frames_wan( + viewmats[0].cpu().detach().numpy(), + current_frame_idx, + memory_frames=16, + temporal_context_size=12, + pred_latent_size=4, + points_local=self.points_local, + device=self.device, + ) + + elif use_memory: + selected_frame_indices = list(range(0, current_frame_idx)) + + with self.progress_bar(total=4) as progress_bar: + for i, t in enumerate(timesteps[:-1]): + if self.interrupt: + continue + + self._current_timestep = t + + current_model = self.transformer + current_guidance_scale = guidance_scale + + if chunk_i > 0: # context latent frame + t_now = torch.full((batch_size, chunk_size), t, device=device, dtype=timesteps.dtype) + t_ctx = torch.full((batch_size, first_chunk_size + (chunk_i - 1) * chunk_size), + stabilization_level - 1, + device=device, + dtype=timesteps.dtype, + ) # [B, F] + timestep = torch.cat([t_ctx, t_now], dim=1) # [B, F+1]: predict the next frame in timestep t + else: + if first_image_condition is not None: + t_now = torch.full((batch_size, chunk_size - 1), t, device=device, dtype=timesteps.dtype) + t_ctx = torch.full((batch_size, 1), stabilization_level - 1, + device=device, + dtype=timesteps.dtype, + ) # [B, F] + timestep = torch.cat([t_ctx, t_now], dim=1) + else: + timestep = torch.full((batch_size, first_chunk_size), t, + device=device, + dtype=timesteps.dtype, + ) # [B, 1] + + all_window_latent_num = latents_curr.shape[2] + + ## kv cache + if chunk_i > 0 and i == 0 and self.use_kv_cache: + latents_cache = latents_curr.clone() + latents_cache = latents_cache[:, :, selected_frame_indices] + t_cache = timestep[:, selected_frame_indices] + + latent_model_cache = latents_cache.to(transformer_dtype) + # flatten the timestep for the transformer + timestep_cache = t_cache.flatten() + if action is not None: + action_cache = action[:, selected_frame_indices] + viewmats_cache = viewmats[:, selected_frame_indices] + Ks_cache = Ks[:, selected_frame_indices] + + kv_start_rope = 0 + kv_end_rope = len(selected_frame_indices) + + torch.cuda.nvtx.range_push("kv_cache") + self._kv_cache = current_model( + hidden_states=latent_model_cache, + timestep=timestep_cache, + encoder_hidden_states=prompt_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + current_start=kv_start_rope * 880, + current_end=kv_end_rope * 880, + kv_cache=self._kv_cache, + is_cache=True, + window_frames=n_latent, + viewmats=(viewmats_cache if viewmats is not None else None), + Ks=Ks_cache if Ks is not None else None, + action=action_cache if action is not None else None, + ) + + torch.cuda.nvtx.range_pop() + + if not self.use_kv_cache: + generate_latent_num = all_window_latent_num + + if selected_frame_indices is not None: + now_window_latent_num = len(selected_frame_indices) + chunk_size + else: + now_window_latent_num = chunk_size + + latent_model_input = latents_curr.clone() + latent_model_input = latent_model_input[:, :, -generate_latent_num:].to(transformer_dtype) + timestep = timestep[:, -generate_latent_num:] + timestep = timestep.flatten() + + generate_rope_start = (now_window_latent_num - generate_latent_num) * 880 + generate_rope_end = now_window_latent_num * 880 + + torch.cuda.nvtx.range_push("noise_pred") + noise_pred = current_model( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + current_start=generate_rope_start, + current_end=generate_rope_end, + kv_cache=self._kv_cache if self.use_kv_cache else None, + is_cache=False, + window_frames=n_latent, + viewmats=viewmats[:, all_window_latent_num - generate_latent_num: all_window_latent_num], + Ks=Ks[:, all_window_latent_num - generate_latent_num: all_window_latent_num], + action=action[:, all_window_latent_num - generate_latent_num: all_window_latent_num], + )[0] + + torch.cuda.nvtx.range_pop() + + # compute the previous noisy sample x_t -> x_t-1 + if chunk_i == 0: + if not few_step: + latent_model_input = self.scheduler.step( noise_pred, t, latent_model_input, return_dict=False)[0] + if first_image_condition is not None: + latents_curr[:, :, -first_chunk_size + 1:] = latent_model_input[:, :, 1:] + else: + latents_curr[:, :, -first_chunk_size:] = latent_model_input + else: + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + dt = sigma_next - sigma + prev_sample = latent_model_input + dt * noise_pred + if first_image_condition is not None: + latents_curr[:, :, -first_chunk_size + 1:] = prev_sample[:, :, 1:] + else: + latents_curr[:, :, -first_chunk_size:] = prev_sample + else: + noise_pred_chunk = noise_pred[:, :, -chunk_size:] # [B, C, chunk_size, H, W + latents_curr_chunk = latent_model_input[:, :, -chunk_size:] # [B, C, chunk_size, H, W] + if not few_step: + latent_curr_pred_chunk = self.scheduler.step( + noise_pred_chunk, + t, + latents_curr_chunk, + return_dict=False, + )[0] # [B, C, chunk_size, H, W] + else: + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + dt = sigma_next - sigma + latent_curr_pred_chunk = latents_curr_chunk + dt * noise_pred_chunk + + latents_curr[:, :, -chunk_size:] = latent_curr_pred_chunk # update the latents + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) + + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + latents[:, :, :already_generate_num, :, :] = latents_curr + + # ========================================================== + # Save the state updated during this call back to self.ctx + self.ctx["latents"] = latents + self.ctx["kv_cache"] = self._kv_cache + self.ctx["kv_cache_neg"] = self._kv_cache_neg + # ========================================================== + + re_latents = latents + self._current_timestep = None + latents = latents[:, :, start_idx:end_idx, :, :] + + # Automatically initialize the decoding state (for subsequent frame-by-frame decoding) + # Note: Must be initialized here because the latents will be modified later + if output_type == "latent": + self.init_decode_state(latents, chunk_i) + + torch.cuda.synchronize() + vae_decode_begin_time = time.time() + if not output_type == "latent": + if not self.par_vae_decode: + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( + 1, self.vae.config.z_dim, 1, 1, 1 + ).to(latents.device, latents.dtype) + latents = latents / latents_std + latents_mean + video = self.vae.decode( + latents, return_dict=False, is_first_chunk=chunk_i == 0 + )[0] + else: + latents = latents.to(self.dist_vae.pipe.dtype) + latents_mean = ( + torch.tensor(self.dist_vae.pipe.config.latents_mean) + .view(1, self.dist_vae.pipe.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor( + self.dist_vae.pipe.config.latents_std + ).view(1, self.dist_vae.pipe.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + + # Accelerated VAE + decode_input = distribute_data_to_gpus( + latents, + dim=4, + rank=rank_in_sp_group_vae, + local_rank=local_rank_in_sp_group_vae, + num_gpus=sp_world_size_vae, + dtype=latents.dtype, + ) + video = self.dist_vae.pipe.decode( + decode_input, return_dict=False, is_first_chunk=chunk_i == 0 + ) + video = gather_results( + video[0], + dim=4, + world_size=sp_world_size_vae, + ) + video = self.video_processor.postprocess_video( + video, output_type=output_type + ) + else: + return None, None + + torch.cuda.synchronize() + print("Vae decode time", time.time() - vae_decode_begin_time) + torch.cuda.nvtx.range_pop() + + if not return_dict: + return ( + video, + re_latents, + ) + + return WanPipelineOutput(frames=video) diff --git a/wan/logging_utils/__init__.py b/wan/logging_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a7f50aab6b7c9c9294684b5e348408ea47c12c5 --- /dev/null +++ b/wan/logging_utils/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 + +from logging_utils.formatter import NewLineFormatter +from logging_utils.formatter import setup_for_distributed + +__all__ = [ + "NewLineFormatter", + "setup_for_distributed", +] diff --git a/wan/logging_utils/formatter.py b/wan/logging_utils/formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..53b547890c4059435c5ed438303e5d4168301038 --- /dev/null +++ b/wan/logging_utils/formatter.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logging_utils/formatter.py + +import logging + + +class NewLineFormatter(logging.Formatter): + """Adds logging prefix to newlines to align multi-line messages.""" + + def __init__(self, fmt, datefmt=None, style="%"): + logging.Formatter.__init__(self, fmt, datefmt, style) + + def format(self, record): + msg = logging.Formatter.format(self, record) + if record.message != "": + parts = msg.split(record.message) + msg = msg.replace("\n", "\r\n" + parts[0]) + return msg + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print diff --git a/wan/models/dits/arwan_w_action_w_mem_relative_rope.py b/wan/models/dits/arwan_w_action_w_mem_relative_rope.py new file mode 100644 index 0000000000000000000000000000000000000000..eb13e65d3c0d9caaf2b30eabba8cfac3f74a300f --- /dev/null +++ b/wan/models/dits/arwan_w_action_w_mem_relative_rope.py @@ -0,0 +1,1039 @@ +# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from platform import processor +from typing import Any, Dict, Optional, Tuple, Union, List + +import os +import torch +import json +import datetime +import torch.nn as nn +import torch.nn.functional as F + +from hyvideo.prope.camera_rope import prope_qkv +from distributed.communication_op import ( + all_to_all_sp, + sequence_model_parallel_all_gather, +) +from distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin +from diffusers.utils import ( + USE_PEFT_BACKEND, + logging, + scale_lora_layers, + unscale_lora_layers, +) +from diffusers.utils.torch_utils import maybe_allow_in_graph +from diffusers.models.attention import FeedForward +from diffusers.models.attention_processor import Attention +from diffusers.models.cache_utils import CacheMixin +from diffusers.models.embeddings import ( + PixArtAlphaTextProjection, + TimestepEmbedding, + Timesteps, + get_1d_rotary_pos_embed, +) +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import FP32LayerNorm + +from sageattention import sageattn + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class CausalCameraPRopeWanAttnProcessor2_0: + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "WanAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0." + ) + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + rotary_emb: Optional[torch.Tensor] = None, + kv_cache: Optional[dict] = None, + is_cache: Optional[bool] = False, + idx: Optional[int] = None, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + context_frames_list: Optional[List[int]] = None, + ) -> torch.Tensor: + encoder_hidden_states_img = None + if attn.add_k_proj is not None: + # 512 is the context length of the text encoder, hardcoded for now + image_context_length = encoder_hidden_states.shape[1] - 512 + encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] + encoder_hidden_states = encoder_hidden_states[:, image_context_length:] + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2) + key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) + value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + if rotary_emb is not None: + + def apply_rotary_emb( + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + x = hidden_states.view(*hidden_states.shape[:-1], -1, 2) + x1, x2 = x[..., 0], x[..., 1] + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + query_rope = apply_rotary_emb(query, *rotary_emb) + key_rope = apply_rotary_emb(key, *rotary_emb) + + else: + query_rope = query + key_rope = key + + pad_size = 0 + if os.getenv("WORLD_SIZE", "0") != "0": + world_size = get_sp_world_size() + head_num = query.shape[1] + if head_num % world_size != 0: + # pad head_num + pad_size = world_size - head_num % world_size + query = torch.nn.functional.pad(query, (0, 0, 0, 0, 0, pad_size)) + key = torch.nn.functional.pad(key, (0, 0, 0, 0, 0, pad_size)) + value = torch.nn.functional.pad(value, (0, 0, 0, 0, 0, pad_size)) + query_rope = torch.nn.functional.pad( + query_rope, (0, 0, 0, 0, 0, pad_size) + ) + key_rope = torch.nn.functional.pad(key_rope, (0, 0, 0, 0, 0, pad_size)) + + query = all_to_all_sp(query, scatter_dim=1, gather_dim=2) + key = all_to_all_sp(key, scatter_dim=1, gather_dim=2) + value = all_to_all_sp(value, scatter_dim=1, gather_dim=2) + query_rope = all_to_all_sp(query_rope, scatter_dim=1, gather_dim=2) + key_rope = all_to_all_sp(key_rope, scatter_dim=1, gather_dim=2) + + value_rope = value + + query_prope, key_prope, value_prope, apply_fn_o = prope_qkv( + query, + key, + value, + viewmats=viewmats, + Ks=Ks, + patches_x=40, # hard code + patches_y=22, # hard code + ) # [batch, num_heads, seqlen, head_dim] + + + kv_cache_return = {} + cache_key = kv_cache["k"] + cache_value = kv_cache["v"] + + + if cache_value is not None and not is_cache: + cache_key_rope, cache_key_prope = cache_key.chunk(2, dim=-1) + cache_value_rope, cache_value_prope = cache_value.chunk(2, dim=-1) + + key_rope = torch.cat([cache_key_rope, key_rope], dim=-2) + value_rope = torch.cat([cache_value_rope, value_rope], dim=-2) + + key_prope = torch.cat([cache_key_prope, key_prope], dim=-2) + value_prope = torch.cat([cache_value_prope, value_prope], dim=-2) + + if is_cache: + kv_cache_return["k"] = torch.cat([key_rope, key_prope], dim=-1) + kv_cache_return["v"] = torch.cat([value_rope, value_prope], dim=-1) + + query_all = torch.cat([query_rope, query_prope], dim=0) + key_all = torch.cat([key_rope, key_prope], dim=0) + value_all = torch.cat([value_rope, value_prope], dim=0) + + hidden_states_all = sageattn( + query_all, key_all, value_all, tensor_layout="HND", is_causal=False + ) + + # hidden_states_all = F.scaled_dot_product_attention( + # query_all, key_all, value_all, dropout_p=0.0 + # ) + + hidden_states_all = hidden_states_all.transpose( + 1, 2 + ) # [batch * 2, seqlen, per_sp_num_heads, head_dim] + + hidden_states_rope, hidden_states_prope = hidden_states_all.chunk(2, dim=0) + hidden_states_prope = apply_fn_o(hidden_states_prope.transpose(1, 2)).transpose( + 1, 2 + ) + + if os.getenv("WORLD_SIZE", "0") != "0": + hidden_states_rope = all_to_all_sp( + hidden_states_rope, scatter_dim=1, gather_dim=2 + ) + hidden_states_prope = all_to_all_sp( + hidden_states_prope, scatter_dim=1, gather_dim=2 + ) + + if pad_size != 0: + hidden_states_rope, _ = hidden_states_rope.split( + [hidden_states_rope.shape[2] - pad_size, pad_size], dim=2 + ) + hidden_states_prope, _ = hidden_states_prope.split( + [hidden_states_prope.shape[2] - pad_size, pad_size], dim=2 + ) + + hidden_states_prope = hidden_states_prope.flatten(2, 3) + hidden_states_prope = hidden_states_prope.type_as(query) + hidden_states_prope = attn.to_out_prope[0](hidden_states_prope) + + hidden_states_rope = hidden_states_rope.flatten(2, 3) + hidden_states_rope = hidden_states_rope.type_as(query) + hidden_states_rope = attn.to_out[0](hidden_states_rope) + hidden_states_rope = attn.to_out[1](hidden_states_rope) + + hidden_states = hidden_states_rope + hidden_states_prope + + return hidden_states, kv_cache_return + + +class WanAttnProcessor2_0: + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "WanAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0." + ) + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + rotary_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + encoder_hidden_states_img = None + if attn.add_k_proj is not None: + # 512 is the context length of the text encoder, hardcoded for now + image_context_length = encoder_hidden_states.shape[1] - 512 + encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] + encoder_hidden_states = encoder_hidden_states[:, image_context_length:] + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2) + key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) + value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + if rotary_emb is not None: + + def apply_rotary_emb( + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + x = hidden_states.view(*hidden_states.shape[:-1], -1, 2) + x1, x2 = x[..., 0], x[..., 1] + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + query = apply_rotary_emb(query, *rotary_emb) + key = apply_rotary_emb(key, *rotary_emb) + + # I2V task + hidden_states_img = None + if encoder_hidden_states_img is not None: + assert False + key_img = attn.add_k_proj(encoder_hidden_states_img) + key_img = attn.norm_added_k(key_img) + value_img = attn.add_v_proj(encoder_hidden_states_img) + + key_img = key_img.unflatten(2, (attn.heads, -1)).transpose(1, 2) + value_img = value_img.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + hidden_states_img = F.scaled_dot_product_attention( + query, + key_img, + value_img, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + hidden_states_img = hidden_states_img.transpose(1, 2).flatten(2, 3) + hidden_states_img = hidden_states_img.type_as(query) + + # SP 16 Pad + + assert attention_mask is None + + hidden_states = sageattn( + query, key, value, tensor_layout="HND", is_causal=False + ) + # hidden_states = F.scaled_dot_product_attention( + # query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + # ) + + hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + if hidden_states_img is not None: + hidden_states = hidden_states + hidden_states_img + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class WanImageEmbedding(torch.nn.Module): + def __init__(self, in_features: int, out_features: int, pos_embed_seq_len=None): + super().__init__() + + self.norm1 = FP32LayerNorm(in_features) + self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu") + self.norm2 = FP32LayerNorm(out_features) + if pos_embed_seq_len is not None: + self.pos_embed = nn.Parameter( + torch.zeros(1, pos_embed_seq_len, in_features) + ) + else: + self.pos_embed = None + + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + if self.pos_embed is not None: + batch_size, seq_len, embed_dim = encoder_hidden_states_image.shape + encoder_hidden_states_image = encoder_hidden_states_image.view( + -1, 2 * seq_len, embed_dim + ) + encoder_hidden_states_image = encoder_hidden_states_image + self.pos_embed + + hidden_states = self.norm1(encoder_hidden_states_image) + hidden_states = self.ff(hidden_states) + hidden_states = self.norm2(hidden_states) + return hidden_states + + +# add the discrete action +class WanActionTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + action_embed_dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + image_embed_dim: Optional[int] = None, + pos_embed_seq_len: Optional[int] = None, + ): + super().__init__() + self.action_embed_dim = action_embed_dim + self.dim = dim + self.time_freq_dim = time_freq_dim + + self.timesteps_proj = Timesteps( + num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.time_embedder = TimestepEmbedding( + in_channels=time_freq_dim, time_embed_dim=dim + ) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection( + text_embed_dim, dim, act_fn="gelu_tanh" + ) + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding( + image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len + ) + + def forward( + self, + action: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + ): + timestep = self.timesteps_proj(timestep) + action = self.timesteps_proj(action.squeeze(0)) + action_embedder_dtype = next(iter(self.action_embedder.parameters())).dtype + if ( + action.dtype != action_embedder_dtype + and action_embedder_dtype != torch.int8 + ): + action = action.to(action_embedder_dtype) + action = self.action_embedder(action).type_as(encoder_hidden_states) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + + temb = temb + action + + timestep_proj = self.time_proj(self.act_fn(temb)) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image is not None: + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image + ) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + +class WanTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + image_embed_dim: Optional[int] = None, + pos_embed_seq_len: Optional[int] = None, + ): + super().__init__() + + self.timesteps_proj = Timesteps( + num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.time_embedder = TimestepEmbedding( + in_channels=time_freq_dim, time_embed_dim=dim + ) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection( + text_embed_dim, dim, act_fn="gelu_tanh" + ) + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding( + image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len + ) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + ): + timestep = self.timesteps_proj(timestep) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image is not None: + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image + ) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: Tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + freqs_dtype = ( + torch.float32 if torch.backends.mps.is_available() else torch.float64 + ) + + freqs_cos = [] + freqs_sin = [] + + for dim in [t_dim, h_dim, w_dim]: + freq_cos, freq_sin = get_1d_rotary_pos_embed( + dim, + max_seq_len, + theta, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + freqs_cos.append(freq_cos) + freqs_sin.append(freq_sin) + + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w + + split_sizes = [ + self.attention_head_dim - 2 * (self.attention_head_dim // 3), + self.attention_head_dim // 3, + self.attention_head_dim // 3, + ] + + freqs_cos = self.freqs_cos.split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.split(split_sizes, dim=1) + + freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape( + 1, 1, ppf * pph * ppw, -1 + ) + freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape( + 1, 1, ppf * pph * ppw, -1 + ) + + return freqs_cos, freqs_sin + + +@maybe_allow_in_graph +class WanTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: Optional[int] = None, + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = Attention( + query_dim=dim, + heads=num_heads, + kv_heads=num_heads, + dim_head=dim // num_heads, + qk_norm=qk_norm, + eps=eps, + bias=True, + cross_attention_dim=None, + out_bias=True, + processor=CausalCameraPRopeWanAttnProcessor2_0(), + ) + + # 2. Cross-attention + self.attn2 = Attention( + query_dim=dim, + heads=num_heads, + kv_heads=num_heads, + dim_head=dim // num_heads, + qk_norm=qk_norm, + eps=eps, + bias=True, + cross_attention_dim=None, + out_bias=True, + added_kv_proj_dim=added_kv_proj_dim, + added_proj_bias=True, + processor=WanAttnProcessor2_0(), + ) + self.norm2 = ( + FP32LayerNorm(dim, eps, elementwise_affine=True) + if cross_attn_norm + else nn.Identity() + ) + + # 3. Feed-forward + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + @torch.compile + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: torch.Tensor, + kv_cache, + is_cache, + idx, + viewmats, + Ks, + context_frames_list=None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb.float() + ).chunk(6, dim=1) + + # Turn bidrectional self-attention to causal self-attention + + # design special adaln for diffusion forcing + if hidden_states.shape[0] != temb.shape[0]: + # print("Diffusion forcing trainning! Adjust the AdaLN operation!") + latent_length = temb.shape[0] // hidden_states.shape[0] # latent length + token_length = hidden_states.shape[1] // latent_length + if os.getenv("WORLD_SIZE", "0") != "0": + token_length = token_length * get_sp_world_size() + shift_msa = shift_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + scale_msa = scale_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + gate_msa = gate_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + c_shift_msa = c_shift_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + c_scale_msa = c_scale_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + c_gate_msa = c_gate_msa.reshape( + hidden_states.shape[0], -1, hidden_states.shape[2] + ) + shift_msa = shift_msa.repeat_interleave(token_length, dim=1) + scale_msa = scale_msa.repeat_interleave(token_length, dim=1) + gate_msa = gate_msa.repeat_interleave(token_length, dim=1) + c_shift_msa = c_shift_msa.repeat_interleave(token_length, dim=1) + c_scale_msa = c_scale_msa.repeat_interleave(token_length, dim=1) + c_gate_msa = c_gate_msa.repeat_interleave(token_length, dim=1) + + if os.getenv("WORLD_SIZE", "0") != "0": + world_size = get_sp_world_size() + rank = get_sp_parallel_rank() + shift_msa = shift_msa.chunk(world_size, dim=1)[rank] + scale_msa = scale_msa.chunk(world_size, dim=1)[rank] + gate_msa = gate_msa.chunk(world_size, dim=1)[rank] + c_shift_msa = c_shift_msa.chunk(world_size, dim=1)[rank] + c_scale_msa = c_scale_msa.chunk(world_size, dim=1)[rank] + c_gate_msa = c_gate_msa.chunk(world_size, dim=1)[rank] + + # 1. Self-attention + norm_hidden_states = ( + self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa + ).type_as(hidden_states) + attn_output, kv_cache_return = self.attn1( + hidden_states=norm_hidden_states, + rotary_emb=rotary_emb, + kv_cache=kv_cache, + is_cache=is_cache, + idx=idx, + viewmats=viewmats, + Ks=Ks, + context_frames_list=context_frames_list, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as( + hidden_states + ) + + # 2. Cross-attention + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + hidden_states=norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + ) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = ( + self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa + ).type_as(hidden_states) + ff_output = self.ffn(norm_hidden_states) + hidden_states = ( + hidden_states.float() + ff_output.float() * c_gate_msa + ).type_as(hidden_states) + + return hidden_states, kv_cache_return + + +class WanTransformer3DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin +): + r""" + A Transformer model for video-like data used in the Wan model. + + Args: + patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). + num_attention_heads (`int`, defaults to `40`): + Fixed length for text embeddings. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + text_dim (`int`, defaults to `512`): + Input dimension for text embeddings. + freq_dim (`int`, defaults to `256`): + Dimension for sinusoidal time embeddings. + ffn_dim (`int`, defaults to `13824`): + Intermediate dimension in feed-forward network. + num_layers (`int`, defaults to `40`): + The number of layers of transformer blocks to use. + window_size (`Tuple[int]`, defaults to `(-1, -1)`): + Window size for local attention (-1 indicates global attention). + cross_attn_norm (`bool`, defaults to `True`): + Enable cross-attention normalization. + qk_norm (`bool`, defaults to `True`): + Enable query/key normalization. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + add_img_emb (`bool`, defaults to `False`): + Whether to use img_emb. + added_kv_proj_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the added key and value projections. If `None`, no projection is used. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] + _no_split_modules = ["WanTransformerBlock"] + _keep_in_fp32_modules = [ + "time_embedder", + "scale_shift_table", + "norm1", + "norm2", + "norm3", + ] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock"] + + @register_to_config + def __init__( + self, + patch_size: Tuple[int] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 16, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + qk_norm: Optional[str] = "rms_norm_across_heads", + eps: float = 1e-6, + image_dim: Optional[int] = None, + added_kv_proj_dim: Optional[int] = None, + rope_max_seq_len: int = 1024, + pos_embed_seq_len: Optional[int] = None, + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + # 1. Patch & position embedding + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding = nn.Conv3d( + in_channels, inner_dim, kernel_size=patch_size, stride=patch_size + ) + + # 2. Condition embeddings + # image_embedding_dim=1280 for I2V model + self.condition_embedder = WanActionTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + image_embed_dim=image_dim, + pos_embed_seq_len=pos_embed_seq_len, + action_embed_dim=8, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList( + [ + WanTransformerBlock( + inner_dim, + ffn_dim, + num_attention_heads, + qk_norm, + cross_attn_norm, + eps, + added_kv_proj_dim, + ) + for _ in range(num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.scale_shift_table = nn.Parameter( + torch.randn(1, 2, inner_dim) / inner_dim**0.5 + ) + + self.gradient_checkpointing = False + self.inner_dim = inner_dim + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + current_start=None, + current_end=None, + kv_cache: dict = None, + is_cache: bool = False, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + action: Optional[torch.Tensor] = None, + window_frames: int = None, + context_frames_list: List[int] = None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + hidden_states = hidden_states.to(encoder_hidden_states.dtype) + viewmats = ( + viewmats.to(encoder_hidden_states.dtype) if viewmats is not None else None + ) + Ks = Ks.to(encoder_hidden_states.dtype) if Ks is not None else None + action = action.to(encoder_hidden_states.dtype) if action is not None else None + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if ( + attention_kwargs is not None + and attention_kwargs.get("scale", None) is not None + ): + logger.warning( + "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." + ) + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + + ## if hidden_states is smaller than the window size, we need to pad it + if post_patch_num_frames < window_frames: + padding = window_frames - post_patch_num_frames + # pad hidden_state to the window size in the dim=2: [B, C, T, H, W] -> [B, C, T + padding, H, W] + padding_hidden_states = F.pad( + hidden_states, (0, 0, 0, 0, 0, padding), value=0.0 + ) # [B, C, T + padding, H, W] + elif post_patch_num_frames == window_frames: + padding_hidden_states = hidden_states.clone() + else: + # if hidden_states is larger than the window size, we need to slice it + padding_hidden_states = hidden_states[:, :, :window_frames, :, :] + + # 得到所有的rope参数 + rotary_emb = self.rope(padding_hidden_states) # 2 * [1, ,1 , T' * H' * W']([1, 1, 21 * 30 * 52]) + + hidden_states = self.patch_embedding(hidden_states) # [B, C, T, H, W] -> [B, C', T', H', W'] + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( + self.condition_embedder( + action, timestep, encoder_hidden_states, encoder_hidden_states_image + ) + ) + timestep_proj = timestep_proj.unflatten(1, (6, -1)) # [B, 9216] -> [B, 6, 1536] + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1 + ) + + if is_cache: + _kv_cache_now = [] + transformer_num_layers = 30 + for i in range(transformer_num_layers): + _kv_cache_now.append({"k": None, "v": None}) + + rotary_emb0 = rotary_emb[0][:, :, current_start:current_end, :] + rotary_emb1 = rotary_emb[1][:, :, current_start:current_end, :] + + # SP + if os.getenv("WORLD_SIZE", "0") != "0": + world_size = get_sp_world_size() + rank = get_sp_parallel_rank() + hidden_states = hidden_states.chunk(chunks=world_size, dim=1)[rank] + rotary_emb0 = rotary_emb0.chunk(chunks=world_size, dim=2)[rank] + rotary_emb1 = rotary_emb1.chunk(chunks=world_size, dim=2)[rank] + rotary_emb = (rotary_emb0, rotary_emb1) + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for idx, block in enumerate(self.blocks): + hidden_states, t_kv = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + kv_cache[idx] if kv_cache is not None else None, + is_cache, + idx, + viewmats, + Ks, + context_frames_list=context_frames_list, + ) + if is_cache: + _kv_cache_now[idx] = t_kv + else: + for idx, block in enumerate(self.blocks): + hidden_states, t_kv = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + kv_cache[idx] if kv_cache is not None else None, + is_cache, + idx, + viewmats, + Ks, + context_frames_list=context_frames_list, + ) + if is_cache: + _kv_cache_now[idx] = t_kv + + if is_cache: + return _kv_cache_now + + if os.getenv("WORLD_SIZE", "0") != "0": + hidden_states = sequence_model_parallel_all_gather( + hidden_states.contiguous(), dim=1 + ) + + # 5. Output norm, projection & unpatchify + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + + # Move the shift and scale tensors to the same device as hidden_states. + # When using multi-GPU inference via accelerate these will be on the + # first device rather than the last device, which hidden_states ends up on. + shift = shift.to(hidden_states.device) + scale = scale.to(hidden_states.device) + + if shift.shape[0] != hidden_states.shape[0]: + # Diffusion forcing training, adjust the AdaLN operation + latent_length = temb.shape[0] // hidden_states.shape[0] + token_length = hidden_states.shape[1] // latent_length + shift = shift.reshape(hidden_states.shape[0], -1, hidden_states.shape[2]) + scale = scale.reshape(hidden_states.shape[0], -1, hidden_states.shape[2]) + # operate on the hidden_states + hidden_states = ( + self.norm_out(hidden_states.float()) + * (1 + scale.repeat_interleave(token_length, dim=1)) + + shift.repeat_interleave(token_length, dim=1) + ).type_as(hidden_states) + else: + hidden_states = ( + self.norm_out(hidden_states.float()) * (1 + scale) + shift + ).type_as(hidden_states) + + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def add_discrete_action_parameters(self): + self.condition_embedder.action_embedder = TimestepEmbedding( + in_channels=self.condition_embedder.time_freq_dim, + time_embed_dim=self.condition_embedder.dim, + ) + nn.init.zeros_(self.condition_embedder.action_embedder.linear_2.weight) + + if self.condition_embedder.action_embedder.linear_2.bias is not None: + nn.init.zeros_(self.condition_embedder.action_embedder.linear_2.bias) + + # prope + for block in self.blocks: + block.attn1.to_out_prope = torch.nn.ModuleList( + [ + torch.nn.Linear(block.attn1.inner_dim, self.inner_dim, bias=True), + ] + ) + nn.init.zeros_(block.attn1.to_out_prope[0].weight) + + if block.attn1.to_out_prope[0].bias is not None: + nn.init.zeros_(block.attn1.to_out_prope[0].bias) diff --git a/wan/models/par_vae/context_parallel/plugins.py b/wan/models/par_vae/context_parallel/plugins.py new file mode 100644 index 0000000000000000000000000000000000000000..8b1989e9b344127a39d46243d7d53d82f4ce6a23 --- /dev/null +++ b/wan/models/par_vae/context_parallel/plugins.py @@ -0,0 +1,453 @@ +import torch +import torch.distributed as dist +import math + +import time + + +class ModulePlugin: + def __init__(self, module, module_id, global_state=None): + self.module = module + self.module_id = module_id + self.global_state = global_state + self.enable = True + self.implement_forward() + + @property + def is_log_node(self): + return ( + self.global_state.get("dist_controller").rank == 0 + and self.module_id[1] == 0 + ) + + @property + def t(self): + return self.global_state.get("timestep") + + @property + def p(self): + return self.t / 1000 + + def implement_forward(self): + module = self.module + if not hasattr(module, "old_forward"): + module.old_forward = module.forward + self.new_forward = self.get_new_forward() + + def forward(*args, **kwargs): + self.update_config() # update config + return ( + self.new_forward(*args, **kwargs) + if self.enable + else self.old_forward(*args, **kwargs) + ) + + module.forward = forward + + def set_enable(self, enable=True): + self.enable = enable + + def get_new_forward(self): + raise NotImplementedError + + def update_config(self, config: dict = None): + if config is None: + config = self.global_state.get("plugin_configs", {}).get( + self.module_id[0], {} + ) + for key, value in config.items(): + setattr(self, key, value) + + def _calculate_padding(self, kernel_size, chunk_dim): + """Calculate padding parameters for convolution layers. + + Args: + kernel_size: kernel size tuple or int + chunk_dim: dimension to chunk/pad along + + Returns: + tuple: (padding_left, padding_right, padding_flag) + """ + if isinstance(kernel_size, int): + kernel_size = tuple([kernel_size] * (chunk_dim - 1)) + + # For 3D conv, use width dimension (index 2) + # For 2D conv, calculate based on chunk_dim + if len(kernel_size) == 3: + kernel_width = kernel_size[2] + else: + pad_kernel_dim = chunk_dim - 3 + kernel_width = kernel_size[pad_kernel_dim] + + d = kernel_width - 1 + padding_left = d // 2 + padding_right = d - padding_left + padding_flag = padding_left if d > 0 else 0 + + return padding_left, padding_right, padding_flag + + def _pad_context_bilateral_impl( + self, h, padding_left, padding_right, slice_dim, cat_dim + ): + """Core implementation for bilateral padding with context from adjacent ranks. + + Args: + h: input tensor to pad + padding_left: number of elements to pad on left + padding_right: number of elements to pad on right + slice_dim: dimension to slice from (3 for 3D, 3 for 2D height/width) + cat_dim: dimension to concatenate along + + Returns: + padded tensor with context from adjacent ranks + """ + if padding_left == 0: + return h + + # Extract slices to share with neighbors + if slice_dim == 3: + share_to_left = h[:, :, :, :padding_left].contiguous() + share_to_right = h[:, :, :, -padding_right:].contiguous() + elif slice_dim == 4: + share_to_left = h[:, :, :, :, :padding_left].contiguous() + share_to_right = h[:, :, :, :, -padding_right:].contiguous() + elif slice_dim == 2: + share_to_left = h[:, :, :padding_left].contiguous() + share_to_right = h[:, :, -padding_right:].contiguous() + else: # slice_dim == 1 + share_to_left = h[:, :, :, :padding_left].contiguous() + share_to_right = h[:, :, :, -padding_right:].contiguous() + + rank = dist.get_rank() + world_size = dist.get_world_size() + adj_groups = self.global_state.get("dist_controller").adj_groups + + if rank % 2: + # Odd rank: pad left first, then right + if rank: + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to( + h.device, non_blocking=True + ) + + if rank != world_size - 1: + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=adj_groups[rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to( + h.device, non_blocking=True + ) + else: + # Even rank: pad right first, then left + if rank != world_size - 1: + padding_list = [torch.zeros_like(share_to_right) for _ in range(2)] + dist.all_gather(padding_list, share_to_right, group=adj_groups[rank]) + right_context = padding_list[1].to(h.device, non_blocking=True) + else: + right_context = torch.zeros_like(share_to_right).to( + h.device, non_blocking=True + ) + + if rank: + padding_list = [torch.zeros_like(share_to_left) for _ in range(2)] + dist.all_gather(padding_list, share_to_left, group=adj_groups[rank - 1]) + left_context = padding_list[0].to(h.device, non_blocking=True) + else: + left_context = torch.zeros_like(share_to_left).to( + h.device, non_blocking=True + ) + + h_with_context = torch.cat([left_context, h, right_context], dim=cat_dim) + return h_with_context + + def _pad_context_bilateral(self, h, padding_left, padding_right, chunk_dim): + """Pad tensor with context from adjacent ranks using bilateral communication. + + Args: + h: input tensor to pad + padding_left: number of elements to pad on left + padding_right: number of elements to pad on right + chunk_dim: dimension to pad along + + Returns: + padded tensor with context from adjacent ranks + """ + slice_dim = 3 if chunk_dim == 3 else 4 + return self._pad_context_bilateral_impl( + h, padding_left, padding_right, slice_dim, chunk_dim + ) + + def _pad_context_bilateral_2d(self, h, padding_left, padding_right, chunk_dim): + """Pad 2D tensor with context from adjacent ranks using bilateral communication. + + Similar to _pad_context_bilateral but pads on chunk_dim-1 dimension for 2D convolutions. + + Args: + h: input tensor to pad + padding_left: number of elements to pad on left + padding_right: number of elements to pad on right + chunk_dim: base dimension reference (actual padding on chunk_dim-1) + + Returns: + padded tensor with context from adjacent ranks + """ + slice_dim = 2 if chunk_dim == 3 else 1 + return self._pad_context_bilateral_impl( + h, padding_left, padding_right, slice_dim, chunk_dim - 1 + ) + + +class GroupNormPlugin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + def get_new_forward(self): + module = self.module + + def new_forward(x): + shape = x.shape + N, C, G = shape[0], shape[1], module.num_groups + assert C % G == 0 + + x = x.reshape(N, G, -1) + + mean = x.mean(-1, keepdim=True).to(torch.float32) + dist.all_reduce(mean) + + mean = mean / dist.get_world_size() + + var = ((x - mean.to(x.dtype)) ** 2).mean(-1, keepdim=True).to(torch.float32) + + dist.all_reduce(var) + var = var / dist.get_world_size() + + x = (x - mean.to(x.dtype)) / (var.to(x.dtype) + module.eps).sqrt() + x = x.view(shape) + + new_shape = [1 for _ in shape] + new_shape[1] = -1 + + return x * module.weight.view(new_shape) + module.bias.view(new_shape) + + return new_forward + + +class Conv3DSafeNewPligin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, "kernel_size", (1, 1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size, self.kernel_size) + + self.chunk_dim = global_state.get("chunk_dim") + self.padding_left, self.padding_right, self.padding_flag = ( + self._calculate_padding(self.kernel_size, self.chunk_dim) + ) + + def pad_context(self, h): + return self._pad_context_bilateral( + h, self.padding_left, self.padding_right, self.chunk_dim + ) + + def get_new_forward(self): + module = self.module + + def new_forward(hidden_states, cache_x=None, *args, **kwargs): + if self.padding_flag == 0: + # print(f"padding=0, return old_forward") + return module.old_forward(hidden_states, cache_x, *args, **kwargs) + + hidden_states = self.pad_context(hidden_states) + if cache_x is not None: + cache_x = self.pad_context(cache_x) + + result = module.old_forward(hidden_states, cache_x, *args, **kwargs) + if self.chunk_dim == 3: + result = result[ + :, + :, + :, + self.padding_left : ( + -self.padding_right if self.padding_right > 0 else None + ), + ] + else: + result = result[ + :, + :, + :, + :, + self.padding_left : ( + -self.padding_right if self.padding_right > 0 else None + ), + ] + + return result + + return new_forward + + +class Conv2DSafeNewPligin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, "kernel_size", (1, 1)) + self.stride = getattr(module, "stride", (1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size) + + self.chunk_dim = global_state.get("chunk_dim") + pad_kernel_dim = self.chunk_dim - 3 + kernel_height = self.kernel_size[pad_kernel_dim] + d = kernel_height - 1 + self.padding_left = d // 2 + self.padding_right = d - self.padding_left + self.padding = self.padding_left if d > 0 else 0 + + def pad_context(self, h): + return self._pad_context_bilateral_2d( + h, self.padding_left, self.padding_right, self.chunk_dim + ) + + def get_new_forward(self): + module = self.module + + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + if self.padding == 0: + return module.old_forward(hidden_states) + + hidden_states = self.pad_context(hidden_states) + + if self.chunk_dim == 3: + hidden_states = module.old_forward(hidden_states)[ + :, + :, + self.padding_left : ( + -self.padding_right if self.padding_right > 0 else None + ), + ] + else: + hidden_states = module.old_forward(hidden_states)[ + :, + :, + :, + self.padding_left : ( + -self.padding_right if self.padding_right > 0 else None + ), + ] + return hidden_states + + return new_forward + + +class Conv2DSafeNewPliginStride2(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + super().__init__(module, module_id, global_state) + + self.kernel_size = getattr(module, "kernel_size", (1, 1)) + self.stride = getattr(module, "stride", (1, 1)) + + if isinstance(self.kernel_size, int): + self.kernel_size = (self.kernel_size, self.kernel_size) + + self.chunk_dim = global_state.get("chunk_dim") + assert self.kernel_size[0] == self.kernel_size[1] + + kernel_height = self.kernel_size[0] + d = kernel_height - 1 + self.padding_left = d // 2 + self.padding_right = d - self.padding_left + self.padding = self.padding_left if d > 0 else 0 + + def pad_context(self, h): + if self.padding == 0: + return h + + rank = dist.get_rank() + world_size = dist.get_world_size() + + if self.chunk_dim == 3: + share_to_left = h[:, :, : self.padding_left].contiguous() + else: + share_to_left = h[:, :, :, : self.padding_left].contiguous() + + if rank < world_size - 1: + right_context = torch.zeros_like(share_to_left) + + dist.recv(right_context, src=rank + 1) + if rank > 0: + dist.send(share_to_left, dst=rank - 1) + # torch.cuda.synchronize() + if rank < world_size - 1: + h_with_context = torch.cat([h, right_context], dim=self.chunk_dim - 1) + else: + h_with_context = h + return h_with_context + + def get_new_forward(self): + module = self.module + + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + # assert self.global_state["chunk_dim"] == 3 # only support chunk height now + if self.padding == 0: + return module.old_forward(hidden_states) + + if self.chunk_dim == 3: + hidden_states = hidden_states[:, :, :-1, :] + else: + hidden_states = hidden_states[:, :, :, :-1] + + hidden_states = self.pad_context(hidden_states) + + if self.chunk_dim == 3: + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, 1)) + else: + hidden_states = torch.nn.functional.pad(hidden_states, (0, 1)) + hidden_states = module.old_forward( + hidden_states + ) # [:,:,self.padding_left:-self.padding_right if self.padding_right > 0 else None] + return hidden_states + + return new_forward + + +class WanAttentionPlugin(ModulePlugin): + def __init__(self, module, module_id, global_state=None): + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + self.gather_dim = global_state.get("chunk_dim") + + super().__init__(module, module_id, global_state) + + def get_new_forward(self): + module = self.module + rank = self.rank + world_size = self.world_size + + def new_forward(hidden_states: torch.Tensor) -> torch.Tensor: + gathered_tensors = [ + torch.zeros_like(hidden_states) for _ in range(world_size) + ] + dist.all_gather(gathered_tensors, hidden_states) + + combined_tensor = torch.cat(gathered_tensors, dim=self.gather_dim) + forward_output = module.old_forward(combined_tensor) + + # chunk_sizes = [t.size(self.gather_dim) for t in gathered_tensors] + # start_idx = sum(chunk_sizes[:rank]) + # end_idx = start_idx + chunk_sizes[rank] + # local_output = forward_output[:, :, :, start_idx:end_idx].contiguous() + local_output = forward_output.chunk(chunks=world_size, dim=self.gather_dim)[ + rank + ] + + return local_output + + return new_forward diff --git a/wan/models/par_vae/context_parallel/wrapper_vae.py b/wan/models/par_vae/context_parallel/wrapper_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..d9c725a95e333bd9411a0bccc5e906bfab0c7426 --- /dev/null +++ b/wan/models/par_vae/context_parallel/wrapper_vae.py @@ -0,0 +1,209 @@ +from ..tools import GlobalState, DistController + +from .plugins import ( + torch, + ModulePlugin, + GroupNormPlugin, + Conv3DSafeNewPligin, + Conv2DSafeNewPligin, + WanAttentionPlugin, + Conv2DSafeNewPliginStride2, +) +from diffusers.models.autoencoders.autoencoder_kl_wan import ( + WanCausalConv3d, + WanAttentionBlock, +) + +# from ...modules.vae import CausalConv3d, AttentionBlock +# from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanAttentionBlock + + +class DistWrapper(object): + def __init__( + self, pipe, dist_controller: DistController, config, chunk_dim + ) -> None: + super().__init__() + self.pipe = pipe + self.dist_controller = dist_controller + self.config = config + # Ensure chunk_dim is either 3 (temporal) or 4 (spatial) for proper parallel processing + assert chunk_dim == 3 or chunk_dim == 4 + # Initialize global state to share distributed configuration across plugins + self.global_state = GlobalState( + {"dist_controller": dist_controller, "chunk_dim": chunk_dim} + ) + self.plugin_mount() + + # Configure plugin-specific hyperparameters for attention and convolution layers + plugin_configs = { + "attn": { + "padding": 24, + "top_k": 24, + "top_k_chunk_size": 24, + "attn_scale": 1.0, + "token_num_scale": True, + "dynamic_scale": True, + }, + "conv_3d": { + "padding": 1, + }, + "conv_layer": {}, + } + self.global_state.set("plugin_configs", plugin_configs) + + # torch.compile + # torch._dynamo.config.recompile_limit = 20 + # Compile decoder with torch.compile for improved inference performance + with torch.no_grad(): + # self.pipe.encoder = torch.compile(self.pipe.encoder) + self.pipe.decoder = torch.compile(self.pipe.decoder) + + def plugin_mount(self): + self.plugins = {} + self.group_norm_plugin_mount() + self.conv_3d_plugin_mount() + self.conv_2d_plugin_stride2_mount() ##only for wan vae encoder + self.conv_2d_plugin_mount() + self.wanattention_plugin_mount() + + def wanattention_plugin_mount(self): + # Mount attention plugins to WanAttentionBlock modules for distributed processing + self.plugins["wanattention"] = {} + wanattention_s = [] + for module in self.pipe.encoder.named_modules(): + # if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock': + # print("Encoder attn: ", module[0]) + if ( + "mid_block." in module[0] + and module[1].__class__.__name__ == "WanAttentionBlock" + ): + wanattention_s.append(module[1]) + for module in self.pipe.decoder.named_modules(): + # if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock': + # print("Decoder attn: ", module[0]) + if ( + "mid_block." in module[0] + and module[1].__class__.__name__ == "WanAttentionBlock" + ): + wanattention_s.append(module[1]) + if self.dist_controller.is_master: + print(f"Found {len(wanattention_s)} wanattention_s") + for i, wanattention in enumerate(wanattention_s): + plugin_id = "wanattention", i + self.plugins["wanattention"][plugin_id] = WanAttentionPlugin( + wanattention, plugin_id, self.global_state + ) + + def group_norm_plugin_mount(self): + self.plugins["group_norm"] = {} + group_norms = [] + for module in self.pipe.decoder.named_modules(): + if ("norm_layer" in module[0]) and module[ + 1 + ].__class__.__name__ == "GroupNorm": + group_norms.append(module[1]) + if self.dist_controller.is_master: + print(f"Found {len(group_norms)} group norms") + for i, group_norm in enumerate(group_norms): + plugin_id = "group_norm", i + self.plugins["group_norm"][plugin_id] = GroupNormPlugin( + group_norm, plugin_id, self.global_state + ) + + def conv_3d_plugin_mount(self): + self.plugins["conv_3d"] = {} + conv3d_s = [] + for module in self.pipe.encoder.named_modules(): + # if isinstance(module[1], CausalConv3d): + # print("Encoder conv3d: ", module[0], module[1].kernel_size[1]) + if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1: + # print(f"Found conv3d: {module[1]}") + conv3d_s.append(module[1]) + for module in self.pipe.decoder.named_modules(): + # if isinstance(module[1], CausalConv3d): + # print("Decoder conv3d: ", module[0], module[1].kernel_size[1]) + if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1: + # print(f"Found conv3d: {module[1]}") + conv3d_s.append(module[1]) + if self.dist_controller.is_master: + print(f"Found {len(conv3d_s)} conv3d_s") + for i, conv in enumerate(conv3d_s): + plugin_id = "conv_3d", i + self.plugins["conv_3d"][plugin_id] = Conv3DSafeNewPligin( + conv, plugin_id, self.global_state + ) + + def conv_2d_plugin_stride2_mount(self): + self.plugins["conv_2d_stride2"] = {} + conv2d_stride2_s = [] + for module in self.pipe.encoder.named_modules(): + if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d": + conv2d_stride2_s.append(module[1]) + if self.dist_controller.is_master: + print(f"Found {len(conv2d_stride2_s)} conv2d_stride2_s") + for i, conv in enumerate(conv2d_stride2_s): + plugin_id = "conv_2d_stride2", i + self.plugins["conv_2d_stride2"][plugin_id] = Conv2DSafeNewPliginStride2( + conv, plugin_id, self.global_state + ) + + def conv_2d_plugin_mount(self): + self.plugins["conv_2d"] = {} + conv2d_s = [] + for module in self.pipe.decoder.named_modules(): + if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d": + conv2d_s.append(module[1]) + if self.dist_controller.is_master: + print(f"Found {len(conv2d_s)} conv2d_s") + for i, conv in enumerate(conv2d_s): + plugin_id = "conv_2d", i + self.plugins["conv_2d"][plugin_id] = Conv2DSafeNewPligin( + conv, plugin_id, self.global_state + ) + + def inference( + self, + local_pose_image, + local_latents, + config={}, + pipe_configs={ + "steps": 50, + "guidance_scale": 12, + "fps": 60, + "num_frames": 24 * 1, + "height": 320, + "width": 512, + "export_fps": 12, + "base_path": "./work/output", + "file_name": None, + }, + plugin_configs={ + "attn": { + "padding": 24, + "top_k": 24, + "top_k_chunk_size": 24, + "attn_scale": 1.0, + "token_num_scale": True, + "dynamic_scale": True, + }, + "conv_3d": { + "padding": 1, + }, + "conv_layer": {}, + }, + additional_info={}, + ): + # Run VAE encode-decode inference on local pose images and latents + self.plugin_mount() + # print("self.config seed: ", self.config["seed"]) + + self.global_state.set("plugin_configs", plugin_configs) + self.pipe = self.pipe.to(device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + local_pose_image = local_pose_image.to(device="cuda", dtype=torch.bfloat16) + local_latents = local_latents.to(device="cuda", dtype=torch.bfloat16) + + tmp_latents = self.pipe.encode(local_pose_image).latent_dist.mode() + + latents = self.pipe.decode(local_latents, return_dict=False)[0] + return latents diff --git a/wan/models/par_vae/tools.py b/wan/models/par_vae/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..262ccb61f7284e79847d17bc63050f3326058806 --- /dev/null +++ b/wan/models/par_vae/tools.py @@ -0,0 +1,117 @@ +import json +import numpy as np +import imageio +import os + +import torch +import torch.distributed as dist + + +def export_to_video(video_frames, output_video_path, fps=12): + # Export a sequence of video frames to an MP4 file with specified frame rate + # Ensure all frames are NumPy arrays and determine video dimensions from the first frame + assert all( + isinstance(frame, np.ndarray) for frame in video_frames + ), "All video frames must be NumPy arrays." + # Ensure output_video_path is ending with .mp4 + if not output_video_path.endswith(".mp4"): + output_video_path += ".mp4" + # Create a video file at the specified path and write frames to it + with imageio.get_writer(output_video_path, fps=fps, format="mp4") as writer: + for frame in video_frames: + writer.append_data((frame * 255).astype(np.uint8)) + + +def save_generation(video_frames, configs, base_path, file_name=None): + if not os.path.exists(base_path): + os.makedirs(base_path) + p_config = configs["pipe_configs"] + frames, steps, fps = p_config["num_frames"], p_config["steps"], p_config["fps"] + if not file_name: + index = [int(each.split("_")[0]) for each in os.listdir(base_path)] + max_idex = max(index) if index else 0 + idx_str = str(max_idex + 1).zfill(6) + + key_info = "_".join([str(frames), str(steps), str(fps)]) + file_name = f"{idx_str}_{key_info}" + + with open(f"{base_path}/{file_name}.json", "w") as f: + json.dump(configs, f, indent=4) + + export_to_video( + video_frames, f"{base_path}/{file_name}.mp4", fps=p_config["export_fps"] + ) + + return file_name + + +class GlobalState: + # Shared state container for passing configuration across distributed components + def __init__(self, state={}) -> None: + self.init_state(state) + + def init_state(self, state={}): + self.state = state + + def set(self, key, value): + self.state[key] = value + + def get(self, key, default=None): + return self.state.get(key, default) + + +class DistController(object): + # Controller for distributed training/inference across multiple GPUs + def __init__(self, rank, local_rank, world_size, config=None) -> None: + super().__init__() + self.rank = rank + self.local_rank = local_rank + self.world_size = world_size + self.config = config + # Designate rank 0 as the master process for logging and coordination + self.is_master = rank == 0 + # print("DistController is master: ", self.is_master) + # self.init_dist() + self.init_group() + # self.device = torch.device(f"cuda:{config['devices'][dist.get_rank()]}") + self.device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(self.device) + + def init_dist(self): + # print(f"Rank {self.rank} is running.") + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(self.config.get("master_port") or "29500") + dist.init_process_group("nccl", rank=self.rank, world_size=self.world_size) + + def init_group(self): + # Create communication groups for adjacent GPU pairs to enable ring communication + self.adj_groups = [ + dist.new_group([i, i + 1]) for i in range(self.world_size - 1) + ] + + +def distribute_data_to_gpus(data, dim, rank, local_rank, num_gpus, dtype): + # Partition data along a specified dimension and distribute chunks to multiple GPUs + dim_size = data.shape[dim] + assert ( + dim_size % num_gpus == 0 + ), f"Dim {dim_size} of data is not a multiply of {num_gpus}" + + chunk_size = dim_size // num_gpus + + slices = [slice(None)] * len(data.shape) + slices[dim] = slice(rank * chunk_size, (rank + 1) * chunk_size) + + local_data = data[slices].to(device=f"cuda:{local_rank}", dtype=dtype) + # print(f"Rank {rank} local_data: {local_data.shape}") + return local_data + + +def gather_results(local_result, dim, world_size): + dist.barrier() + + gathered_results = [torch.zeros_like(local_result) for _ in range(world_size)] + dist.all_gather(gathered_results, local_result) + + gathered_result = torch.cat(gathered_results, dim=dim) + return gathered_result diff --git a/wan/models/utils.py b/wan/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..080d2e6fd7d01888a576e3de232245aa45380911 --- /dev/null +++ b/wan/models/utils.py @@ -0,0 +1,327 @@ +from typing import List +import numpy as np +import torch +import math +from einops import rearrange +from scipy.spatial.transform import Rotation as R_scipy +from hyvideo.utils.retrieval_context import calculate_fov_overlap_similarity + +from typing import List, Dict, Tuple +from distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size + + +def shard_latents_dim_across_sp(latents: torch.Tensor, total_dim) -> torch.Tensor: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + # latents = latents[:, :, :num_latent_t] + if sp_world_size > 1: + if total_dim == 5: + latents = rearrange( + latents, "b c (n s) h w -> b c n s h w", n=sp_world_size + ).contiguous() + latents = latents[:, :, rank_in_sp_group, :, :, :] + elif total_dim == 4: + latents = rearrange( + latents, "b (n s) h w -> b n s h w", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :, :, :] + elif total_dim == 3: + latents = rearrange( + latents, "b (n s) c -> b n s c", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :, :] + elif total_dim == 2: + latents = rearrange( + latents, "b (n s)-> b n s", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :] + elif total_dim == 1: + latents = rearrange(latents, "(n s)-> n s", n=sp_world_size).contiguous() + latents = latents[rank_in_sp_group, :] + else: + raise NotImplementedError(f"total_dim {total_dim} not supported") + return latents + + +def select_mem_frames_wan( + w2c_list: List[np.ndarray], + current_frame_idx: int, + memory_frames: int, + temporal_context_size: int, + pred_latent_size: int, + pos_weight: float = 1.0, + ang_weight: float = 1.0, + device=None, + points_local=None, +) -> List[int]: + """ + 为给定帧选择记忆帧和上下文帧，基于复杂的四帧片段距离计算。 + + 参数: + w2c_list (List[np.ndarray]): 包含所有N个4x4外参矩阵的列表。 + current_frame_idx (int): 当前要处理的帧的索引。 + memory_frames (int): 需要选择的记忆帧总数。 + context_size (int): 需要选择的上下文帧总数。 + pos_weight (float): 空间距离的权重。 + ang_weight (float): 角度距离的权重。 + + 返回: + List[int]: 包含选定记忆帧和上下文帧索引的列表。 + """ + num_total_frames = len(w2c_list) + # 检查当前帧是否能构成一个完整的4帧片段 + if current_frame_idx >= num_total_frames or current_frame_idx < 3: + raise ValueError("当前帧索引必须在 w2c_list 的有效范围内，且至少为3。") + + # 1. 选择上下文帧 (Context Frames) + start_context_idx = max(0, current_frame_idx - temporal_context_size) + context_frames_indices = list(range(start_context_idx, current_frame_idx)) + + # 2. 计算记忆帧 (Memory Frames) 的候选池 + candidate_distances = [] + query_clip_indices = list( + range( + current_frame_idx, + ( + current_frame_idx + pred_latent_size + if current_frame_idx + pred_latent_size <= num_total_frames + else num_total_frames + ), + ) + ) + + historical_clip_indices = list( + range(0, current_frame_idx - temporal_context_size, 4) + ) + + memory_frames_indices = [] # add the first latent frame as context + memory_frames = memory_frames - temporal_context_size + + for hist_idx in historical_clip_indices: + total_dist = 0 + hist_w2c_1 = w2c_list[hist_idx] + hist_w2c_2 = w2c_list[hist_idx + 2] + for query_idx in query_clip_indices: + dist_1_for_query_idx = 1.0 - calculate_fov_overlap_similarity( + w2c_list[query_idx], + hist_w2c_1, + fov_h_deg=60.0, + fov_v_deg=35.0, + device=device, + points_local=points_local, + ) + dist_2_for_query_idx = 1.0 - calculate_fov_overlap_similarity( + w2c_list[query_idx], + hist_w2c_2, + fov_h_deg=60.0, + fov_v_deg=35.0, + device=device, + points_local=points_local, + ) + dist_for_query_idx = (dist_1_for_query_idx + dist_2_for_query_idx) / 2.0 + total_dist += dist_for_query_idx + + final_clip_distance = total_dist / len(query_clip_indices) + candidate_distances.append((hist_idx, final_clip_distance)) + + candidate_distances.sort(key=lambda x: x[1]) + + # 遍历排序后的候选片段，直到收集到足够的记忆帧 + for start_idx, _ in candidate_distances: + if start_idx not in memory_frames_indices: + memory_frames_indices.extend(range(start_idx, start_idx + 4)) + + # 检查是否已达到 memory_size 的要求 + if len(memory_frames_indices) >= memory_frames: + break + + # 4. 组合并去重，以确保没有重复的帧 + selected_frames_set = set(context_frames_indices) + selected_frames_set.update(memory_frames_indices) + + final_selected_frames = sorted(list(selected_frames_set)) + assert len(final_selected_frames) == memory_frames + temporal_context_size + return final_selected_frames + + +# --- 核心函数重构 --- +# def select_memory_frames(w2c_list: List[np.ndarray], current_frame_idx: int, memory_size: int, +# context_size: int, pos_weight: float = 1.0, ang_weight: float = 1.0) -> List[int]: +# """ +# 为给定帧选择记忆帧和上下文帧，基于复杂的四帧片段距离计算。 + +# 参数: +# w2c_list (List[np.ndarray]): 包含所有N个4x4外参矩阵的列表。 +# current_frame_idx (int): 当前要处理的帧的索引。 +# memory_size (int): 需要选择的记忆帧总数。 +# context_size (int): 需要选择的上下文帧总数。 +# pos_weight (float): 空间距离的权重。 +# ang_weight (float): 角度距离的权重。 + +# 返回: +# List[int]: 包含选定记忆帧和上下文帧索引的列表。 +# """ +# num_total_frames = len(w2c_list) +# # 检查当前帧是否能构成一个完整的4帧片段 +# if current_frame_idx >= num_total_frames or current_frame_idx < 3: +# raise ValueError("当前帧索引必须在 w2c_list 的有效范围内，且至少为3。") + +# # 1. 选择上下文帧 (Context Frames) +# start_context_idx = max(0, current_frame_idx - context_size) +# context_frames_indices = list(range(start_context_idx, current_frame_idx)) + +# # 2. 计算记忆帧 (Memory Frames) 的候选池 +# candidate_distances = [] +# query_clip_indices = list(range(current_frame_idx, current_frame_idx + 4)) + +# # 遍历所有历史片段，将每个片段作为记忆帧候选 +# # 历史片段的起始索引必须是 4 的倍数，且不能与上下文帧重叠 +# for i in range(0, current_frame_idx, 4): +# historical_clip_indices = list(range(i, i + 4)) + +# # 排除上下文帧，如果历史片段与上下文帧有任何交集，就跳过 +# is_context_clip = any(idx in context_frames_indices for idx in historical_clip_indices) +# if is_context_clip: +# continue + +# avg_distance = calculate_complex_clip_distance(w2c_list, query_clip_indices, historical_clip_indices, +# pos_weight, ang_weight) + +# # 存储 (片段起始帧索引, 平均距离) +# candidate_distances.append((i, avg_distance)) + +# # 按平均距离从小到大排序 +# candidate_distances.sort(key=lambda x: x[1]) + +# # 3. 选取最相似的 `memory_size` 个帧 +# memory_frames_indices = [] + +# # 遍历排序后的候选片段，直到收集到足够的记忆帧 +# for start_idx, _ in candidate_distances: +# # 将整个片段的帧都添加到记忆帧列表中 +# memory_frames_indices.extend(range(start_idx, start_idx + 4)) + +# # 检查是否已达到 memory_size 的要求 +# if len(memory_frames_indices) >= memory_size: +# break + +# # 4. 组合并去重，以确保没有重复的帧 +# selected_frames_set = set(context_frames_indices) +# selected_frames_set.update(memory_frames_indices) + +# final_selected_frames = sorted(list(selected_frames_set)) + +# return final_selected_frames + + +def calculate_complex_clip_distance( + w2c_list: List[np.ndarray], + query_clip_indices: List[int], + historical_clip_indices: List[int], + pos_weight: float = 1.0, + ang_weight: float = 1.0, +) -> float: + """ + 计算查询片段与历史片段之间的复杂姿态距离。 + + 该距离是基于查询片段的第二帧和第四帧与历史片段的每一帧的平均距离。 + """ + if len(query_clip_indices) < 4 or len(historical_clip_indices) < 4: + raise ValueError("片段索引列表必须包含4个元素。") + + # 提取查询片段的第二帧和第四帧的w2c矩阵 + query_2nd_w2c = w2c_list[query_clip_indices[1]] # 第二帧 + query_4th_w2c = w2c_list[query_clip_indices[3]] # 第四帧 + + # 1. 计算查询片段第二帧与历史片段每帧的平均距离 + dists_from_2nd = [] + for hist_idx in historical_clip_indices: + hist_w2c = w2c_list[hist_idx] + dist = calculate_pose_distance_from_w2c( + query_2nd_w2c, hist_w2c, pos_weight, ang_weight + ) + dists_from_2nd.append(dist) + avg_dist_from_2nd = np.mean(dists_from_2nd) + + # 2. 计算查询片段第四帧与历史片段每帧的平均距离 + dists_from_4th = [] + for hist_idx in historical_clip_indices: + hist_w2c = w2c_list[hist_idx] + dist = calculate_pose_distance_from_w2c( + query_4th_w2c, hist_w2c, pos_weight, ang_weight + ) + dists_from_4th.append(dist) + avg_dist_from_4th = np.mean(dists_from_4th) + + # 3. 将两个平均值再次取平均，得到最终的片段距离 + final_clip_distance = (avg_dist_from_2nd + avg_dist_from_4th) / 2.0 + + return final_clip_distance + + +def calculate_pose_distance_from_w2c( + w2c_1: np.ndarray, + w2c_2: np.ndarray, + pos_weight: float = 1.0, + ang_weight: float = 1.0, +) -> float: + """ + 根据两个 4x4 W2C (World-to-Camera) 矩阵计算它们之间的综合姿态距离。 + + 该距离量化了两个相机姿态的相似度，类似其 FOV 的重叠程度。 + + 参数: + w2c_1 (np.ndarray): 第一个相机的 4x4 World-to-Camera 矩阵。 + w2c_2 (np.ndarray): 第二个相机的 4x4 World-to-Camera 矩阵。 + pos_weight (float): 空间距离的权重。 + ang_weight (float): 角度距离的权重。 + + 返回: + float: 两个姿态之间的综合距离。 + """ + + def w2c_to_6d_pose(w2c_matrix: np.ndarray) -> np.ndarray: + """ + 将 4x4 World-to-Camera (W2C) 矩阵转换为 6D 姿态。 + + 6D 姿态元组为 (x, y, z, pitch, yaw, roll)。 + """ + # 提取旋转矩阵 R 和平移向量 t + R_cw = w2c_matrix[:3, :3] + t_cw = w2c_matrix[:3, 3] + + # 计算相机在世界坐标系下的位置 C_world + # C_world = -R_cw.T @ t_cw + C_world = -np.dot(R_cw.T, t_cw) + + # 将旋转矩阵转换为欧拉角 (pitch, yaw, roll) + # 注意: scipy 默认的欧拉角顺序是 ZYX，对应 yaw, pitch, roll + # 为了与常见的 (pitch, yaw, roll) 顺序匹配，我们手动转换 + r = R_scipy.from_matrix(R_cw) + pitch, yaw, roll = r.as_euler("yxz", degrees=True) + + return np.array([C_world[0], C_world[1], C_world[2], pitch, yaw, roll]) + + # 1. 将两个 W2C 矩阵转换为 6D 姿态 + pose1_6d = w2c_to_6d_pose(w2c_1) + pose2_6d = w2c_to_6d_pose(w2c_2) + + # 2. 计算空间距离 (欧几里得距离) + pos1 = pose1_6d[:3] + pos2 = pose2_6d[:3] + spatial_distance = np.linalg.norm(pos1 - pos2) + + # 3. 计算角度距离 (考虑圆周特性) + angles1 = pose1_6d[3:] + angles2 = pose2_6d[3:] + + angle_diff = np.abs(angles1 - angles2) + # 修正角度差，确保是最小的圆周距离 + angular_distance_vector = np.minimum(angle_diff, 360 - angle_diff) + # 使用欧几里得范数作为综合角度距离 + angular_distance = np.linalg.norm(angular_distance_vector) + + # 4. 结合两种距离得到综合姿态距离 + total_distance = pos_weight * spatial_distance + ang_weight * angular_distance + + return total_distance diff --git a/wan/platforms/__init__.py b/wan/platforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c9d5991b67b591589ee3f1f176d50fa3279231a --- /dev/null +++ b/wan/platforms/__init__.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/__init__.py + +import traceback +from typing import TYPE_CHECKING + +from distributed.logger import init_logger + +# imported by other files, do not remove +from platforms.interface import AttentionBackendEnum # noqa: F401 +from platforms.interface import Platform, PlatformEnum +from utils import resolve_obj_by_qualname + +logger = init_logger(__name__) + + +def cuda_platform_plugin() -> str | None: + is_cuda = False + + try: + from utils import import_pynvml + + pynvml = import_pynvml() # type: ignore[no-untyped-call] + pynvml.nvmlInit() + try: + # NOTE: Edge case: fastvideo cpu build on a GPU machine. + # Third-party pynvml can be imported in cpu build, + # we need to check if fastvideo is built with cpu too. + # Otherwise, fastvideo will always activate cuda plugin + # on a GPU machine, even if in a cpu build. + is_cuda = pynvml.nvmlDeviceGetCount() > 0 + finally: + pynvml.nvmlShutdown() + except Exception as e: + if "nvml" not in e.__class__.__name__.lower(): + # If the error is not related to NVML, re-raise it. + raise e + + # CUDA is supported on Jetson, but NVML may not be. + import os + + def cuda_is_jetson() -> bool: + return os.path.isfile("/etc/nv_tegra_release") or os.path.exists( + "/sys/class/tegra-firmware" + ) + + if cuda_is_jetson(): + is_cuda = True + if is_cuda: + logger.info("CUDA is available") + + return "platforms.cuda.CudaPlatform" if is_cuda else None + + +def mps_platform_plugin() -> str | None: + """Detect if MPS (Metal Performance Shaders) is available on macOS.""" + is_mps = False + + try: + import torch + + if torch.backends.mps.is_available(): + is_mps = True + logger.info("MPS (Metal Performance Shaders) is available") + except Exception as e: + logger.info("MPS detection failed: %s", e) + + return "platforms.mps.MpsPlatform" if is_mps else None + + +def cpu_platform_plugin() -> str | None: + """Detect if CPU platform should be used.""" + # CPU is always available as a fallback + return "platforms.cpu.CpuPlatform" + + +def rocm_platform_plugin() -> str | None: + is_rocm = False + + try: + import amdsmi + + amdsmi.amdsmi_init() + try: + if len(amdsmi.amdsmi_get_processor_handles()) > 0: + is_rocm = True + logger.info("ROCm platform is available") + finally: + amdsmi.amdsmi_shut_down() + except Exception as e: + logger.info("ROCm platform is unavailable: %s", e) + + return "platforms.rocm.RocmPlatform" if is_rocm else None + + +builtin_platform_plugins = { + "cuda": cuda_platform_plugin, + "rocm": rocm_platform_plugin, + "mps": mps_platform_plugin, + "cpu": cpu_platform_plugin, +} + + +def resolve_qualname() -> str: + # TODO(will): if we need to support other platforms, we should consider if + # vLLM's plugin architecture is suitable for our needs. + + # Try MPS first on macOS + platform_cls_qualname = mps_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to ROCm + platform_cls_qualname = rocm_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CUDA + platform_cls_qualname = cuda_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CPU as last resort + platform_cls_qualname = cpu_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + raise RuntimeError("No platform plugin found. Please check your " "installation.") + + +_current_platform = None +_init_trace: str = "" + +if TYPE_CHECKING: + current_platform: Platform + + +def __getattr__(name: str): + if name == "current_platform": + # lazy init current_platform. + # 1. out-of-tree platform plugins need `from fastvideo.platforms import + # Platform` so that they can inherit `Platform` class. Therefore, + # we cannot resolve `current_platform` during the import of + # `fastvideo.platforms`. + # 2. when users use out-of-tree platform plugins, they might run + # `import fastvideo`, some fastvideo internal code might access + # `current_platform` during the import, and we need to make sure + # `current_platform` is only resolved after the plugins are loaded + # (we have tests for this, if any developer violate this, they will + # see the test failures). + global _current_platform + if _current_platform is None: + platform_cls_qualname = resolve_qualname() + _current_platform = resolve_obj_by_qualname(platform_cls_qualname)() + global _init_trace + _init_trace = "".join(traceback.format_stack()) + return _current_platform + elif name in globals(): + return globals()[name] + else: + raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") + + +__all__ = ["Platform", "PlatformEnum", "current_platform", "_init_trace"] diff --git a/wan/platforms/cpu.py b/wan/platforms/cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..f778180074344f13e9a5f15807f771fbbca36969 --- /dev/null +++ b/wan/platforms/cpu.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py + +import platform + +import torch + +from fastvideo.platforms.interface import CpuArchEnum, Platform, PlatformEnum + + +class CpuPlatform(Platform): + _enum = PlatformEnum.CPU + device_name = "CPU" + device_type = "cpu" + dispatch_key = "CPU" + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture.""" + machine = platform.machine().lower() + if machine in ("x86_64", "amd64", "i386", "i686"): + return CpuArchEnum.X86 + elif machine in ("arm64", "aarch64"): + return CpuArchEnum.ARM + else: + return CpuArchEnum.UNSPECIFIED + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return platform.processor() + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + return platform.machine() + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + # This is a rough estimate for CPU memory + # In practice, you might want to use psutil or similar + return 0 + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + return True + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + # For CPU, we can't easily get memory usage without additional libraries + return 0.0 + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cpu_communicator.CpuCommunicator" diff --git a/wan/platforms/cuda.py b/wan/platforms/cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..36f0e22a5bfc99af626639eff17b5bb08b89ce8f --- /dev/null +++ b/wan/platforms/cuda.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py +"""Code inside this file can safely assume cuda platform, e.g. importing +pynvml. However, it should not initialize cuda context. +""" + +import os +from collections.abc import Callable +from functools import lru_cache, wraps +from typing import TypeVar + +import torch +from typing_extensions import ParamSpec + +import distributed.envs as envs +from distributed.logger import init_logger +from platforms.interface import ( + AttentionBackendEnum, + DeviceCapability, + Platform, + PlatformEnum, +) +from utils import import_pynvml + +logger = init_logger(__name__) + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +pynvml = import_pynvml() # type: ignore[no-untyped-call] + +# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models +# see https://github.com/huggingface/diffusers/issues/9704 for details +torch.backends.cuda.enable_cudnn_sdp(False) + + +def device_id_to_physical_device_id(device_id: int) -> int: + if "CUDA_VISIBLE_DEVICES" in os.environ: + device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",") + if device_ids == [""]: + msg = ( + "CUDA_VISIBLE_DEVICES is set to empty string, which means" + " GPU support is disabled. If you are using ray, please unset" + " the environment variable `CUDA_VISIBLE_DEVICES` inside the" + " worker/actor. " + "Check https://github.com/vllm-project/vllm/issues/8402 for" + " more information." + ) + raise RuntimeError(msg) + physical_device_id = device_ids[device_id] + return int(physical_device_id) + else: + return device_id + + +def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: + + @wraps(fn) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + pynvml.nvmlInit() + try: + return fn(*args, **kwargs) + finally: + pynvml.nvmlShutdown() + + return wrapper + + +class CudaPlatformBase(Platform): + _enum = PlatformEnum.CUDA + device_name: str = "cuda" + device_type: str = "cuda" + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used" + ) + return False + return True + + @classmethod + def is_full_nvlink(cls, device_ids: list[int]) -> bool: + raise NotImplementedError + + @classmethod + def log_warnings(cls) -> None: + pass + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + # TODO(will): maybe come up with a more general interface for local attention + # if distributed is False, we always try to use Flash attn + + logger.info( + "Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND + ) + if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: + try: + from st_attn import sliding_tile_attention # noqa: F401 + + from fastvideo.attention.backends.sliding_tile_attn import ( # noqa: F401 + SlidingTileAttentionBackend, + ) + + logger.info("Using Sliding Tile Attention backend.") + + return "fastvideo.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Sliding Tile Attention backend: %s", str(e) + ) + raise ImportError( + "Sliding Tile Attention backend is not installed. " + ) from e + elif selected_backend == AttentionBackendEnum.SAGE_ATTN: + try: + from sageattention import sageattn # noqa: F401 + + from fastvideo.attention.backends.sage_attn import ( # noqa: F401 + SageAttentionBackend, + ) + + logger.info("Using Sage Attention backend.") + + return "fastvideo.attention.backends.sage_attn.SageAttentionBackend" + except ImportError as e: + logger.info(e) + logger.info( + "Sage Attention backend is not installed. Fall back to Flash Attention." + ) + elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + try: + from vsa import block_sparse_attn # noqa: F401 + + from fastvideo.attention.backends.video_sparse_attn import ( # noqa: F401 + VideoSparseAttentionBackend, + ) + + logger.info("Using Video Sparse Attention backend.") + + return "fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Video Sparse Attention backend: %s", str(e) + ) + raise ImportError( + "Video Sparse Attention backend is not installed. " + ) from e + elif selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + elif ( + selected_backend == AttentionBackendEnum.FLASH_ATTN + or selected_backend is None + ): + pass + elif selected_backend: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + + target_backend = AttentionBackendEnum.FLASH_ATTN + if not cls.has_device_capability(80): + logger.info( + "Cannot use FlashAttention-2 backend for Volta and Turing " "GPUs." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + elif dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention-2 backend for dtype other than " + "torch.float16 or torch.bfloat16." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + # FlashAttn is valid for the model, checking if the package is + # installed. + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from fastvideo.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend, + ) + + supported_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info( + "Cannot use FlashAttention-2 backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "fastvideo.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + + +# NVML utils +# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, +# all the related functions work on real physical device ids. +# the major benefit of using NVML is that it will not initialize CUDA +class NvmlCudaPlatform(CudaPlatformBase): + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: + try: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) + return DeviceCapability(major=major, minor=minor) + except RuntimeError: + return None + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + try: + return bool(super().has_device_capability(capability, device_id)) + except RuntimeError: + return False + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_name(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + return cls._get_physical_device_name(physical_device_id) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_uuid(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return str(pynvml.nvmlDeviceGetUUID(handle)) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_total_memory(cls, device_id: int = 0) -> int: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + + @classmethod + @with_nvml_context + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + """ + query if the set of gpus are fully connected by nvlink (1 hop) + """ + handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i < j: + try: + p2p_status = pynvml.nvmlDeviceGetP2PStatus( + handle, + peer_handle, + pynvml.NVML_P2P_CAPS_INDEX_NVLINK, + ) + if p2p_status != pynvml.NVML_P2P_STATUS_OK: + return False + except pynvml.NVMLError: + logger.exception( + "NVLink detection failed. This is normal if" + " your machine has no NVLink equipped." + ) + return False + return True + + @classmethod + def _get_physical_device_name(cls, device_id: int = 0) -> str: + handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) + return str(pynvml.nvmlDeviceGetName(handle)) + + @classmethod + @with_nvml_context + def log_warnings(cls) -> None: + device_ids: int = pynvml.nvmlDeviceGetCount() + if device_ids > 1: + device_names = [cls._get_physical_device_name(i) for i in range(device_ids)] + if ( + len(set(device_names)) > 1 + and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID" + ): + logger.warning( + "Detected different devices in the system: %s. Please" + " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to " + "avoid unexpected behavior.", + ", ".join(device_names), + ) + + +class NonNvmlCudaPlatform(CudaPlatformBase): + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + device_props = torch.cuda.get_device_properties(device_id) + return int(device_props.total_memory) + + @classmethod + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + logger.exception( + "NVLink detection not possible, as context support was" + " not found. Assuming no NVLink available." + ) + return False + + +# Autodetect either NVML-enabled or non-NVML platform +# based on whether NVML is available. +nvml_available = False +try: + try: + pynvml.nvmlInit() + nvml_available = True + except Exception: + # On Jetson, NVML is not supported. + nvml_available = False +finally: + if nvml_available: + pynvml.nvmlShutdown() + +CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform + +try: + from sphinx.ext.autodoc.mock import _MockModule + + if not isinstance(pynvml, _MockModule): + CudaPlatform.log_warnings() +except ModuleNotFoundError: + CudaPlatform.log_warnings() diff --git a/wan/platforms/interface.py b/wan/platforms/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..1559985c63e7a2904dab15a68561e98b2fd6cef9 --- /dev/null +++ b/wan/platforms/interface.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py + +import enum +import random +from typing import NamedTuple + +import numpy as np +import torch + +from distributed.logger import init_logger + +logger = init_logger(__name__) + + +class AttentionBackendEnum(enum.Enum): + FLASH_ATTN = enum.auto() + SLIDING_TILE_ATTN = enum.auto() + TORCH_SDPA = enum.auto() + SAGE_ATTN = enum.auto() + VIDEO_SPARSE_ATTN = enum.auto() + NO_ATTENTION = enum.auto() + + +class PlatformEnum(enum.Enum): + CUDA = enum.auto() + ROCM = enum.auto() + TPU = enum.auto() + CPU = enum.auto() + MPS = enum.auto() + OOT = enum.auto() + UNSPECIFIED = enum.auto() + + +class CpuArchEnum(enum.Enum): + X86 = enum.auto() + ARM = enum.auto() + UNSPECIFIED = enum.auto() + + +class DeviceCapability(NamedTuple): + major: int + minor: int + + def as_version_str(self) -> str: + return f"{self.major}.{self.minor}" + + def to_int(self) -> int: + """ + Express device capability as an integer ````. + + It is assumed that the minor version is always a single digit. + """ + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class Platform: + _enum: PlatformEnum + device_name: str + device_type: str + + # available dispatch keys: + # use "CPU" as a fallback for platforms not registered in PyTorch + dispatch_key: str = "CPU" + + # The torch.compile backend for compiling simple and + # standalone functions. The default value is "inductor" to keep + # the same behavior as PyTorch. + # NOTE: for the forward part of the model, vLLM has another separate + # compilation strategy. + simple_compile_backend: str = "inductor" + + supported_quantization: list[str] = [] + + def is_cuda(self) -> bool: + return self._enum == PlatformEnum.CUDA + + def is_rocm(self) -> bool: + return self._enum == PlatformEnum.ROCM + + def is_tpu(self) -> bool: + return self._enum == PlatformEnum.TPU + + def is_cpu(self) -> bool: + return self._enum == PlatformEnum.CPU + + def is_out_of_tree(self) -> bool: + return self._enum == PlatformEnum.OOT + + def is_cuda_alike(self) -> bool: + """Stateless version of :func:`torch.cuda.is_available`.""" + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + + def is_mps(self) -> bool: + return self._enum == PlatformEnum.MPS + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + """Get the attention backend class of a device.""" + return "" + + @classmethod + def get_device_capability( + cls, + device_id: int = 0, + ) -> DeviceCapability | None: + """Stateless version of :func:`torch.cuda.get_device_capability`.""" + return None + + @classmethod + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + """ + Test whether this platform is compatible with a device capability. + + The ``capability`` argument can either be: + + - A tuple ``(major, minor)``. + - An integer ````. (See :meth:`DeviceCapability.to_int`) + """ + current_capability = cls.get_device_capability(device_id=device_id) + if current_capability is None: + return False + + if isinstance(capability, tuple): + return current_capability >= capability + + return current_capability.to_int() >= capability + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + """Get the name of a device.""" + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + """Get the uuid of a device, e.g. the PCI bus ID.""" + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + """Get the total memory of a device in bytes.""" + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + """ + Check if the current platform supports async output. + """ + raise NotImplementedError + + @classmethod + def inference_mode(cls): + """A device-specific wrapper of `torch.inference_mode`. + + This wrapper is recommended because some hardware backends such as TPU + do not support `torch.inference_mode`. In such a case, they will fall + back to `torch.no_grad` by overriding this method. + """ + return torch.inference_mode(mode=True) + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """ + Set the seed of each random module. + `torch.manual_seed` will set seed on all devices. + """ + if seed is not None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + @classmethod + def verify_model_arch(cls, model_arch: str) -> None: + """ + Verify whether the current platform supports the specified model + architecture. + + - This will raise an Error or Warning based on the model support on + the current platform. + - By default all models are considered supported. + """ + pass + + @classmethod + def verify_quantization(cls, quant: str) -> None: + """ + Verify whether the quantization is supported by the current platform. + """ + if cls.supported_quantization and quant not in cls.supported_quantization: + raise ValueError( + f"{quant} quantization is currently not supported in " + f"{cls.device_name}." + ) + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + """ + Return the memory usage in bytes. + """ + raise NotImplementedError + + @classmethod + def get_device_communicator_cls(cls) -> str: + """ + Get device specific communicator class for distributed communication. + """ + return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture of the current platform.""" + return CpuArchEnum.UNSPECIFIED + + +class UnspecifiedPlatform(Platform): + _enum = PlatformEnum.UNSPECIFIED + device_type = "" diff --git a/wan/platforms/mps.py b/wan/platforms/mps.py new file mode 100644 index 0000000000000000000000000000000000000000..796030e101b840e8ddeedba6125112e1ac027026 --- /dev/null +++ b/wan/platforms/mps.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from fastvideo.logger import init_logger +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.platforms.interface import DeviceCapability, Platform, PlatformEnum + +logger = init_logger(__name__) + + +class MpsPlatform(Platform): + _enum = PlatformEnum.MPS + device_name: str = "mps" + device_type: str = "mps" + dispatch_key: str = "MPS" + device_control_env_var: str = "MPS_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable MPS " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used" + ) + return False + return True + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + return 0.0 + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible + logger.info("Using Torch SDPA backend for MPS.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + # Use base communicator for MPS + return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """Set the seed for MPS device.""" + if seed is not None: + import random + + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + # MPS doesn't have manual_seed_all like CUDA + # The manual_seed above should be sufficient diff --git a/wan/platforms/rocm.py b/wan/platforms/rocm.py new file mode 100644 index 0000000000000000000000000000000000000000..2f4396187a90fc0a14b76faf77250f15cdf79e2d --- /dev/null +++ b/wan/platforms/rocm.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py +""" +This file is a platform abstraction for ROCm GPUs, +adjusted to match the structure and interface of `cuda.py`. +""" + +import torch + +import fastvideo.envs as envs +from fastvideo.logger import init_logger +from fastvideo.platforms.interface import ( + AttentionBackendEnum, + DeviceCapability, + Platform, + PlatformEnum, +) + +logger = init_logger(__name__) + + +# ROCm uses the same torch.cuda interface +class RocmPlatform(Platform): + _enum = PlatformEnum.ROCM + device_name: str = "rocm" + device_type: str = "cuda" # torch uses 'cuda' backend string + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + return torch.cuda.get_device_properties(device_id).total_memory + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA graph. " + "Since enforce-eager is enabled, async output processor cannot be used" + ) + return False + return True + + @classmethod + def log_warnings(cls) -> None: + pass # ROCm-specific warnings can be added here + + @classmethod + def get_current_memory_usage(cls, device: torch.device | None = None) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + logger.info( + "Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND + ) + + if selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + + elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None): + pass + + elif selected_backend in ( + AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.SAGE_ATTN, + ): + raise ValueError( + f"{selected_backend.name} is not supported on {cls.device_name}." + ) + elif selected_backend: + raise ValueError( + f"Invalid attention backend for {cls.device_name}: {selected_backend}" + ) + + target_backend = AttentionBackendEnum.FLASH_ATTN + if dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention backend for dtype other than " + "torch.float16 or torch.bfloat16." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from fastvideo.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend, + ) + + supported_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info( + "Cannot use FlashAttention backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "fastvideo.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too diff --git a/wan/third_party/__init__.py b/wan/third_party/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/wan/third_party/pynvml.py b/wan/third_party/pynvml.py new file mode 100644 index 0000000000000000000000000000000000000000..41ddb4ed7095c73f61558f26fd049e10447848ac --- /dev/null +++ b/wan/third_party/pynvml.py @@ -0,0 +1,7225 @@ +# SPDX-License-Identifier: Apache-2.0 +# copied from https://pypi.org/project/nvidia-ml-py +# version 12.570.86 + +##### +# Copyright (c) 2011-2023, NVIDIA Corporation. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA Corporation nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +##### + +## +# Python bindings for the NVML library +## +from ctypes import * +from ctypes.util import find_library +from functools import wraps +import sys +import os +import threading +import string + +## C Type mappings ## +## Enums +_nvmlEnableState_t = c_uint +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 + +_nvmlBrandType_t = c_uint +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = ( + 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +) +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 17 + +_nvmlTemperatureThresholds_t = c_uint +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 + +_nvmlTemperatureSensors_t = c_uint +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 + + +_nvmlComputeMode_t = c_uint +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 +NVML_COMPUTEMODE_COUNT = 4 + +_nvmlMemoryLocation_t = c_uint +NVML_MEMORY_LOCATION_L1_CACHE = 0 +NVML_MEMORY_LOCATION_L2_CACHE = 1 +NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 +NVML_MEMORY_LOCATION_DRAM = 2 +NVML_MEMORY_LOCATION_REGISTER_FILE = 3 +NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 +NVML_MEMORY_LOCATION_TEXTURE_SHM = 5 +NVML_MEMORY_LOCATION_CBU = 6 +NVML_MEMORY_LOCATION_SRAM = 7 +NVML_MEMORY_LOCATION_COUNT = 8 + +NVML_NVLINK_MAX_LINKS = 18 + +# For backwards compatibility, maintain the incorrectly-named "LANES" define +NVML_NVLINK_MAX_LANES = NVML_NVLINK_MAX_LINKS + +_nvmlNvLinkErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_REPLAY = 0 +NVML_NVLINK_ERROR_DL_RECOVERY = 1 +NVML_NVLINK_ERROR_DL_CRC_FLIT = 2 +NVML_NVLINK_ERROR_DL_CRC_DATA = 3 +NVML_NVLINK_ERROR_DL_ECC_DATA = 4 +NVML_NVLINK_ERROR_COUNT = 5 + +_nvmlNvLinkEccLaneErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_ECC_LANE0 = 0 +NVML_NVLINK_ERROR_DL_ECC_LANE1 = 1 +NVML_NVLINK_ERROR_DL_ECC_LANE2 = 2 +NVML_NVLINK_ERROR_DL_ECC_LANE3 = 3 +NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 + +_nvmlNvLinkCapability_t = c_uint +NVML_NVLINK_CAP_P2P_SUPPORTED = 0 +NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 + +_nvmlNvLinkUtilizationCountPktTypes_t = c_uint +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF + +_nvmlNvLinkUtilizationCountUnits_t = c_uint +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 + +_nvmlNvLinkDeviceType_t = c_uint +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF + +# These are deprecated, instead use _nvmlMemoryErrorType_t +_nvmlEccBitType_t = c_uint +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 +NVML_ECC_ERROR_TYPE_COUNT = 2 + +_nvmlEccCounterType_t = c_uint +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 +NVML_ECC_COUNTER_TYPE_COUNT = 2 + +_nvmlMemoryErrorType_t = c_uint +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 + +_nvmlClockType_t = c_uint +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 + +_nvmlClockId_t = c_uint +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 +NVML_CLOCK_ID_COUNT = 4 + +_nvmlDriverModel_t = c_uint +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 + +NVML_MAX_GPU_PERF_PSTATES = 16 + +_nvmlPstates_t = c_uint +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 + +_nvmlInforomObject_t = c_uint +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 + +_nvmlReturn_t = c_uint +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_UNKNOWN = 999 + +_nvmlFanState_t = c_uint +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 + +_nvmlFanControlPolicy_t = c_uint +NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 +NVML_FAN_POLICY_MANUAL = 1 + +_nvmlLedColor_t = c_uint +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 + +_nvmlGpuOperationMode_t = c_uint +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 + +_nvmlPageRetirementCause_t = c_uint +NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 + +_nvmlRestrictedAPI_t = c_uint +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 + +_nvmlBridgeChipType_t = c_uint +NVML_BRIDGE_CHIP_PLX = 0 +NVML_BRIDGE_CHIP_BRO4 = 1 +NVML_MAX_PHYSICAL_BRIDGE = 128 + +_nvmlValueType_t = c_uint +NVML_VALUE_TYPE_DOUBLE = 0 +NVML_VALUE_TYPE_UNSIGNED_INT = 1 +NVML_VALUE_TYPE_UNSIGNED_LONG = 2 +NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 +NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4 +NVML_VALUE_TYPE_SIGNED_INT = 5 +NVML_VALUE_TYPE_UNSIGNED_SHORT = 6 +NVML_VALUE_TYPE_COUNT = 7 + +_nvmlNvlinkVersion_t = c_uint +NVML_NVLINK_VERSION_INVALID = 0 +NVML_NVLINK_VERSION_1_0 = 1 +NVML_NVLINK_VERSION_2_0 = 2 +NVML_NVLINK_VERSION_2_2 = 3 +NVML_NVLINK_VERSION_3_0 = 4 +NVML_NVLINK_VERSION_3_1 = 5 +NVML_NVLINK_VERSION_4_0 = 6 +NVML_NVLINK_VERSION_5_0 = 7 + +_nvmlPerfPolicyType_t = c_uint +NVML_PERF_POLICY_POWER = 0 +NVML_PERF_POLICY_THERMAL = 1 +NVML_PERF_POLICY_SYNC_BOOST = 2 +NVML_PERF_POLICY_BOARD_LIMIT = 3 +NVML_PERF_POLICY_LOW_UTILIZATION = 4 +NVML_PERF_POLICY_RELIABILITY = 5 +NVML_PERF_POLICY_TOTAL_APP_CLOCKS = 10 +NVML_PERF_POLICY_TOTAL_BASE_CLOCKS = 11 +NVML_PERF_POLICY_COUNT = 12 + +_nvmlEncoderQueryType_t = c_uint +NVML_ENCODER_QUERY_H264 = 0 +NVML_ENCODER_QUERY_HEVC = 1 +NVML_ENCODER_QUERY_AV1 = 2 +NVML_ENCODER_QUERY_UNKNOWN = 255 + +_nvmlFBCSessionType_t = c_uint +NVML_FBC_SESSION_TYPE_UNKNOWN = 0 +NVML_FBC_SESSION_TYPE_TOSYS = 1 +NVML_FBC_SESSION_TYPE_CUDA = 2 +NVML_FBC_SESSION_TYPE_VID = 3 +NVML_FBC_SESSION_TYPE_HWENC = 4 + +_nvmlDetachGpuState_t = c_uint +NVML_DETACH_GPU_KEEP = 0 +NVML_DETACH_GPU_REMOVE = 1 + +_nvmlPcieLinkState_t = c_uint +NVML_PCIE_LINK_KEEP = 0 +NVML_PCIE_LINK_SHUT_DOWN = 1 + +_nvmlSamplingType_t = c_uint +NVML_TOTAL_POWER_SAMPLES = 0 +NVML_GPU_UTILIZATION_SAMPLES = 1 +NVML_MEMORY_UTILIZATION_SAMPLES = 2 +NVML_ENC_UTILIZATION_SAMPLES = 3 +NVML_DEC_UTILIZATION_SAMPLES = 4 +NVML_PROCESSOR_CLK_SAMPLES = 5 +NVML_MEMORY_CLK_SAMPLES = 6 +NVML_MODULE_POWER_SAMPLES = 7 +NVML_JPG_UTILIZATION_SAMPLES = 8 +NVML_OFA_UTILIZATION_SAMPLES = 9 +NVML_SAMPLINGTYPE_COUNT = 10 + +_nvmlPcieUtilCounter_t = c_uint +NVML_PCIE_UTIL_TX_BYTES = 0 +NVML_PCIE_UTIL_RX_BYTES = 1 +NVML_PCIE_UTIL_COUNT = 2 + +_nvmlGpuTopologyLevel_t = c_uint +NVML_TOPOLOGY_INTERNAL = 0 +NVML_TOPOLOGY_SINGLE = 10 +NVML_TOPOLOGY_MULTIPLE = 20 +NVML_TOPOLOGY_HOSTBRIDGE = 30 +NVML_TOPOLOGY_NODE = 40 +NVML_TOPOLOGY_CPU = NVML_TOPOLOGY_NODE +NVML_TOPOLOGY_SYSTEM = 50 + +_nvmlGpuP2PCapsIndex_t = c_uint +NVML_P2P_CAPS_INDEX_READ = (0,) +NVML_P2P_CAPS_INDEX_WRITE = 1 +NVML_P2P_CAPS_INDEX_NVLINK = 2 +NVML_P2P_CAPS_INDEX_ATOMICS = 3 +# +# NVML_P2P_CAPS_INDEX_PROP is deprecated. +# Use NVML_P2P_CAPS_INDEX_PCI instead. +# +NVML_P2P_CAPS_INDEX_PROP = 4 +NVML_P2P_CAPS_INDEX_PCI = 4 +NVML_P2P_CAPS_INDEX_UNKNOWN = 5 + +_nvmlGpuP2PStatus_t = c_uint +NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED +NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4 +NVML_P2P_STATUS_NOT_SUPPORTED = 5 +NVML_P2P_STATUS_UNKNOWN = 6 + +_nvmlDeviceArchitecture_t = c_uint +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_T23X = 11 +NVML_DEVICE_ARCH_UNKNOWN = 0xFFFFFFFF + +# PCI bus Types +_nvmlBusType_t = c_uint +NVML_BUS_TYPE_UNKNOWN = 0 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 + +_nvmlPowerSource_t = c_uint +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 + +_nvmlAdaptiveClockInfoStatus_t = c_uint +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000 +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 + +_nvmlClockLimitId_t = c_uint +NVML_CLOCK_LIMIT_ID_RANGE_START = 0xFFFFFF00 +NVML_CLOCK_LIMIT_ID_TDP = 0xFFFFFF01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xFFFFFF02 + +_nvmlPcieLinkMaxSpeed_t = c_uint +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 +NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 +NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 + +_nvmlPcieAtomicsCapability_t = c_uint +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 + +_nvmlAffinityScope_t = c_uint +NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_SOCKET = 1 + +_nvmlDeviceGpuRecoveryAction_t = c_uint +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 + +# C preprocessor defined values +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 + +NVML_MAX_GPC_COUNT = 32 + +# buffer size +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = ( + 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +) +NVML_PERF_MODES_BUFFER_SIZE = 2048 + +# Format strings +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" + +NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) +NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) + +""" + Field Identifiers. + + All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. +""" +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive + +# ECC Count Totals +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors +# Individual ECC locations +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_DEV = ( + 22 # Device memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_DEV = ( + 23 # Device memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_REG = ( + 24 # Register File single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_REG = ( + 25 # Register File double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_TEX = ( + 26 # Texture memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_TEX = ( + 27 # Texture memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors + +# Page Retirement +NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors +NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors +NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = ( + 32 # NVLink flow control CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = ( + 33 # NVLink flow control CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = ( + 34 # NVLink flow control CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = ( + 35 # NVLink flow control CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = ( + 36 # NVLink flow control CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = ( + 37 # NVLink flow control CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = ( + 38 # NVLink flow control CRC Error Counter total for all Lanes +) + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = ( + 39 # NVLink data CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = ( + 40 # NVLink data CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = ( + 41 # NVLink data CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = ( + 42 # NVLink data CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = ( + 43 # NVLink data CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = ( + 44 # NVLink data CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = ( + 45 # NvLink data CRC Error Counter total for all Lanes +) + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = ( + 52 # NVLink Replay Error Counter total for all Lanes +) + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = ( + 53 # NVLink Recovery Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = ( + 54 # NVLink Recovery Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = ( + 55 # NVLink Recovery Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = ( + 56 # NVLink Recovery Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = ( + 57 # NVLink Recovery Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = ( + 58 # NVLink Recovery Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = ( + 59 # NVLink Recovery Error Counter total for all Lanes +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = ( + 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = ( + 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = ( + 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = ( + 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = ( + 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = ( + 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = ( + 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = ( + 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = ( + 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = ( + 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = ( + 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = ( + 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = ( + 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = ( + 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes +) + +# Perf Policy Counters +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = ( + 78 # Perf Policy Counter for Low GPU Utilization Policy +) +NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = ( + 80 # Perf Policy Counter for Total App Clock Policy +) +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = ( + 81 # Perf Policy Counter for Total Base Clocks Policy +) + +# Memory temperatures +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device + +# Energy Counter +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = ( + 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded +) + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 + +# NVLink Link Count +NVML_FI_DEV_NVLINK_LINK_COUNT = 91 + +# Page Retirement pending fields +NVML_FI_DEV_RETIRED_PENDING_SBE = 92 +NVML_FI_DEV_RETIRED_PENDING_DBE = 93 + +# PCIe replay and replay rollover counters +NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 +NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = ( + 96 # NVLink flow control CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = ( + 97 # NVLink flow control CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = ( + 98 # NVLink flow control CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = ( + 99 # NVLink flow control CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = ( + 100 # NVLink flow control CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = ( + 101 # NVLink flow control CRC Error Counter for Lane 11 +) + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = ( + 102 # NVLink data CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = ( + 103 # NVLink data CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = ( + 104 # NVLink data CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = ( + 105 # NVLink data CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = ( + 106 # NVLink data CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = ( + 107 # NVLink data CRC Error Counter for Lane 11 +) + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = ( + 112 # NVLink Replay Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = ( + 113 # NVLink Replay Error Counter for Lane 11 +) + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = ( + 114 # NVLink Recovery Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = ( + 115 # NVLink Recovery Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = ( + 116 # NVLink Recovery Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = ( + 117 # NVLink Recovery Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = ( + 118 # NVLink Recovery Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = ( + 119 # NVLink Recovery Error Counter for Lane 11 +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = ( + 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = ( + 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = ( + 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = ( + 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = ( + 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = ( + 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = ( + 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = ( + 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = ( + 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = ( + 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = ( + 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = ( + 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 +) + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 + +# NVLink Throughput Counters +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB + +# Row Remapper +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 + +# Remote device NVLink ID +NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 + +# Number of NVLinks connected to NVSwitch +NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 + +# NvLink ECC Data Error Counters +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = ( + 148 # < NVLink data ECC Error Counter for Link 0 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = ( + 149 # < NVLink data ECC Error Counter for Link 1 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = ( + 150 # < NVLink data ECC Error Counter for Link 2 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = ( + 151 # < NVLink data ECC Error Counter for Link 3 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = ( + 152 # < NVLink data ECC Error Counter for Link 4 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = ( + 153 # < NVLink data ECC Error Counter for Link 5 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = ( + 154 # < NVLink data ECC Error Counter for Link 6 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = ( + 155 # < NVLink data ECC Error Counter for Link 7 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = ( + 156 # < NVLink data ECC Error Counter for Link 8 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = ( + 157 # < NVLink data ECC Error Counter for Link 9 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = ( + 158 # < NVLink data ECC Error Counter for Link 10 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = ( + 159 # < NVLink data ECC Error Counter for Link 11 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = ( + 160 # < NvLink data ECC Error Counter total for all Links +) + +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 + +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 + +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 + +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 + +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 + +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 + +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 + +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 + +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 + +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 + +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 + +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 + +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = ( + 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +) +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 + +NVML_FI_DEV_RESET_STATUS = ( + 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = ( + 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 + +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = ( + 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = ( + 258 # HW Circuitry % lifetime remaining +) +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = ( + 261 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = ( + 262 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = ( + 263 # Ramp down hysteresis value in ms for a given profile +) +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = ( + 265 # % TMP floor for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = ( + 266 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = ( + 267 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = ( + 268 # Ramp down hysteresis value in ms for a given profile +) + +NVML_FI_MAX = 269 # One greater than the largest field ID defined above + +# NVML_FI_DEV_NVLINK_GET_STATE state enums +NVML_NVLINK_STATE_INACTIVE = 0x0 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 + +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = ( + 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = ( + 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) + +## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = ( + 1 # Device is associated with GPU-Passthorugh +) +NVML_GPU_VIRTUALIZATION_MODE_VGPU = ( + 2 # Device is associated with vGPU inside virtual machine. +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = ( + 3 # Device is associated with VGX hypervisor in vGPU mode +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = ( + 4 # Device is associated with VGX hypervisor in vSGA mode +) + +## Lib loading ## +nvmlLib = None +libLoadLock = threading.Lock() +_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown + +## vGPU Management +_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuInstance_t = c_uint + +_nvmlVgpuVmIdType_t = c_uint +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 + +_nvmlGridLicenseFeatureCode_t = c_uint +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = ( + 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +) +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 + +_nvmlGridLicenseExpiryStatus_t = c_uint8 +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = (0,) # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = (1,) # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = (2,) # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = (3,) # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = (4,) # Permanent expiry + +_nvmlVgpuCapability_t = c_uint +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = ( + 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +) +NVML_VGPU_CAP_EXCLUSIVE_TYPE = ( + 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +) +NVML_VGPU_CAP_EXCLUSIVE_SIZE = ( + 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +) +NVML_VGPU_CAP_COUNT = 5 + +_nvmlVgpuDriverCapability_t = c_uint +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = ( + 0 # Supports mixing of different vGPU profiles within one guest VM +) +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 + +_nvmlDeviceVgpuCapability_t = c_uint +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = ( + 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +) +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = ( + 6 # Set/Get support of mini-quarter vGPU profiles +) +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = ( + 7 # Set/Get support for compute media engine vGPU profiles +) +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = ( + 8 # Query whether the GPU supports FSR and warm update +) +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_COUNT = 10 + +_nvmlVgpuGuestInfoState_t = c_uint +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 + +_nvmlVgpuVmCompatibility_t = c_uint +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 + +_nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 + +_nvmlHostVgpuMode_t = c_uint +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 + +_nvmlConfComputeGpusReadyState_t = c_uint +NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 +NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1 + +_nvmlConfComputeGpuCaps_t = c_uint +NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0 +NVML_CC_SYSTEM_GPUS_CC_CAPABLE = 1 + +_nvmlConfComputeCpuCaps_t = c_uint +NVML_CC_SYSTEM_CPU_CAPS_NONE = 0 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1 +NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV_SNP = 3 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SNP_VTOM = 4 + +_nvmlConfComputeDevToolsMode_t = c_uint +NVML_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0 +NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 + +NVML_CC_SYSTEM_MULTIGPU_NONE = 0 +NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 + +NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 +NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 +NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 + +_nvmlConfComputeCcFeature_t = c_uint +NVML_CC_SYSTEM_FEATURE_DISABLED = 0 +NVML_CC_SYSTEM_FEATURE_ENABLED = 1 + +_nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50 +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 + +# GSP firmware +NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + + +class NVMLLibraryMismatchError(Exception): + pass + + +## Error Checking ## +class NVMLError(Exception): + _valClassMapping = dict() + # List of currently known error codes + _errcode_to_string = { + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + + def __new__(typ, value): + """ + Maps value to a proper subclass of NVMLError. + See _extractNVMLErrorsAsClasses function for more details + """ + if typ == NVMLError: + typ = NVMLError._valClassMapping.get(value, typ) + obj = Exception.__new__(typ) + obj.value = value + return obj + + def __str__(self): + try: + if self.value not in NVMLError._errcode_to_string: + NVMLError._errcode_to_string[self.value] = str( + nvmlErrorString(self.value) + ) + return NVMLError._errcode_to_string[self.value] + except NVMLError: + return "NVML Error with code %d" % self.value + + def __eq__(self, other): + return self.value == other.value + + +def nvmlExceptionClass(nvmlErrorCode): + if nvmlErrorCode not in NVMLError._valClassMapping: + raise ValueError("nvmlErrorCode %s is not valid" % nvmlErrorCode) + return NVMLError._valClassMapping[nvmlErrorCode] + + +def _extractNVMLErrorsAsClasses(): + """ + Generates a hierarchy of classes on top of NVMLError class. + + Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate + exceptions more easily. + + NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. + e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized + """ + this_module = sys.modules[__name__] + nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")] + for err_name in nvmlErrorsNames: + # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized + class_name = "NVMLError_" + string.capwords( + err_name.replace("NVML_ERROR_", ""), "_" + ).replace("_", "") + err_val = getattr(this_module, err_name) + + def gen_new(val): + def new(typ): + obj = NVMLError.__new__(typ, val) + return obj + + return new + + new_error_class = type(class_name, (NVMLError,), {"__new__": gen_new(err_val)}) + new_error_class.__module__ = __name__ + setattr(this_module, class_name, new_error_class) + NVMLError._valClassMapping[err_val] = new_error_class + + +_extractNVMLErrorsAsClasses() + + +def _nvmlCheckReturn(ret): + if ret != NVML_SUCCESS: + raise NVMLError(ret) + return ret + + +## Function access ## +_nvmlGetFunctionPointer_cache = ( + dict() +) # function pointers are cached to prevent unnecessary libLoadLock locking + + +def _nvmlGetFunctionPointer(name): + global nvmlLib + + if name in _nvmlGetFunctionPointer_cache: + return _nvmlGetFunctionPointer_cache[name] + + libLoadLock.acquire() + try: + # ensure library was loaded + if nvmlLib == None: + raise NVMLError(NVML_ERROR_UNINITIALIZED) + try: + _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) + return _nvmlGetFunctionPointer_cache[name] + except AttributeError: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + + +## Alternative object +# Allows the object to be printed +# Allows mismatched types to be assigned +# - like None when the Structure variant requires c_uint +class nvmlFriendlyObject(object): + def __init__(self, dictionary): + for x in dictionary: + setattr(self, x, dictionary[x]) + + def __str__(self): + return self.__dict__.__str__() + + +def nvmlStructToFriendlyObject(struct): + d = {} + for x in struct._fields_: + key = x[0] + value = getattr(struct, key) + # only need to convert from bytes if bytes, no need to check python version. + d[key] = value.decode() if isinstance(value, bytes) else value + obj = nvmlFriendlyObject(d) + return obj + + +# pack the object so it can be passed to the NVML library +def nvmlFriendlyObjectToStruct(obj, model): + for x in model._fields_: + key = x[0] + value = obj.__dict__[key] + # any c_char_p in python3 needs to be bytes, default encoding works fine. + if sys.version_info >= (3,): + setattr(model, key, value.encode()) + else: + setattr(model, key, value) + return model + + +## Unit structures +class struct_c_nvmlUnit_t(Structure): + pass # opaque handle + + +c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + + +class _PrintableStructure(Structure): + """ + Abstract class that produces nicer __str__ output than ctypes.Structure. + e.g. instead of: + >>> print str(obj) + + this class will print + class_name(field_name: formatted_value, field_name: formatted_value) + + _fmt_ dictionary of -> + e.g. class that has _field_ 'hex_value', c_uint could be formatted with + _fmt_ = {"hex_value" : "%08X"} + to produce nicer output. + Default fomratting string for all fields can be set with key "" like: + _fmt_ = {"" : "%d MHz"} # e.g all values are numbers in MHz. + If not set it's assumed to be just "%s" + + Exact format of returned str from this class is subject to change in the future. + """ + + _fmt_ = {} + + def __str__(self): + result = [] + for x in self._fields_: + key = x[0] + value = getattr(self, key) + fmt = "%s" + if key in self._fmt_: + fmt = self._fmt_[key] + elif "" in self._fmt_: + fmt = self._fmt_[""] + result.append(("%s: " + fmt) % (key, value)) + return self.__class__.__name__ + "(" + ", ".join(result) + ")" + + def __getattribute__(self, name): + res = super(_PrintableStructure, self).__getattribute__(name) + # need to convert bytes to unicode for python3 don't need to for python2 + # Python 2 strings are of both str and bytes + # Python 3 strings are not of type bytes + # ctypes should convert everything to the correct values otherwise + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + def __setattr__(self, name, value): + if isinstance(value, str): + # encoding a python2 string returns the same value, since python2 strings are bytes already + # bytes passed in python3 will be ignored. + value = value.encode() + super(_PrintableStructure, self).__setattr__(name, value) + + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), + ] + + +class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): + _fields_ = [("isC2cEnabled", c_uint)] + + +nvmlC2cModeInfo_v1 = 0x1000008 + + +class c_nvmlLedState_t(_PrintableStructure): + _fields_ = [ + ("cause", c_char * 256), + ("color", _nvmlLedColor_t), + ] + + +class c_nvmlPSUInfo_t(_PrintableStructure): + _fields_ = [ + ("state", c_char * 256), + ("current", c_uint), + ("voltage", c_uint), + ("power", c_uint), + ] + + +class c_nvmlUnitFanInfo_t(_PrintableStructure): + _fields_ = [ + ("speed", c_uint), + ("state", _nvmlFanState_t), + ] + + +class c_nvmlUnitFanSpeeds_t(_PrintableStructure): + _fields_ = [("fans", c_nvmlUnitFanInfo_t * 24), ("count", c_uint)] + + +## Device structures +class struct_c_nvmlDevice_t(Structure): + pass # opaque handle + + +c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + + +class nvmlPciInfoExt_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + ("pciSubSystemId", c_uint), + ("baseClass", c_uint), + ("subClass", c_uint), + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + "version": "0x%04X", + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + "baseClass": "0x%01X", + "subClass": "0x%01X", + } + + +nvmlPciInfoExt_v1 = 0x1000040 + + +# Legacy pciInfo used for _v1 and _v2 +class nvmlPciInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + # Added in 2.285 + ("pciSubSystemId", c_uint), + ("reserved0", c_uint), + ("reserved1", c_uint), + ("reserved2", c_uint), + ("reserved3", c_uint), + ] + _fmt_ = { + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + + +class nvmlPciInfo_t(_PrintableStructure): + _fields_ = [ + # Moved to the new busId location below + ("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + # Added in 2.285 + ("pciSubSystemId", c_uint), + # New busId replaced the long deprecated and reserved fields with a + # field of the same size in 9.0 + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + "domain": "0x%08X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + + +class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ] + + +SystemDriverBranchInfo_v1 = 0x1000054 + + +class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): + _fields_ = [("pci", nvmlPciInfo_t), ("uuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE)] + + +class nvmlNvLinkUtilizationControl_t(_PrintableStructure): + _fields_ = [ + ("units", _nvmlNvLinkUtilizationCountUnits_t), + ("pktfilter", _nvmlNvLinkUtilizationCountPktTypes_t), + ] + + +class c_nvmlMemory_t(_PrintableStructure): + _fields_ = [ + ("total", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +class c_nvmlMemory_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("total", c_ulonglong), + ("reserved", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +nvmlMemory_v2 = 0x02000028 + + +class c_nvmlBAR1Memory_t(_PrintableStructure): + _fields_ = [ + ("bar1Total", c_ulonglong), + ("bar1Free", c_ulonglong), + ("bar1Used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +class nvmlClkMonFaultInfo_t(Structure): + _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)] + + +MAX_CLK_DOMAINS = 32 + + +class nvmlClkMonStatus_t(Structure): + _fields_ = [ + ("bGlobalStatus", c_uint), + ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS), + ] + + +# On Windows with the WDDM driver, usedGpuMemory is reported as None +# Code that processes this structure should check for None, I.E. +# +# if (info.usedGpuMemory == None): +# # TODO handle the error +# pass +# else: +# print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) +# endif +# +# See NVML documentation for more information +class c_nvmlProcessInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ] + _fmt_ = {"usedGpuMemory": "%d B"} + + +c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t + +c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t + +_nvmlProcessMode_t = c_uint +NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_GRAPHICS = 1 +NVML_PROCESS_MODE_MPS = 2 + + +class c_nvmlProcessDetail_v1_t(Structure): + _fields_ = [ + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ("usedGpuCcProtectedMemory", c_ulonglong), + ] + + +class c_nvmlProcessDetailList_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("mode", _nvmlProcessMode_t), + ("numProcArrayEntries", c_uint), + ("procArray", POINTER(c_nvmlProcessDetail_v1_t)), + ] + _fmt_ = {"numProcArrayEntries": "%d B"} + + +c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t + +nvmlProcessDetailList_v1 = 0x1000018 + + +class c_nvmlBridgeChipInfo_t(_PrintableStructure): + _fields_ = [ + ("type", _nvmlBridgeChipType_t), + ("fwVersion", c_uint), + ] + + +class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): + _fields_ = [ + ("bridgeCount", c_uint), + ("bridgeChipInfo", c_nvmlBridgeChipInfo_t * 128), + ] + + +class c_nvmlEccErrorCounts_t(_PrintableStructure): + _fields_ = [ + ("l1Cache", c_ulonglong), + ("l2Cache", c_ulonglong), + ("deviceMemory", c_ulonglong), + ("registerFile", c_ulonglong), + ] + + +class c_nvmlUtilization_t(_PrintableStructure): + _fields_ = [ + ("gpu", c_uint), + ("memory", c_uint), + ] + _fmt_ = {"": "%d %%"} + + +# Added in 2.285 +class c_nvmlHwbcEntry_t(_PrintableStructure): + _fields_ = [ + ("hwbcId", c_uint), + ("firmwareVersion", c_char * 32), + ] + + +class c_nvmlValue_t(Union): + _fields_ = [ + ("dVal", c_double), + ("uiVal", c_uint), + ("ulVal", c_ulong), + ("ullVal", c_ulonglong), + ("sllVal", c_longlong), + ("siVal", c_int), + ("usVal", c_ushort), + ] + + +class c_nvmlSample_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("sampleValue", c_nvmlValue_t), + ] + + +class c_nvmlViolationTime_t(_PrintableStructure): + _fields_ = [ + ("referenceTime", c_ulonglong), + ("violationTime", c_ulonglong), + ] + + +class c_nvmlFieldValue_t(_PrintableStructure): + _fields_ = [ + ("fieldId", c_uint32), + ("scopeId", c_uint32), + ("timestamp", c_int64), + ("latencyUsec", c_int64), + ("valueType", _nvmlValueType_t), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_nvmlValue_t), + ] + + +NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 + +nvmlNvlinkSupportedBwModes_v1 = 0x100001C + + +class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("bwModes", c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ("totalBwModes", c_uint8), + ] + + def __init__(self): + super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__( + version=nvmlNvlinkSupportedBwModes_v1 + ) + + +nvmlNvlinkGetBwMode_v1 = 0x100000C + + +class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): + _fields_ = [("version", c_uint), ("bIsBest", c_uint), ("bwMode", c_uint8)] + + def __init__(self): + super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) + + +nvmlNvlinkSetBwMode_v1 = 0x100000C + + +class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): + _fields_ = [("version", c_uint), ("bSetBest", c_uint), ("bwMode", c_uint8)] + + def __init__(self): + super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) + + +class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("mode", c_uint), + ] + + +VgpuHeterogeneousMode_v1 = 0x1000008 + + +class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("placementId", c_uint), + ] + + +VgpuPlacementId_v1 = 0x1000008 + + +class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("count", c_uint), + ("placementSize", c_uint), + ("placementIds", POINTER(c_uint)), + ] + + +VgpuPlacementList_v1 = 0x1000018 + +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + + +class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("placementSize", c_uint), + ("count", c_uint), + ("placementIds", POINTER(c_uint)), + ("mode", c_uint), + ] + + +VgpuPlacementList_v2 = 0x2000020 + + +class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("bar1Size", c_ulonglong), + ] + + +VgpuTypeBar1Info_v1 = 0x1000010 + + +class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("vgpuInstance", _nvmlVgpuInstance_t), + ("timeStamp", c_ulonglong), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), + ] + + +class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), + ("jpgUtil", c_nvmlValue_t), + ("ofaUtil", c_nvmlValue_t), + ] + + +class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("sampleValType", _nvmlValueType_t), + ("vgpuInstanceCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuUtilArray", POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), + ] + + +VgpuInstancesUtilizationInfo_v1 = 0x01000020 + + +class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ] + + +class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), + ] + + +class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("vgpuProcessCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuProcUtilArray", POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), + ] + + +VgpuProcessesUtilizationInfo_v1 = 0x01000018 + + +class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("size", c_ulonglong), + ] + + +VgpuRuntimeState_v1 = 0x1000010 + + +class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), + ] + + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 +NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + + +class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): + _fields_ = [ + ("isLicensed", c_uint8), + ("licenseExpiry", c_nvmlVgpuLicenseExpiry_t), + ("currentState", c_uint), + ] + + +class c_nvmlEncoderSession_t(_PrintableStructure): + _fields_ = [ + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("codecType", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFps", c_uint), + ("encodeLatency", c_uint), + ] + + +class c_nvmlProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("pid", c_uint), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ] + + +class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), + ] + + +class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("processSamplesCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("procUtilArray", POINTER(c_nvmlProcessUtilizationInfo_v1_t)), + ] + + +ProcessesUtilizationInfo_v1 = 0x01000018 + + +class c_nvmlGridLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), + ] + + +class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), + ("licenseExpiry", c_nvmlGridLicenseExpiry_t), + ] + + +class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), + ] + + +class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + + +class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + + +class c_nvmlGridLicensableFeatures_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlMarginTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("marginTemperature", c_int), + ] + + +nvmlMarginTemperature_v1 = 0x1000008 + + +## Event structures +class struct_c_nvmlEventSet_t(Structure): + pass # opaque handle + + +c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) + +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = ( + nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction +) + +## Clock Event Reasons defines +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = ( + nvmlClocksEventReasonNone + | nvmlClocksEventReasonGpuIdle + | nvmlClocksEventReasonApplicationsClocksSetting + | nvmlClocksEventReasonSwPowerCap + | nvmlClocksEventReasonHwSlowdown + | nvmlClocksEventReasonSyncBoost + | nvmlClocksEventReasonSwThermalSlowdown + | nvmlClocksEventReasonHwThermalSlowdown + | nvmlClocksEventReasonHwPowerBrakeSlowdown + | nvmlClocksEventReasonDisplayClockSetting +) + +## Following have been deprecated +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = ( + nvmlClocksThrottleReasonNone + | nvmlClocksThrottleReasonGpuIdle + | nvmlClocksThrottleReasonApplicationsClocksSetting + | nvmlClocksThrottleReasonSwPowerCap + | nvmlClocksThrottleReasonHwSlowdown + | nvmlClocksThrottleReasonSyncBoost + | nvmlClocksThrottleReasonSwThermalSlowdown + | nvmlClocksThrottleReasonHwThermalSlowdown + | nvmlClocksThrottleReasonHwPowerBrakeSlowdown + | nvmlClocksThrottleReasonDisplayClockSetting +) + + +class c_nvmlEventData_t(_PrintableStructure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("eventType", c_ulonglong), + ("eventData", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ] + _fmt_ = {"eventType": "0x%08X"} + + +class c_nvmlAccountingStats_t(_PrintableStructure): + _fields_ = [ + ("gpuUtilization", c_uint), + ("memoryUtilization", c_uint), + ("maxMemoryUsage", c_ulonglong), + ("time", c_ulonglong), + ("startTime", c_ulonglong), + ("isRunning", c_uint), + ("reserved", c_uint * 5), + ] + + +class c_nvmlVgpuVersion_t(Structure): + _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)] + + +class c_nvmlVgpuMetadata_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("guestInfoState", _nvmlVgpuGuestInfoState_t), + ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), + ("guestVgpuVersion", c_uint), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE), + ] + + +class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("pgpuVirtualizationCaps", c_uint), + ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE), + ] + + +class c_nvmlVgpuPgpuCompatibility_t(Structure): + _fields_ = [ + ("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t), + ] + + +## vGPU scheduler policy defines +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 + +## Supported vGPU scheduler policy count +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 + +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 + + +class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): + _fields_ = [ + ("avgFactor", c_uint), + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedData_t(_PrintableStructure): + _fields_ = [ + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedulerParams_t(Union): + _fields_ = [ + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedData_t), + ] + + +class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): + _fields_ = [ + ("timestamp", c_ulonglong), + ("timeRunTotal", c_ulonglong), + ("timeRun", c_ulonglong), + ("swRunlistId", c_uint), + ("targetTimeSlice", c_ulonglong), + ("cumulativePreemptionTime", c_ulonglong), + ] + + +class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): + _fields_ = [ + ("engineId", c_uint), + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), + ("entriesCount", c_uint), + ( + "logEntries", + c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES, + ), + ] + + +class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): + _fields_ = [ + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), + ] + + +class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): + _fields_ = [ + ("avgFactor", c_uint), + ("frequency", c_uint), + ] + + +class c_nvmlVgpuSchedSetData_t(_PrintableStructure): + _fields_ = [ + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedulerSetParams_t(Union): + _fields_ = [ + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedSetDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedSetData_t), + ] + + +class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): + _fields_ = [ + ("schedulerPolicy", c_uint), + ("enableARRMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerSetParams_t), + ] + + +class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): + _fields_ = [ + ("supportedSchedulers", c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), + ("maxTimeslice", c_uint), + ("minTimeslice", c_uint), + ("isArrModeSupported", c_uint), + ("maxFrequencyForARR", c_uint), + ("minFrequencyForARR", c_uint), + ("maxAvgFactorForARR", c_uint), + ("minAvgFactorForARR", c_uint), + ] + + +class c_nvmlFBCStats_t(Structure): + _fields_ = [ + ("sessionsCount", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), + ] + + +class c_nvmlFBCSession_t(_PrintableStructure): + _fields_ = [ + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("displayOrdinal", c_uint), + ("sessionType", c_uint), + ("sessionFlags", c_uint), + ("hMaxResolution", c_uint), + ("vMaxResolution", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), + ] + + +NVML_DEVICE_MIG_DISABLE = 0x0 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD + + +class c_nvmlGpuInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), ("size", c_uint)] + + +class c_nvmlGpuInstanceProfileInfo_t(Structure): + _fields_ = [ + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + + +nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + + +class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] + + def __init__(self): + super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__( + version=nvmlGpuInstanceProfileInfo_v2 + ) + + +class c_nvmlGpuInstanceInfo_t(Structure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t), + ] + + +class struct_c_nvmlGpuInstance_t(Structure): + pass # opaque handle + + +c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) + +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 + +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + + +class c_nvmlComputeInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), ("size", c_uint)] + + +class c_nvmlComputeInstanceProfileInfo_t(Structure): + _fields_ = [ + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ] + + +nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + + +class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] + + def __init__(self): + super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__( + version=nvmlComputeInstanceProfileInfo_v2 + ) + + +class c_nvmlComputeInstanceInfo_t(Structure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("gpuInstance", c_nvmlGpuInstance_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t), + ] + + +NVML_MAX_GPU_UTILIZATIONS = 8 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 + + +class c_nvmlGpuDynamicPstatesUtilization_t(Structure): + _fields_ = [ + ("bIsPresent", c_uint, 1), + ("percentage", c_uint), + ("incThreshold", c_uint), + ("decThreshold", c_uint), + ] + + +class c_nvmlGpuDynamicPstatesInfo_t(Structure): + _fields_ = [ + ("flags", c_uint), + ( + "utilization", + c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS, + ), + ] + + +NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 + +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 + +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + + +class c_nvmlGpuThermalSensor_t(Structure): + _fields_ = [ + ("controller", c_int), + ("defaultMinTemp", c_int), + ("defaultMaxTemp", c_int), + ("currentTemp", c_int), + ("target", c_int), + ] + + +class c_nvmlGpuThermalSettings_t(Structure): + _fields_ = [ + ("count", c_uint), + ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU), + ] + + +_nvmlCoolerControl_t = c_uint +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 + +_nvmlCoolerTarget_t = c_uint +NVML_THERMAL_COOLER_TARGET_NONE = 1 << 0 +NVML_THERMAL_COOLER_TARGET_GPU = 1 << 1 +NVML_THERMAL_COOLER_TARGET_MEMORY = 1 << 2 +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = 1 << 3 +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = ( + NVML_THERMAL_COOLER_TARGET_GPU + | NVML_THERMAL_COOLER_TARGET_MEMORY + | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY +) + + +class c_nvmlCoolerInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("index", c_uint), + ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t), + ] + + +nvmlCoolerInfo_v1 = 0x1000010 + + +def nvmlDeviceGetCoolerInfo(handle): + c_coolerInfo = c_nvmlCoolerInfo_t() + c_coolerInfo.version = nvmlCoolerInfo_v1 + c_coolerInfo.index = 0 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCoolerInfo") + ret = fn(handle, byref(c_coolerInfo)) + _nvmlCheckReturn(ret) + return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + + +class struct_c_nvmlComputeInstance_t(Structure): + pass # opaque handle + + +c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + + +class c_nvmlDeviceAttributes(Structure): + _fields_ = [ + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + + +class c_nvmlRowRemapperHistogramValues(Structure): + _fields_ = [ + ("max", c_uint), + ("high", c_uint), + ("partial", c_uint), + ("low", c_uint), + ("none", c_uint), + ] + + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 +NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 +NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + + +class c_nvmlConfComputeSystemState_t(Structure): + _fields_ = [ + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ] + + +nvmlSystemConfComputeSettings_v1 = 0x1000014 + + +class c_nvmlSystemConfComputeSettings_v1_t(Structure): + _fields_ = [ + ("version", c_uint), + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ("multiGpuMode", c_uint), + ] + + def __init__(self): + super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__( + version=nvmlSystemConfComputeSettings_v1 + ) + + +class c_nvmlConfComputeSystemCaps_t(Structure): + _fields_ = [ + ("cpuCaps", c_uint), + ("gpusCaps", c_uint), + ] + + +class c_nvmlConfComputeMemSizeInfo_t(Structure): + _fields_ = [ + ("protectedMemSizeKib", c_ulonglong), + ("unprotectedMemSizeKib", c_ulonglong), + ] + + +class c_nvmlConfComputeGpuCertificate_t(Structure): + _fields_ = [ + ("certChainSize", c_uint), + ("attestationCertChainSize", c_uint), + ("certChain", c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ("attestationCertChain", c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), + ] + + +class c_nvmlConfComputeGpuAttestationReport_t(Structure): + _fields_ = [ + ("isCecAttestationReportPresent", c_uint), + ("attestationReportSize", c_uint), + ("cecAttestationReportSize", c_uint), + ("nonce", c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), + ("attestationReport", c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), + ("cecAttestationReport", c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), + ] + + +class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): + _fields_ = [ + ("version", c_uint), + ("maxAttackerAdvantage", c_ulong), + ] + + +ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + + +class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): + _fields_ = [ + ("version", c_uint), + ("attackerAdvantage", c_ulong), + ] + + +ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 + + +## string/bytes conversion for ease of use +def convertStrBytes(func): + """ + In python 3, strings are unicode instead of bytes, and need to be converted for ctypes + Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) + Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> + ---- + Returned from function: b'returned string' + Returned to caller: 'returned string' + """ + + @wraps(func) + def wrapper(*args, **kwargs): + # encoding a str returns bytes in python 2 and 3 + args = [arg.encode() if isinstance(arg, str) else arg for arg in args] + res = func(*args, **kwargs) + # In python 2, str and bytes are the same + # In python 3, str is unicode and should be decoded. + # Ctypes handles most conversions, this only effects c_char and char arrays. + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + if sys.version_info >= (3,): + return wrapper + return func + + +def throwOnVersionMismatch(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except NVMLError_FunctionNotFound: + raise NVMLLibraryMismatchError( + "Unversioned function called and the " + "pyNVML version does not match the NVML lib version. " + "Either use matching pyNVML and NVML lib versions or " + "use a versioned function such as " + func.__name__ + "_v2" + ) + + return wrapper + + +## C function wrappers ## +def nvmlInitWithFlags(flags): + _LoadNvmlLibrary() + + # + # Initialize the library + # + fn = _nvmlGetFunctionPointer("nvmlInitWithFlags") + ret = fn(flags) + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + _nvmlLib_refcount += 1 + libLoadLock.release() + return None + + +def nvmlInit(): + nvmlInitWithFlags(0) + return None + + +def _LoadNvmlLibrary(): + """ + Load the library if it isn't loaded already + """ + global nvmlLib + + if nvmlLib == None: + # lock to ensure only one caller loads the library + libLoadLock.acquire() + + try: + # ensure the library still isn't loaded + if nvmlLib == None: + try: + if sys.platform[:3] == "win": + # cdecl calling convention + try: + # Check for nvml.dll in System32 first for DCH drivers + nvmlLib = CDLL( + os.path.join( + os.getenv("WINDIR", "C:/Windows"), + "System32/nvml.dll", + ) + ) + except OSError as ose: + # If nvml.dll is not found in System32, it should be in ProgramFiles + # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll + nvmlLib = CDLL( + os.path.join( + os.getenv("ProgramFiles", "C:/Program Files"), + "NVIDIA Corporation/NVSMI/nvml.dll", + ) + ) + else: + # assume linux + nvmlLib = CDLL("libnvidia-ml.so.1") + except OSError as ose: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + if nvmlLib == None: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + + +def nvmlShutdown(): + # + # Leave the library loaded, but shutdown the interface + # + fn = _nvmlGetFunctionPointer("nvmlShutdown") + ret = fn() + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + if 0 < _nvmlLib_refcount: + _nvmlLib_refcount -= 1 + libLoadLock.release() + return None + + +# Added in 2.285 +@convertStrBytes +def nvmlErrorString(result): + fn = _nvmlGetFunctionPointer("nvmlErrorString") + fn.restype = c_char_p # otherwise return is an int + ret = fn(result) + return ret + + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetNVMLVersion(): + c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +def nvmlSystemGetCudaDriverVersion(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + + +def nvmlSystemGetCudaDriverVersion_v2(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetProcessName(pid): + c_name = create_string_buffer(1024) + fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") + ret = fn(c_uint(pid), c_name, c_uint(1024)) + _nvmlCheckReturn(ret) + return c_name.value + + +@convertStrBytes +def nvmlSystemGetDriverVersion(): + c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 2.285 +def nvmlSystemGetHicVersion(): + c_count = c_uint(0) + hics = None + fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") + + # get the count + ret = fn(byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # If there are no hics + if c_count.value == 0: + return [] + + hic_array = c_nvmlHwbcEntry_t * c_count.value + hics = hic_array() + ret = fn(byref(c_count), hics) + _nvmlCheckReturn(ret) + return hics + + +def nvmlSystemGetDriverBranch(): + c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) + c_branchInfo.version = SystemDriverBranchInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_branchInfo + + +## Unit get functions +def nvmlUnitGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlUnitGetHandleByIndex(index): + c_index = c_uint(index) + unit = c_nvmlUnit_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") + ret = fn(c_index, byref(unit)) + _nvmlCheckReturn(ret) + return unit + + +def nvmlUnitGetUnitInfo(unit): + c_info = c_nvmlUnitInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlUnitGetLedState(unit): + c_state = c_nvmlLedState_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") + ret = fn(unit, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + + +def nvmlUnitGetPsuInfo(unit): + c_info = c_nvmlPSUInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlUnitGetTemperature(unit, type): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") + ret = fn(unit, c_uint(type), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlUnitGetFanSpeedInfo(unit): + c_speeds = c_nvmlUnitFanSpeeds_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") + ret = fn(unit, byref(c_speeds)) + _nvmlCheckReturn(ret) + return c_speeds + + +# added to API +def nvmlUnitGetDeviceCount(unit): + c_count = c_uint(0) + # query the unit to determine device count + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), None) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = NVML_SUCCESS + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlUnitGetDevices(unit): + c_count = c_uint(nvmlUnitGetDeviceCount(unit)) + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return c_devices + + +## Device get functions +def nvmlDeviceGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetHandleByIndex(index): + c_index = c_uint(index) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") + ret = fn(c_index, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleBySerial(serial): + c_serial = c_char_p(serial) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") + ret = fn(c_serial, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleByUUID(uuid): + c_uuid = c_char_p(uuid) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") + ret = fn(c_uuid, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleByPciBusId(pciBusId): + c_busId = c_char_p(pciBusId) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") + ret = fn(c_busId, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetName(handle): + c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") + ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_name.value + + +class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + + +nvmlDevicePerfModes_v1 = 0x1000804 + + +@convertStrBytes +def nvmlDeviceGetPerformanceModes(handle): + perfModes = c_nvmlDevicePerfModes_v1_t() + perfModes.version = nvmlDevicePerfModes_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceModes") + ret = fn(handle, byref(perfModes)) + _nvmlCheckReturn(ret) + return perfModes.str + + +class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + + +nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + + +@convertStrBytes +def nvmlDeviceGetCurrentClockFreqs(handle): + currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() + currentClockFreqs.version = nvmlDeviceCurrentClockFreqs_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClockFreqs") + ret = fn(handle, byref(currentClockFreqs)) + _nvmlCheckReturn(ret) + return currentClockFreqs.str + + +def nvmlDeviceGetBoardId(handle): + c_id = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") + ret = fn(handle, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + + +def nvmlDeviceGetMultiGpuBoard(handle): + c_multiGpu = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") + ret = fn(handle, byref(c_multiGpu)) + _nvmlCheckReturn(ret) + return c_multiGpu.value + + +def nvmlDeviceGetBrand(handle): + c_type = _nvmlBrandType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") + ret = fn(handle, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + + +def nvmlDeviceGetC2cModeInfoV1(handle): + c_info = c_nvmlC2cModeInfo_v1_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlDeviceGetC2cModeInfoV(handle): + return nvmlDeviceGetC2cModeInfoV1(handle) + + +@convertStrBytes +def nvmlDeviceGetBoardPartNumber(handle): + c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardPartNumber") + ret = fn(handle, c_part_number, c_uint(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_part_number.value + + +@convertStrBytes +def nvmlDeviceGetSerial(handle): + c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") + ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_serial.value + + +def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): + isReference = type(moduleId) is not c_uint + moduleIdRef = moduleId if isReference else byref(moduleId) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetModuleId") + ret = fn(handle, moduleIdRef) + if isReference: + return ret + else: + _nvmlCheckReturn(ret) + return moduleId.value + + +def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): + affinity_array = c_ulonglong * nodeSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryAffinity") + ret = fn(handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinityWithinScope") + ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") + ret = fn(handle, cpuSetSize, byref(c_affinity)) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceSetCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNumaNodeId(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") + node = c_int() + ret = fn(handle, byref(node)) + _nvmlCheckReturn(ret) + return node.value + + +def nvmlDeviceGetMinorNumber(handle): + c_minor_number = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") + ret = fn(handle, byref(c_minor_number)) + _nvmlCheckReturn(ret) + return c_minor_number.value + + +@convertStrBytes +def nvmlDeviceGetUUID(handle): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") + ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlDeviceGetInforomVersion(handle, infoRomObject): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") + ret = fn( + handle, + _nvmlInforomObject_t(infoRomObject), + c_version, + c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE), + ) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 4.304 +@convertStrBytes +def nvmlDeviceGetInforomImageVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 4.304 +def nvmlDeviceGetInforomConfigurationChecksum(handle): + c_checksum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") + ret = fn(handle, byref(c_checksum)) + _nvmlCheckReturn(ret) + return c_checksum.value + + +# Added in 4.304 +def nvmlDeviceValidateInforom(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetLastBBXFlushTime(handle): + c_timestamp = c_ulonglong() + c_durationUs = c_ulong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetLastBBXFlushTime") + ret = fn(handle, byref(c_timestamp), byref(c_durationUs)) + _nvmlCheckReturn(ret) + return [c_timestamp.value, c_durationUs.value] + + +def nvmlDeviceGetDisplayMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetDisplayActive(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetPersistenceMode(handle): + c_state = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") + ret = fn(handle, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +def nvmlDeviceGetPciInfoExt(handle, c_info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") + ret = fn(handle, c_info) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetPciInfo_v3(handle): + c_info = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlDeviceGetPciInfo(handle): + return nvmlDeviceGetPciInfo_v3(handle) + + +def nvmlDeviceGetClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 2.285 +def nvmlDeviceGetMaxClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 4.304 +def nvmlDeviceGetApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +def nvmlDeviceGetMaxCustomerBoostClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +def nvmlDeviceGetClock(handle, type, id): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") + ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 5.319 +def nvmlDeviceGetDefaultApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 4.304 +def nvmlDeviceGetSupportedMemoryClocks(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no clocks + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + + +# Added in 4.304 +def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no clocks + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetFanSpeed(handle): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") + ret = fn(handle, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetFanSpeed_v2(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +class c_nvmlFanSpeedInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("fan", c_uint), + ("speed", c_uint), + ] + + +nvmlFanSpeedInfo_v1 = 0x100000C + + +def nvmlDeviceGetFanSpeedRPM(handle): + c_fanSpeed = c_nvmlFanSpeedInfo_t() + c_fanSpeed.fan = 0 + c_fanSpeed.version = nvmlFanSpeedInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeedRPM") + ret = fn(handle, byref(c_fanSpeed)) + _nvmlCheckReturn(ret) + return c_fanSpeed.speed + + +def nvmlDeviceGetTargetFanSpeed(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetNumFans(device): + c_numFans = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") + ret = fn(device, byref(c_numFans)) + _nvmlCheckReturn(ret) + return c_numFans.value + + +def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2") + ret = fn(handle, index) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): + isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) + minSpeedRef = minSpeed if isReference else byref(minSpeed) + maxSpeedRef = maxSpeed if isReference else byref(maxSpeed) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxFanSpeed") + ret = fn(handle, minSpeedRef, maxSpeedRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + + +def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): + isReference = type(fanControlPolicy) is not c_uint + fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanControlPolicy_v2") + ret = fn(handle, fan, fanControlPolicyRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else fanControlPolicy.value + + +def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") + ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +class c_nvmlTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("sensorType", _nvmlTemperatureSensors_t), + ("temperature", c_int), + ] + + +nvmlTemperature_v1 = 0x100000C + + +def nvmlDeviceGetTemperatureV1(handle, sensor): + c_temp = c_nvmlTemperature_v1_t() + c_temp.version = nvmlTemperature_v1 + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") + ret = fn(handle, byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.temperature + + +def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): + if version == nvmlTemperature_v1: + return nvmlDeviceGetTemperatureV1(handle, sensor) + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + +# DEPRECATED use nvmlDeviceGetTemperatureV instead +def nvmlDeviceGetTemperature(handle, sensor): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") + ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlDeviceGetTemperatureThreshold(handle, threshold): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): + c_temp = c_uint() + c_temp.value = temp + fn = _nvmlGetFunctionPointer("nvmlDeviceSetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetMarginTemperature(handle): + c_marginTempInfo = c_nvmlMarginTemperature_v1_t() + c_marginTempInfo.version = nvmlMarginTemperature_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMarginTemperature") + ret = fn(handle, byref(c_marginTempInfo)) + _nvmlCheckReturn(ret) + return c_marginTempInfo.marginTemperature + + +# DEPRECATED use nvmlDeviceGetPerformanceState +def nvmlDeviceGetPowerState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + + +def nvmlDeviceGetPerformanceState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + + +def nvmlDeviceGetPowerManagementMode(handle): + c_pcapMode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") + ret = fn(handle, byref(c_pcapMode)) + _nvmlCheckReturn(ret) + return c_pcapMode.value + + +def nvmlDeviceGetPowerManagementLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 4.304 +def nvmlDeviceGetPowerManagementLimitConstraints(handle): + c_minLimit = c_uint() + c_maxLimit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") + ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) + _nvmlCheckReturn(ret) + return [c_minLimit.value, c_maxLimit.value] + + +# Added in 4.304 +def nvmlDeviceGetPowerManagementDefaultLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 331 +def nvmlDeviceGetEnforcedPowerLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +def nvmlDeviceGetPowerUsage(handle): + c_watts = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") + ret = fn(handle, byref(c_watts)) + _nvmlCheckReturn(ret) + return c_watts.value + + +def nvmlDeviceGetTotalEnergyConsumption(handle): + c_millijoules = c_uint64() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") + ret = fn(handle, byref(c_millijoules)) + _nvmlCheckReturn(ret) + return c_millijoules.value + + +# Added in 4.304 +def nvmlDeviceGetGpuOperationMode(handle): + c_currState = _nvmlGpuOperationMode_t() + c_pendingState = _nvmlGpuOperationMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + + +# Added in 4.304 +def nvmlDeviceGetCurrentGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[0] + + +# Added in 4.304 +def nvmlDeviceGetPendingGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[1] + + +def nvmlDeviceGetMemoryInfo(handle, version=None): + if not version: + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") + else: + c_memory = c_nvmlMemory_v2_t() + c_memory.version = version + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo_v2") + ret = fn(handle, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + + +def nvmlDeviceGetBAR1MemoryInfo(handle): + c_bar1_memory = c_nvmlBAR1Memory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") + ret = fn(handle, byref(c_bar1_memory)) + _nvmlCheckReturn(ret) + return c_bar1_memory + + +def nvmlDeviceGetComputeMode(handle): + c_mode = _nvmlComputeMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetCudaComputeCapability(handle): + c_major = c_int() + c_minor = c_int() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") + ret = fn(handle, byref(c_major), byref(c_minor)) + _nvmlCheckReturn(ret) + return (c_major.value, c_minor.value) + + +def nvmlDeviceGetEccMode(handle): + c_currState = _nvmlEnableState_t() + c_pendingState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + + +# added to API +def nvmlDeviceGetCurrentEccMode(handle): + return nvmlDeviceGetEccMode(handle)[0] + + +# added to API +def nvmlDeviceGetPendingEccMode(handle): + return nvmlDeviceGetEccMode(handle)[1] + + +def nvmlDeviceGetDefaultEccMode(handle): + c_defaultState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") + ret = fn(handle, byref(c_defaultState)) + _nvmlCheckReturn(ret) + return [c_defaultState.value] + + +def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_count), + ) + _nvmlCheckReturn(ret) + return c_count.value + + +# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter +def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): + c_counts = c_nvmlEccErrorCounts_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_counts), + ) + _nvmlCheckReturn(ret) + return c_counts + + +# Added in 4.304 +def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), + byref(c_count), + ) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetUtilizationRates(handle): + c_util = c_nvmlUtilization_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") + ret = fn(handle, byref(c_util)) + _nvmlCheckReturn(ret) + return c_util + + +def nvmlDeviceGetEncoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetDecoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetJpgUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetJpgUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetOfaUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetOfaUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetPcieReplayCounter(handle): + c_replay = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") + ret = fn(handle, byref(c_replay)) + _nvmlCheckReturn(ret) + return c_replay.value + + +def nvmlDeviceGetDriverModel(handle): + c_currModel = _nvmlDriverModel_t() + c_pendingModel = _nvmlDriverModel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") + ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) + _nvmlCheckReturn(ret) + return [c_currModel.value, c_pendingModel.value] + + +# added to API +def nvmlDeviceGetCurrentDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[0] + + +# added to API +def nvmlDeviceGetPendingDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[1] + + +# Added in 2.285 +@convertStrBytes +def nvmlDeviceGetVbiosVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +@throwOnVersionMismatch +def nvmlDeviceGetComputeRunningProcesses(handle): + return nvmlDeviceGetComputeRunningProcesses_v3(handle) + + +def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +@throwOnVersionMismatch +def nvmlDeviceGetGraphicsRunningProcesses(handle): + return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + + +@throwOnVersionMismatch +def nvmlDeviceGetMPSComputeRunningProcesses(handle): + return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + + +def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): + c_processDetailList = c_nvmlProcessDetailList_t() + c_processDetailList.version = version + c_processDetailList.mode = mode + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRunningProcessDetailList") + + # first call to get the size + ret = fn(handle, byref(c_processDetailList)) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries + c_processDetailList.procArray = cast( + (c_procs)(), POINTER(c_nvmlProcessDetail_v1_t) + ) + + # make the call again + ret = fn(handle, byref(c_processDetailList)) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_processDetailList.numProcArrayEntries): + # use an alternative struct for this object + obj = c_processDetailList.procArray[i] + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + obj.usedGpuMemory = None + if obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + obj.usedGpuCcProtectedMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetAutoBoostedClocksEnabled(handle): + c_isEnabled = _nvmlEnableState_t() + c_defaultIsEnabled = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") + ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) + _nvmlCheckReturn(ret) + return [c_isEnabled.value, c_defaultIsEnabled.value] + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +## Set functions +def nvmlUnitSetLedState(unit, color): + fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") + ret = fn(unit, _nvmlLedColor_t(color)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetPersistenceMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetComputeMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") + ret = fn(handle, _nvmlComputeMode_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetEccMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearEccErrorCounts(handle, counterType): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") + ret = fn(handle, _nvmlEccCounterType_t(counterType)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetDriverModel(handle, model): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") + ret = fn(handle, _nvmlDriverModel_t(model)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled)) + _nvmlCheckReturn(ret) + return None + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) + _nvmlCheckReturn(ret) + return None + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") + ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetGpuLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") + ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetMemoryLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): + isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t + c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClkMonStatus") + ret = fn(handle, c_clkMonInfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_clkMonInfo + + +# Added in 4.304 +def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") + ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceResetApplicationsClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceSetPowerManagementLimit(handle, limit): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") + ret = fn(handle, c_uint(limit)) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceSetGpuOperationMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") + ret = fn(handle, _nvmlGpuOperationMode_t(mode)) + _nvmlCheckReturn(ret) + return None + + +# Added in 2.285 +def nvmlEventSetCreate(): + fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") + eventSet = c_nvmlEventSet_t() + ret = fn(byref(eventSet)) + _nvmlCheckReturn(ret) + return eventSet + + +# Added in 2.285 +def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): + fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") + ret = fn(handle, c_ulonglong(eventTypes), eventSet) + _nvmlCheckReturn(ret) + return None + + +# Added in 2.285 +def nvmlDeviceGetSupportedEventTypes(handle): + c_eventTypes = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") + ret = fn(handle, byref(c_eventTypes)) + _nvmlCheckReturn(ret) + return c_eventTypes.value + + +# raises NVML_ERROR_TIMEOUT exception on timeout +def nvmlEventSetWait_v2(eventSet, timeoutms): + fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") + data = c_nvmlEventData_t() + ret = fn(eventSet, byref(data), c_uint(timeoutms)) + _nvmlCheckReturn(ret) + return data + + +def nvmlEventSetWait(eventSet, timeoutms): + return nvmlEventSetWait_v2(eventSet, timeoutms) + + +# Added in 2.285 +def nvmlEventSetFree(eventSet): + fn = _nvmlGetFunctionPointer("nvmlEventSetFree") + ret = fn(eventSet) + _nvmlCheckReturn(ret) + return None + + +# Added in 3.295 +def nvmlDeviceOnSameBoard(handle1, handle2): + fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") + onSameBoard = c_int() + ret = fn(handle1, handle2, byref(onSameBoard)) + _nvmlCheckReturn(ret) + return onSameBoard.value != 0 + + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + + +def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 4.304 +def nvmlDeviceGetSupportedClocksThrottleReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +def nvmlDeviceGetSupportedClocksEventReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +# Added in 4.304 +def nvmlDeviceGetCurrentClocksThrottleReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +def nvmlDeviceGetCurrentClocksEventReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +# Added in 5.319 +def nvmlDeviceGetIndex(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") + c_index = c_uint() + ret = fn(handle, byref(c_index)) + _nvmlCheckReturn(ret) + return c_index.value + + +# Added in 5.319 +def nvmlDeviceGetAccountingMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceSetAccountingMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearAccountingPids(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetAccountingStats(handle, pid): + stats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") + ret = fn(handle, c_uint(pid), byref(stats)) + _nvmlCheckReturn(ret) + if stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + stats.maxMemoryUsage = None + return stats + + +def nvmlDeviceGetAccountingPids(handle): + count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) + pids = (c_uint * count.value)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") + ret = fn(handle, byref(count), pids) + _nvmlCheckReturn(ret) + return list(map(int, pids[0 : count.value])) + + +def nvmlDeviceGetAccountingBufferSize(handle): + bufferSize = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") + ret = fn(handle, byref(bufferSize)) + _nvmlCheckReturn(ret) + return int(bufferSize.value) + + +def nvmlDeviceGetRetiredPages(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + ret = fn(device, c_source, byref(c_count), c_pages) + _nvmlCheckReturn(ret) + return list(map(int, c_pages[0 : c_count.value])) + + +def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages_v2") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + times_array = c_ulonglong * c_count.value + c_times = times_array() + ret = fn(device, c_source, byref(c_count), c_pages, c_times) + _nvmlCheckReturn(ret) + return [ + {"address": int(c_pages[i]), "timestamp": int(c_times[i])} + for i in range(c_count.value) + ] + + +def nvmlDeviceGetRetiredPagesPendingStatus(device): + c_pending = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") + ret = fn(device, byref(c_pending)) + _nvmlCheckReturn(ret) + return int(c_pending.value) + + +def nvmlDeviceGetAPIRestriction(device, apiType): + c_permission = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") + ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) + _nvmlCheckReturn(ret) + return int(c_permission.value) + + +def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") + ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetBridgeChipInfo(handle): + bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") + ret = fn(handle, byref(bridgeHierarchy)) + _nvmlCheckReturn(ret) + return bridgeHierarchy + + +def nvmlDeviceGetSamples(device, sampling_type, timeStamp): + c_sampling_type = _nvmlSamplingType_t(sampling_type) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_count = c_uint(0) + c_sample_value_type = _nvmlValueType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") + + ## First Call gets the size + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + None, + ) + + # Stop if this fails + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + sampleArray = c_sample_count.value * c_nvmlSample_t + c_samples = sampleArray() + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + c_samples, + ) + _nvmlCheckReturn(ret) + return (c_sample_value_type.value, c_samples[0 : c_sample_count.value]) + + +def nvmlDeviceGetViolationStatus(device, perfPolicyType): + c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) + c_violTime = c_nvmlViolationTime_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") + + ## Invoke the method to get violation time + ret = fn(device, c_perfPolicy_type, byref(c_violTime)) + _nvmlCheckReturn(ret) + return c_violTime + + +def nvmlDeviceGetPcieThroughput(device, counter): + c_util = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") + ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) + _nvmlCheckReturn(ret) + return c_util.value + + +def nvmlSystemGetTopologyGpuSet(cpuNumber): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") + + # First call will get the size + ret = fn(cpuNumber, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(cpuNumber, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0 : c_count.value]) + + +def nvmlDeviceGetTopologyNearestGpus(device, level): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") + + # First call will get the size + ret = fn(device, level, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(device, level, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0 : c_count.value]) + + +def nvmlDeviceGetTopologyCommonAncestor(device1, device2): + c_level = _nvmlGpuTopologyLevel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") + ret = fn(device1, device2, byref(c_level)) + _nvmlCheckReturn(ret) + return c_level.value + + +def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): + c_rxcounter = c_ulonglong() + c_txcounter = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationCounter") + ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter)) + _nvmlCheckReturn(ret) + return (c_rxcounter.value, c_txcounter.value) + + +def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): + fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") + ret = fn(device, link, counter, freeze) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") + ret = fn(device, link, counter) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(control), reset) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): + c_control = nvmlNvLinkUtilizationControl_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(c_control)) + _nvmlCheckReturn(ret) + return c_control + + +def nvmlDeviceGetNvLinkCapability(device, link, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") + ret = fn(device, link, capability, byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): + c_result = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") + ret = fn(device, link, counter, byref(c_result)) + _nvmlCheckReturn(ret) + return c_result.value + + +def nvmlDeviceResetNvLinkErrorCounters(device, link): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") + ret = fn(device, link) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNvLinkRemotePciInfo(device, link): + c_pci = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") + ret = fn(device, link, byref(c_pci)) + _nvmlCheckReturn(ret) + return c_pci + + +def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): + c_type = _nvmlNvLinkDeviceType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") + ret = fn(handle, link, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + + +def nvmlDeviceGetNvLinkState(device, link): + c_isActive = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") + ret = fn(device, link, byref(c_isActive)) + _nvmlCheckReturn(ret) + return c_isActive.value + + +def nvmlDeviceGetNvLinkVersion(device, link): + c_version = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") + ret = fn(device, link, byref(c_version)) + _nvmlCheckReturn(ret) + return c_version.value + + +def nvmlDeviceModifyDrainState(pciInfo, newState): + fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") + ret = fn(pointer(pciInfo), newState) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceQueryDrainState(pciInfo): + c_newState = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") + ret = fn(pointer(pciInfo), byref(c_newState)) + _nvmlCheckReturn(ret) + return c_newState.value + + +def nvmlDeviceRemoveGpu(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceDiscoverGpus(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + + +def nvmlDeviceClearFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceClearFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + + +def nvmlDeviceGetVirtualizationMode(handle): + c_virtualization_mode = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") + ret = fn(handle, byref(c_virtualization_mode)) + _nvmlCheckReturn(ret) + return c_virtualization_mode.value + + +def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") + return fn(handle, virtualization_mode) + + +def nvmlDeviceGetVgpuHeterogeneousMode(handle): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return c_vgpuHeterogeneousMode.mode + + +def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + c_vgpuHeterogeneousMode.mode = heterogeneous_mode + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlVgpuInstanceGetPlacementId(vgpuInstance): + c_placement = c_nvmlVgpuPlacementId_v1_t(0) + c_placement.version = VgpuPlacementId_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetPlacementId") + ret = fn(vgpuInstance, byref(c_placement)) + _nvmlCheckReturn(ret) + return c_placement.placementId + + +def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + c_vgpu_placements.mode = mode + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + + +def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + + +def nvmlGetVgpuDriverCapabilities(capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") + ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceGetVgpuCapabilities(handle, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceSetVgpuCapabilities(handle, capability, state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetSupportedVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no supported vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetCreatableVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no supported vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): + c_profile_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + ret = fn(vgpuTypeId, byref(c_profile_id)) + _nvmlCheckReturn(ret) + return c_profile_id.value + + +@convertStrBytes +def nvmlVgpuTypeGetClass(vgpuTypeId): + c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_class.value + + +@convertStrBytes +def nvmlVgpuTypeGetName(vgpuTypeId): + c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_name.value + + +def nvmlVgpuTypeGetDeviceID(vgpuTypeId): + c_device_id = c_ulonglong(0) + c_subsystem_id = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) + _nvmlCheckReturn(ret) + return (c_device_id.value, c_subsystem_id.value) + + +def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): + c_fb_size = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + ret = fn(vgpuTypeId, byref(c_fb_size)) + _nvmlCheckReturn(ret) + return c_fb_size.value + + +def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): + c_num_heads = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + ret = fn(vgpuTypeId, byref(c_num_heads)) + _nvmlCheckReturn(ret) + return c_num_heads.value + + +def nvmlVgpuTypeGetResolution(vgpuTypeId): + c_xdim = c_uint(0) + c_ydim = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) + _nvmlCheckReturn(ret) + return (c_xdim.value, c_ydim.value) + + +@convertStrBytes +def nvmlVgpuTypeGetLicense(vgpuTypeId): + c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + ret = fn(vgpuTypeId, c_license, c_buffer_size) + _nvmlCheckReturn(ret) + return c_license.value + + +def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): + c_frl_config = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + ret = fn(vgpuTypeId, byref(c_frl_config)) + _nvmlCheckReturn(ret) + return c_frl_config.value + + +def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): + c_gsp_heap = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + ret = fn(vgpuTypeId, byref(c_gsp_heap)) + _nvmlCheckReturn(ret) + return c_gsp_heap.value + + +def nvmlVgpuTypeGetFbReservation(vgpuTypeId): + c_fb_reservation = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + ret = fn(vgpuTypeId, byref(c_fb_reservation)) + _nvmlCheckReturn(ret) + return c_fb_reservation.value + + +def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): + c_runtime_state = nvmlVgpuRuntimeState_v1_t() + c_runtime_state.version = VgpuRuntimeState_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + ret = fn(vgpuInstance, byref(c_runtime_state)) + _nvmlCheckReturn(ret) + return c_runtime_state + + +def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + return c_max_instances.value + + +def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): + c_max_instances_per_vm = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) + _nvmlCheckReturn(ret) + return c_max_instances_per_vm.value + + +def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): + c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) + c_bar1Info.version = VgpuTypeBar1Info_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + ret = fn(vgpuTypeId, byref(c_bar1Info)) + _nvmlCheckReturn(ret) + return c_bar1Info + + +def nvmlDeviceGetActiveVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value + c_vgpu_instances = vgpu_instance_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_instances[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +@convertStrBytes +def nvmlVgpuInstanceGetVmID(vgpuInstance): + c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) + _nvmlCheckReturn(ret) + return (c_vm_id.value, c_vm_id_type.value) + + +@convertStrBytes +def nvmlVgpuInstanceGetUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): + c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) + _nvmlCheckReturn(ret) + return c_driver_version.value + + +def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): + c_license_status = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + ret = fn(vgpuInstance, byref(c_license_status)) + _nvmlCheckReturn(ret) + return c_license_status.value + + +def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + c_license_info = c_nvmlVgpuLicenseInfo_t() + ret = fn(vgpuInstance, byref(c_license_info)) + _nvmlCheckReturn(ret) + return c_license_info + + +def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): + return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + + +def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): + c_frl = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + ret = fn(vgpuInstance, byref(c_frl)) + _nvmlCheckReturn(ret) + return c_frl.value + + +def nvmlVgpuInstanceGetEccMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlVgpuInstanceGetType(vgpuInstance): + c_vgpu_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + ret = fn(vgpuInstance, byref(c_vgpu_type)) + _nvmlCheckReturn(ret) + return c_vgpu_type.value + + +def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): + c_encoder_capacity = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + ret = fn(vgpuInstance, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + + +def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + return fn(vgpuInstance, encoder_capacity) + + +def nvmlVgpuInstanceGetFbUsage(vgpuInstance): + c_fb_usage = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + ret = fn(vgpuInstance, byref(c_fb_usage)) + _nvmlCheckReturn(ret) + return c_fb_usage.value + + +def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): + c_cap_result = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) + _nvmlCheckReturn(ret) + return c_cap_result.value + + +def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): + c_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + ret = fn(vgpuInstance, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + + +@convertStrBytes +def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): + c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") + ret = fn( + vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)) + ) + _nvmlCheckReturn(ret) + return c_vgpuPciId.value + + +def nvmlDeviceGetVgpuUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_value_type = _nvmlValueType_t() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + ret = fn( + handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None + ) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn( + handle, + c_time_stamp, + byref(c_sample_value_type), + byref(c_vgpu_count), + c_samples, + ) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0) + c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1 + c_vgpuUtilInfo.sampleValType = _nvmlValueType_t() + c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) + c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + ret = fn(handle, byref(c_vgpuUtilInfo)) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_vgpuUtilInfo.vgpuUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpuUtilInfo.vgpuInstanceCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): + c_p2pstatus = _nvmlGpuP2PStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") + ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus)) + _nvmlCheckReturn(ret) + return c_p2pstatus.value + + +def nvmlDeviceGetGridLicensableFeatures_v4(handle): + c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") + ret = fn(handle, byref(c_get_grid_licensable_features)) + _nvmlCheckReturn(ret) + + return c_get_grid_licensable_features + + +def nvmlDeviceGetGridLicensableFeatures(handle): + return nvmlDeviceGetGridLicensableFeatures_v4(handle) + + +def nvmlDeviceGetGspFirmwareVersion(handle, version=None): + isUserDefined = version is not None + if not isUserDefined: + version = (c_char * NVML_GSP_FIRMWARE_VERSION_BUF_SIZE)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareVersion") + ret = fn(handle, version) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else version.value + + +def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): + isReference = type(isEnabled) is not c_uint + isEnabledRef = isEnabled if isReference else byref(isEnabled) + defaultModeRef = defaultMode if isReference else byref(defaultMode) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareMode") + ret = fn(handle, isEnabledRef, defaultModeRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + + +def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): + c_encoder_capacity = c_ulonglong(0) + c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderCapacity") + ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + + +def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0) + c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1 + c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) + c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpuProcUtilInfo.vgpuProcessCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetEncoderStats(handle): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderStats") + ret = fn(handle, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + + +def nvmlDeviceGetEncoderSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + ret = fn(handle, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetFBCStats(handle): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") + ret = fn(handle, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + + +def nvmlDeviceGetFBCSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + ret = fn(handle, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + ret = fn( + vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency) + ) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + + +def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetFBCStats(vgpuInstance): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") + ret = fn(vgpuInstance, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + + +def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetProcessUtilization(handle, timeStamp): + # first call to get the size + c_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + ret = fn(handle, None, byref(c_count), c_time_stamp) + + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_samples, byref(c_count), c_time_stamp) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0) + c_processesUtilInfo.version = ProcessesUtilizationInfo_v1 + c_processesUtilInfo.processSamplesCount = c_uint(0) + c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + ret = fn(handle, byref(c_processesUtilInfo)) + + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_processesUtilInfo.procUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_processesUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_processesUtilInfo.processSamplesCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetMetadata(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") + c_vgpuMetadata = c_nvmlVgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuMetadata + + +def nvmlDeviceGetVgpuMetadata(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") + c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuPgpuMetadata + + +def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): + fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") + c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() + ret = fn(byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility)) + _nvmlCheckReturn(ret) + return c_vgpuPgpuCompatibility + + +@convertStrBytes +def nvmlDeviceGetPgpuMetadataString(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") + c_pgpuMetadata = create_string_buffer(NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pgpuMetadata.value, c_bufferSize.value) + + +def nvmlDeviceGetVgpuSchedulerLog(handle): + c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") + ret = fn(handle, byref(c_vgpu_sched_log)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_log + + +def nvmlDeviceGetVgpuSchedulerState(handle): + c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") + ret = fn(handle, byref(c_vgpu_sched_state)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_state + + +def nvmlDeviceGetVgpuSchedulerCapabilities(handle): + c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") + ret = fn(handle, byref(c_vgpu_sched_caps)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_caps + + +def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") + ret = fn(handle, byref(sched_state)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSetVgpuVersion(vgpuVersion): + fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") + ret = fn(byref(vgpuVersion)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGetVgpuVersion(supported=None, current=None): + isUserDefined = (supported is not None) or (current is not None) + if not isUserDefined: + supported = c_nvmlVgpuVersion_t() + current = c_nvmlVgpuVersion_t() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") + ret = fn(byref(supported), byref(current)) + _nvmlCheckReturn(ret) + return ( + NVML_SUCCESS + if isUserDefined + else [ + (supported.minVersion, supported.maxVersion), + (current.minVersion, current.maxVersion), + ] + ) + + +def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): + c_pidCount = c_uint() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") + ret = fn(vgpuInstance, byref(c_pidCount), None) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + sampleArray = c_pidCount.value * c_uint + c_pidArray = sampleArray() + ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pidCount, c_pidArray) + + +def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): + c_accountingStats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") + ret = fn(vgpuInstance, pid, byref(c_accountingStats)) + _nvmlCheckReturn(ret) + return c_accountingStats + + +def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") + ret = fn(vgpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGetExcludedDeviceCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlGetExcludedDeviceInfoByIndex(index): + c_index = c_uint(index) + info = c_nvmlExcludedDeviceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceInfoByIndex") + ret = fn(c_index, byref(info)) + _nvmlCheckReturn(ret) + return info + + +def nvmlDeviceGetHostVgpuMode(handle): + c_host_vgpu_mode = _nvmlHostVgpuMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") + ret = fn(handle, byref(c_host_vgpu_mode)) + _nvmlCheckReturn(ret) + return c_host_vgpu_mode.value + + +def nvmlDeviceSetMigMode(device, mode): + c_activationStatus = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") + ret = fn(device, mode, byref(c_activationStatus)) + _nvmlCheckReturn(ret) + return c_activationStatus.value + + +def nvmlDeviceGetMigMode(device): + c_currentMode = c_uint() + c_pendingMode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigMode") + ret = fn(device, byref(c_currentMode), byref(c_pendingMode)) + _nvmlCheckReturn(ret) + return [c_currentMode.value, c_pendingMode.value] + + +def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): + if version == 2: + c_info = c_nvmlGpuInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlGpuInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +# Define function alias for the API exposed by NVML +nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo + + +def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") + ret = fn(device, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetGpuInstancePossiblePlacements( + device, profileId, placementsRef, countRef +): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") + ret = fn(device, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceCreateGpuInstance(device, profileId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") + ret = fn(device, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") + ret = fn(device, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceDestroy(gpuInstance): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") + ret = fn(gpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") + ret = fn(device, profileId, gpuInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") + ret = fn(device, gpuInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceGetInfo(gpuInstance): + c_info = c_nvmlGpuInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") + ret = fn(gpuInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlGpuInstanceGetComputeInstanceProfileInfo( + device, profile, engProfile, version=2 +): + if version == 2: + c_info = c_nvmlComputeInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlComputeInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, engProfile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +# Define function alias for the API exposed by NVML +nvmlGpuInstanceGetComputeInstanceProfileInfoV = ( + nvmlGpuInstanceGetComputeInstanceProfileInfo +) + + +def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceRemainingCapacity") + ret = fn(gpuInstance, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlGpuInstanceGetComputeInstancePossiblePlacements( + gpuInstance, profileId, placementsRef, countRef +): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") + ret = fn(gpuInstance, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") + ret = fn(gpuInstance, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceCreateComputeInstanceWithPlacement( + gpuInstance, profileId, placement +): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") + ret = fn(gpuInstance, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlComputeInstanceDestroy(computeInstance): + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") + ret = fn(computeInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceGetComputeInstances( + gpuInstance, profileId, computeInstancesRef, countRef +): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") + ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") + ret = fn(gpuInstance, computeInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlComputeInstanceGetInfo_v2(computeInstance): + c_info = c_nvmlComputeInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") + ret = fn(computeInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlComputeInstanceGetInfo(computeInstance): + return nvmlComputeInstanceGetInfo_v2(computeInstance) + + +def nvmlDeviceIsMigDeviceHandle(device): + c_isMigDevice = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") + ret = fn(device, byref(c_isMigDevice)) + _nvmlCheckReturn(ret) + return c_isMigDevice + + +def nvmlDeviceGetGpuInstanceId(device): + c_gpuInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") + ret = fn(device, byref(c_gpuInstanceId)) + _nvmlCheckReturn(ret) + return c_gpuInstanceId.value + + +def nvmlDeviceGetComputeInstanceId(device): + c_computeInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") + ret = fn(device, byref(c_computeInstanceId)) + _nvmlCheckReturn(ret) + return c_computeInstanceId.value + + +def nvmlDeviceGetMaxMigDeviceCount(device): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") + ret = fn(device, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetMigDeviceHandleByIndex(device, index): + c_index = c_uint(index) + migDevice = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigDeviceHandleByIndex") + ret = fn(device, c_index, byref(migDevice)) + _nvmlCheckReturn(ret) + return migDevice + + +def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") + ret = fn(migDevice, byref(device)) + _nvmlCheckReturn(ret) + return device + + +def nvmlDeviceGetAttributes_v2(device): + c_attrs = c_nvmlDeviceAttributes() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") + ret = fn(device, byref(c_attrs)) + _nvmlCheckReturn(ret) + return c_attrs + + +def nvmlDeviceGetAttributes(device): + return nvmlDeviceGetAttributes_v2(device) + + +def nvmlDeviceGetRemappedRows(device): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") + c_corr = c_uint() + c_unc = c_uint() + c_bpending = c_uint() + c_bfailure = c_uint() + ret = fn(device, byref(c_corr), byref(c_unc), byref(c_bpending), byref(c_bfailure)) + _nvmlCheckReturn(ret) + return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + + +def nvmlDeviceGetRowRemapperHistogram(device): + c_vals = c_nvmlRowRemapperHistogramValues() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") + ret = fn(device, byref(c_vals)) + _nvmlCheckReturn(ret) + return c_vals + + +def nvmlDeviceGetArchitecture(device): + arch = _nvmlDeviceArchitecture_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") + ret = fn(device, byref(arch)) + _nvmlCheckReturn(ret) + return arch.value + + +def nvmlDeviceGetBusType(device): + c_busType = _nvmlBusType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") + ret = fn(device, byref(c_busType)) + _nvmlCheckReturn(ret) + return c_busType.value + + +def nvmlDeviceGetIrqNum(device): + c_irqNum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") + ret = fn(device, byref(c_irqNum)) + _nvmlCheckReturn(ret) + return c_irqNum.value + + +def nvmlDeviceGetNumGpuCores(device): + c_numCores = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") + ret = fn(device, byref(c_numCores)) + _nvmlCheckReturn(ret) + return c_numCores.value + + +def nvmlDeviceGetPowerSource(device): + c_powerSource = _nvmlPowerSource_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") + ret = fn(device, byref(c_powerSource)) + _nvmlCheckReturn(ret) + return c_powerSource.value + + +def nvmlDeviceGetMemoryBusWidth(device): + c_memBusWidth = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") + ret = fn(device, byref(c_memBusWidth)) + _nvmlCheckReturn(ret) + return c_memBusWidth.value + + +def nvmlDeviceGetPcieLinkMaxSpeed(device): + c_speed = _nvmlPcieLinkMaxSpeed_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetAdaptiveClockInfoStatus(device): + c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") + ret = fn(device, byref(c_adaptiveClockInfoStatus)) + _nvmlCheckReturn(ret) + return c_adaptiveClockInfoStatus.value + + +def nvmlDeviceGetPcieSpeed(device): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetDynamicPstatesInfo( + device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t() +): + isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t + dynamicpstatesinfoRef = ( + c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + ) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo") + ret = fn(device, dynamicpstatesinfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + + +def nvmlDeviceSetFanSpeed_v2(handle, index, speed): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2") + ret = fn(handle, index, speed) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetThermalSettings( + device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t() +): + isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t + thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings") + ret = fn(device, sensorindex, thermalsettingsRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] + + +def nvmlDeviceGetMinMaxClockOfPState( + device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint() +): + isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) + minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) + maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState") + ret = fn( + device, + _nvmlClockType_t(clockType), + _nvmlClockType_t(pstate), + minClockMHzRef, + maxClockMHzRef, + ) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + + +class c_nvmlClockOffset_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("type", _nvmlClockType_t), + ("pstate", _nvmlPstates_t), + ("clockOffsetMHz", c_int), + ("minClockOffsetMHz", c_int), + ("maxClockOffsetMHz", c_int), + ] + + +nvmlClockOffset_v1 = 0x1000018 + + +def nvmlDeviceGetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets") + ret = fn(device, info) + return NVML_SUCCESS + + +def nvmlDeviceSetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets") + ret = fn(device, info) + return NVML_SUCCESS + + +def nvmlDeviceGetSupportedPerformanceStates(device): + pstates = [] + c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) + c_size = sizeof(c_uint) * c_count.value + + # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. + pstates_array = _nvmlPstates_t * c_count.value + c_pstates = pstates_array() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedPerformanceStates") + ret = fn(device, c_pstates, c_size) + _nvmlCheckReturn(ret) + + for value in c_pstates: + if value != NVML_PSTATE_UNKNOWN: + pstates.append(value) + + return pstates + + +def nvmlDeviceGetGpcClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + + +def nvmlDeviceSetGpcClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + + +def nvmlDeviceGetMemClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + + +def nvmlDeviceSetMemClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + + +def nvmlSystemSetConfComputeGpusReadyState(state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") + ret = fn(c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetConfComputeGpusReadyState(): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +def nvmlSystemGetConfComputeCapabilities(): + c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") + ret = fn(byref(c_ccSysCaps)) + _nvmlCheckReturn(ret) + return c_ccSysCaps + + +def nvmlSystemGetConfComputeState(): + c_state = c_nvmlConfComputeSystemState_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + + +def nvmlSystemGetConfComputeSettings(settings): + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") + return fn(settings) + + +def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") + ret = fn(device, c_ccMemSize) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetConfComputeMemSizeInfo(device): + c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") + ret = fn(device, byref(c_ccMemSize)) + _nvmlCheckReturn(ret) + return c_ccMemSize + + +def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") + ret = fn(device, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + + +def nvmlDeviceGetConfComputeGpuCertificate(device): + c_cert = c_nvmlConfComputeGpuCertificate_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") + ret = fn(device, byref(c_cert)) + _nvmlCheckReturn(ret) + return c_cert + + +def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): + c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() + c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) + setattr(c_attestReport, "nonce", c_nonce_arr) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") + ret = fn(device, byref(c_attestReport)) + _nvmlCheckReturn(ret) + return c_attestReport + + +def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): + c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 + c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): + c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return c_keyRotationThrInfo + + +## GPM ## +######### + +## Enums/defines + +#### GPM Metric Identifiers +NVML_GPM_METRIC_GRAPHICS_UTIL = ( + 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +) +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = ( + 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +) +NVML_GPM_METRIC_INTEGER_UTIL = ( + 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_ANY_TENSOR_UTIL = ( + 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = ( + 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = ( + 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = ( + 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +) +NVML_GPM_METRIC_DRAM_BW_UTIL = ( + 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP64_UTIL = ( + 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP32_UTIL = ( + 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP16_UTIL = ( + 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +) +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = ( + 60 # NvLink read bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = ( + 61 # NvLink write bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = ( + 63 # NvLink write bandwidth for link 0 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = ( + 65 # NvLink write bandwidth for link 1 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = ( + 67 # NvLink write bandwidth for link 2 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = ( + 69 # NvLink write bandwidth for link 3 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = ( + 71 # NvLink write bandwidth for link 4 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = ( + 73 # NvLink write bandwidth for link 5 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = ( + 75 # NvLink write bandwidth for link 6 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = ( + 77 # NvLink write bandwidth for link 7 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = ( + 79 # NvLink write bandwidth for link 8 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = ( + 81 # NvLink write bandwidth for link 9 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = ( + 82 # NvLink read bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = ( + 83 # NvLink write bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = ( + 84 # NvLink read bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = ( + 85 # NvLink write bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = ( + 86 # NvLink read bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = ( + 87 # NvLink write bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = ( + 88 # NvLink read bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = ( + 89 # NvLink write bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = ( + 90 # NvLink read bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = ( + 91 # NvLink write bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = ( + 92 # NvLink read bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = ( + 93 # NvLink write bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = ( + 94 # NvLink read bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = ( + 95 # NvLink write bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = ( + 96 # NvLink read bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = ( + 97 # NvLink write bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_MAX = 98 + +## Structs + + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), + ] + + +class struct_c_nvmlGpmSample_t(Structure): + pass # opaque handle + + +c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + + +class c_metricInfo_t(Structure): + _fields_ = [ + ("shortName", c_char_p), + ("longName", c_char_p), + ("unit", c_char_p), + ] + + +class c_nvmlGpmMetric_t(_PrintableStructure): + _fields_ = [ + ("metricId", c_uint), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_double), + ("metricInfo", c_metricInfo_t), + ] + + +class c_nvmlGpmMetricsGet_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("numMetrics", c_uint), + ("sample1", c_nvmlGpmSample_t), + ("sample2", c_nvmlGpmSample_t), + ("metrics", c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX), + ] + + +NVML_GPM_METRICS_GET_VERSION = 1 + + +class c_nvmlGpmSupport_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("isSupportedDevice", c_uint), + ] + + +NVML_GPM_SUPPORT_VERSION = 1 + +## Functions + + +def nvmlGpmMetricsGet(metricsGet): + fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") + ret = fn(byref(metricsGet)) + _nvmlCheckReturn(ret) + return metricsGet + + +def nvmlGpmSampleFree(gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") + ret = fn(gpmSample) + _nvmlCheckReturn(ret) + return + + +def nvmlGpmSampleAlloc(): + gpmSample = c_nvmlGpmSample_t() + fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") + ret = fn(byref(gpmSample)) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmSampleGet(device, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") + ret = fn(device, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") + ret = fn(device, gpuInstanceId, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmQueryDeviceSupport(device): + gpmSupport = c_nvmlGpmSupport_t() + gpmSupport.version = NVML_GPM_SUPPORT_VERSION + fn = _nvmlGetFunctionPointer("nvmlGpmQueryDeviceSupport") + ret = fn(device, byref(gpmSupport)) + _nvmlCheckReturn(ret) + return gpmSupport + + +def nvmlGpmSetStreamingEnabled(device, state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") + ret = fn(device, c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpmQueryIfStreamingEnabled(device): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") + ret = fn(device, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +# Low Power Structure and Function + +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 + +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + + +class c_nvmlNvLinkPowerThres_t(Structure): + _fields_ = [ + ("lowPwrThreshold", c_uint), + ] + + +def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): + c_info = c_nvmlNvLinkPowerThres_t() + c_info.lowPwrThreshold = l1threshold + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkDeviceLowPowerThreshold") + ret = fn(device, byref(c_info)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +NVML_GPU_FABRIC_UUID_LEN = 16 + +_nvmlGpuFabricState_t = c_uint +NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + + +class c_nvmlGpuFabricInfo_t(_PrintableStructure): + _fields_ = [ + ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ] + + +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 + +nvmlGpuFabricInfo_v2 = 0x02000024 + + +class c_nvmlGpuFabricInfoV_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ("healthMask", c_uint32), + ] + + def __init__(self): + super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2) + + +def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo") + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV") + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +###################### +## Enums/defines +#### NVML GPU NVLINK BW MODE +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + + +def nvmlSystemSetNvlinkBwMode(mode): + fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") + ret = fn(mode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetNvlinkBwMode(): + mode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") + ret = fn(byref(mode)) + _nvmlCheckReturn(ret) + return mode.value + + +_nvmlPowerScopeType_t = c_uint +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + + +class c_nvmlPowerValue_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("powerScope", _nvmlPowerScopeType_t), + ("powerValueMw", c_uint), + ] + _fmt_ = {"": "%d B"} + + +nvmlPowerValue_v2 = 0x0200000C + + +def nvmlDeviceSetPowerManagementLimit_v2( + device, powerScope, powerLimit, version=nvmlPowerValue_v2 +): + c_powerScope = _nvmlPowerScopeType_t(powerScope) + c_powerValue = c_nvmlPowerValue_v2_t() + c_powerValue.version = c_uint(version) + c_powerValue.powerScope = c_powerScope + c_powerValue.powerValueMw = c_uint(powerLimit) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit_v2") + ret = fn(device, byref(c_powerValue)) + return NVML_SUCCESS + + +class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("aggregateUncParity", c_ulonglong), + ("aggregateUncSecDed", c_ulonglong), + ("aggregateCor", c_ulonglong), + ("volatileUncParity", c_ulonglong), + ("volatileUncSecDed", c_ulonglong), + ("volatileCor", c_ulonglong), + ("aggregateUncBucketL2", c_ulonglong), + ("aggregateUncBucketSm", c_ulonglong), + ("aggregateUncBucketPcie", c_ulonglong), + ("aggregateUncBucketMcu", c_ulonglong), + ("aggregateUncBucketOther", c_ulonglong), + ("bThresholdExceeded", c_uint), + ] + + def __init__(self): + super(c_nvmlEccSramErrorStatus_v1_t, self).__init__( + version=nvmlEccSramErrorStatus_v1 + ) + + +nvmlEccSramErrorStatus_v1 = 0x1000068 + + +def nvmlDeviceGetSramEccErrorStatus(device, status): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") + ret = fn(device, status) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +NVML_DEV_CAP_EGM = 1 << 0 +nvmlDeviceCapabilities_v1 = 0x1000008 + + +class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("capMask", c_uint), + ] + + def __init__(self): + super(c_nvmlDeviceCapabilities_v1_t, self).__init__( + version=nvmlDeviceCapabilities_v1 + ) + + +def nvmlDeviceGetCapabilities(device, caps): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") + return fn(device, caps) + + +class c_nvmlPlatformInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("ibGuid", c_char * 16), + ("rackGuid", c_char * 16), + ("chassisPhysicalSlotNumber", c_char), + ("computeSlotIndex", c_char), + ("nodeIndex", c_char), + ("peerType", c_char), + ("moduleId", c_char), + ] + + def __init__(self): + super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) + + +nvmlPlatformInfo_v1 = 0x100002C + + +def nvmlDeviceGetPlatformInfo(device, platformInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") + ret = fn(device, platformInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +class c_nvmlMask255_t(_PrintableStructure): + _fields_ = [ + ("mask", c_uint * 8), + ] + + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 + +nvmlWorkloadPowerProfileInfo_v1 = 0x100002C + + +class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("profileId", c_uint), + ("priority", c_uint), + ("conflictingmask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileInfo_v1 + ) + + +nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002BF8 + + +class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ( + "perfProfile", + c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES, + ), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileProfilesInfo_v1 + ) + + +nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 + + +class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ("requestedProfilesMask", c_nvmlMask255_t), + ("enforcedProfilesMask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileCurrentProfiles_v1 + ) + + +nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 + + +class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("requestedProfilesMask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileRequestedProfiles_v1 + ) + + +def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") + ret = fn(device, profilesInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") + ret = fn(device, currentProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") + ret = fn(device, supportedBwModes) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetNvlinkBwMode(device, getBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") + ret = fn(device, getBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceSetNvlinkBwMode(device, setBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") + ret = fn(device, setBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +nvmlDramEncryptionInfo_v1 = 0x01000008 + + +class c_nvmlDramEncryptionInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("encryptionState", _nvmlEnableState_t), + ] + + def __init__(self): + super(c_nvmlDramEncryptionInfo_t, self).__init__( + version=nvmlDramEncryptionInfo_v1 + ) + + +def nvmlDeviceGetDramEncryptionMode(handle): + c_currState = c_nvmlDramEncryptionInfo_t() + c_pendingState = c_nvmlDramEncryptionInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDramEncryptionMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.encryptionState, c_pendingState.encryptionState] + + +# added to API +def nvmlDeviceGetCurrentDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[0] + + +# added to API +def nvmlDeviceGetPendingDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[1] + + +def nvmlDeviceSetDramEncryptionMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") + c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() + c_dramEncryptionMode.encryptionState = mode + ret = fn(handle, byref(c_dramEncryptionMode)) + _nvmlCheckReturn(ret) + return None + + +# Power Smoothing defines +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 + +nvmlPowerSmoothingState_v1 = 0x1000008 + + +class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("state", c_uint), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingState_v1_t, self).__init__( + version=nvmlPowerSmoothingState_v1 + ) + + +nvmlPowerSmoothingProfile_v1 = 0x1000018 + + +class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("profileId", c_uint), + ("paramId", c_uint), + ("value", c_double), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__( + version=nvmlPowerSmoothingProfile_v1 + ) + + +def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + + +def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + + +def nvmlDevicePowerSmoothingSetState(device, state): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") + ret = fn(device, state) + _nvmlCheckReturn(ret) diff --git a/wan/utils.py b/wan/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b9f8cff9ddc6c3b14e9a4a675db77086f3e47f59 --- /dev/null +++ b/wan/utils.py @@ -0,0 +1,845 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py + +import argparse +import ctypes +import hashlib +import importlib +import importlib.util +import inspect +import json +import math +import os +import signal +import socket +import sys +import tempfile +import threading +import traceback +from collections.abc import Callable +from dataclasses import dataclass, fields, is_dataclass +from functools import lru_cache, partial, wraps +from typing import Any, TypeVar, cast + +import cloudpickle +import filelock +import torch +import yaml +from diffusers.loaders.lora_base import ( + _best_guess_weight_name, +) # watch out for potetential removal from diffusers +from huggingface_hub import snapshot_download +from remote_pdb import RemotePdb +from torch.distributed.fsdp import MixedPrecisionPolicy + +import distributed.envs as envs +from distributed.logger import init_logger + +logger = init_logger(__name__) + +T = TypeVar("T") + +# TODO(will): used to convert fastvideo_args.precision to torch.dtype. Find a +# cleaner way to do this. +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +STR_BACKEND_ENV_VAR: str = "FASTVIDEO_ATTENTION_BACKEND" +STR_ATTN_CONFIG_ENV_VAR: str = "FASTVIDEO_ATTENTION_CONFIG" + + +def find_nccl_library() -> str: + """ + We either use the library file specified by the `VLLM_NCCL_SO_PATH` + environment variable, or we find the library file brought by PyTorch. + After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be + found by `ctypes` automatically. + """ + so_file = envs.FASTVIDEO_NCCL_SO_PATH + + # manually load the nccl library + if so_file: + logger.info( + "Found nccl from environment variable FASTVIDEO_NCCL_SO_PATH=%s", so_file + ) + else: + if torch.version.cuda is not None: + so_file = "libnccl.so.2" + elif torch.version.hip is not None: + so_file = "librccl.so.1" + else: + raise ValueError("NCCL only supports CUDA and ROCm backends.") + logger.info("Found nccl from library %s", so_file) + return str(so_file) + + +prev_set_stream = torch.cuda.set_stream + +_current_stream = None + + +def _patched_set_stream(stream: torch.cuda.Stream | None) -> None: + global _current_stream + _current_stream = stream + if stream is not None: + prev_set_stream(stream) + + +torch.cuda.set_stream = _patched_set_stream + + +def current_stream() -> torch.cuda.Stream | None: + """ + replace `torch.cuda.current_stream()` with `fastvideo.utils.current_stream()`. + it turns out that `torch.cuda.current_stream()` is quite expensive, + as it will construct a new stream object at each call. + here we patch `torch.cuda.set_stream` to keep track of the current stream + directly, so that we can avoid calling `torch.cuda.current_stream()`. + + the underlying hypothesis is that we do not call `torch._C._cuda_setStream` + from C/C++ code. + """ + from platforms import current_platform + + # For non-CUDA platforms, return None + if not current_platform.is_cuda_alike(): + return None + + global _current_stream + if _current_stream is None: + # when this function is called before any stream is set, + # we return the default stream. + # On ROCm using the default 0 stream in combination with RCCL + # is hurting performance. Therefore creating a dedicated stream + # per process + _current_stream = ( + torch.cuda.Stream() + if current_platform.is_rocm() + else torch.cuda.current_stream() + ) + return _current_stream + + +class StoreBoolean(argparse.Action): + + def __init__(self, option_strings, dest, default=False, required=False, help=None): + super().__init__( + option_strings=option_strings, + dest=dest, + nargs="?", + const=True, + default=default, + required=required, + help=help, + ) + + def __call__(self, parser, namespace, values, option_string=None): + if values is None: + setattr(namespace, self.dest, True) + elif isinstance(values, str): + if values.lower() == "true": + setattr(namespace, self.dest, True) + elif values.lower() == "false": + setattr(namespace, self.dest, False) + else: + raise ValueError( + f"Invalid boolean value: {values}. " "Expected 'true' or 'false'." + ) + else: + setattr(namespace, self.dest, bool(values)) + + +class SortedHelpFormatter(argparse.HelpFormatter): + """SortedHelpFormatter that sorts arguments by their option strings.""" + + def add_arguments(self, actions): + actions = sorted(actions, key=lambda x: x.option_strings) + super().add_arguments(actions) + + +class FlexibleArgumentParser(argparse.ArgumentParser): + """ArgumentParser that allows both underscore and dash in names.""" + + def __init__(self, *args, **kwargs) -> None: + # Set the default 'formatter_class' to SortedHelpFormatter + if "formatter_class" not in kwargs: + kwargs["formatter_class"] = SortedHelpFormatter + super().__init__(*args, **kwargs) + + def parse_args( # type: ignore[override] + self, args=None, namespace=None + ) -> argparse.Namespace: + if args is None: + args = sys.argv[1:] + + if "--config" in args: + args = self._pull_args_from_config(args) + + # Convert underscores to dashes and vice versa in argument names + processed_args = [] + for arg in args: + if arg.startswith("--"): + if "=" in arg: + key, value = arg.split("=", 1) + key = "--" + key[len("--") :].replace("_", "-") + processed_args.append(f"{key}={value}") + else: + processed_args.append("--" + arg[len("--") :].replace("_", "-")) + elif arg.startswith("-O") and arg != "-O" and len(arg) == 2: + # allow -O flag to be used without space, e.g. -O3 + processed_args.append("-O") + processed_args.append(arg[2:]) + else: + processed_args.append(arg) + + namespace = super().parse_args(processed_args, namespace) + + # Track which arguments were explicitly provided + namespace._provided = set() + + i = 0 + while i < len(args): + arg = args[i] + if arg.startswith("--"): + # Handle --key=value format + if "=" in arg: + key = arg.split("=")[0][2:].replace("-", "_") + namespace._provided.add(key) + i += 1 + # Handle --key value format + else: + key = arg[2:].replace("-", "_") + namespace._provided.add(key) + # Skip the value if there is one + if i + 1 < len(args) and not args[i + 1].startswith("-"): + i += 2 + else: + i += 1 + else: + i += 1 + + return namespace # type: ignore[no-any-return] + + def _pull_args_from_config(self, args: list[str]) -> list[str]: + """Method to pull arguments specified in the config file + into the command-line args variable. + + The arguments in config file will be inserted between + the argument list. + + example: + ```yaml + port: 12323 + tensor-parallel-size: 4 + ``` + ```python + $: vllm {serve,chat,complete} "facebook/opt-12B" \ + --config config.yaml -tp 2 + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--config', 'config.yaml', + '-tp', '2' + ] + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--port', '12323', + '--tp-size', '4', + '-tp', '2' + ] + ``` + + Please note how the config args are inserted after the sub command. + this way the order of priorities is maintained when these are args + parsed by super(). + """ + assert args.count("--config") <= 1, "More than one config file specified!" + + index = args.index("--config") + if index == len(args) - 1: + raise ValueError( + "No config file specified! \ + Please check your command-line arguments." + ) + + file_path = args[index + 1] + + config_args = self._load_config_file(file_path) + + # 0th index is for {serve,chat,complete} + # followed by model_tag (only for serve) + # followed by config args + # followed by rest of cli args. + # maintaining this order will enforce the precedence + # of cli > config > defaults + if args[0] == "serve": + if index == 1: + raise ValueError( + "No model_tag specified! Please check your command-line" + " arguments." + ) + args = ( + [args[0]] + [args[1]] + config_args + args[2:index] + args[index + 2 :] + ) + else: + args = [args[0]] + config_args + args[1:index] + args[index + 2 :] + + return args + + def _load_config_file(self, file_path: str) -> list[str]: + """Loads a yaml file and returns the key value pairs as a + flattened list with argparse like pattern + ```yaml + port: 12323 + tensor-parallel-size: 4 + vae_config: + load_encoder: false + load_decoder: true + ``` + returns: + processed_args: list[str] = [ + '--port': '12323', + '--tp-size': '4', + '--vae-config.load-encoder': 'false', + '--vae-config.load-decoder': 'true' + ] + """ + + extension: str = file_path.split(".")[-1] + if extension not in ("yaml", "yml", "json"): + raise ValueError( + "Config file must be of a yaml/yml/json type.\ + %s supplied", + extension, + ) + + processed_args: list[str] = [] + + config: dict[str, Any] = {} + try: + with open(file_path) as config_file: + config = yaml.safe_load(config_file) + except Exception as ex: + logger.error( + "Unable to read the config file at %s. \ + Make sure path is correct", + file_path, + ) + raise ex + + store_boolean_arguments = [ + action.dest for action in self._actions if isinstance(action, StoreBoolean) + ] + + def process_dict(prefix: str, d: dict[str, Any]): + for key, value in d.items(): + full_key = f"{prefix}.{key}" if prefix else key + + if isinstance(value, bool) and full_key not in store_boolean_arguments: + if value: + processed_args.append("--" + full_key) + else: + processed_args.append("--" + full_key) + processed_args.append("false") + elif isinstance(value, list): + processed_args.append("--" + full_key) + for item in value: + processed_args.append(str(item)) + elif isinstance(value, dict): + process_dict(full_key, value) + else: + processed_args.append("--" + full_key) + processed_args.append(str(value)) + + process_dict("", config) + + return processed_args + + +def get_lock(model_name_or_path: str): + lock_dir = tempfile.gettempdir() + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +def warn_for_unimplemented_methods(cls: type[T]) -> type[T]: + """ + A replacement for `abc.ABC`. + When we use `abc.ABC`, subclasses will fail to instantiate + if they do not implement all abstract methods. + Here, we only require `raise NotImplementedError` in the + base class, and log a warning if the method is not implemented + in the subclass. + """ + + original_init = cls.__init__ + + def find_unimplemented_methods(self: object): + unimplemented_methods = [] + for attr_name in dir(self): + # bypass inner method + if attr_name.startswith("_"): + continue + + try: + attr = getattr(self, attr_name) + # get the func of callable method + if callable(attr): + attr_func = attr.__func__ + except AttributeError: + continue + src = inspect.getsource(attr_func) + if "NotImplementedError" in src: + unimplemented_methods.append(attr_name) + if unimplemented_methods: + method_names = ",".join(unimplemented_methods) + msg = f"Methods {method_names} not implemented in {self}" + logger.warning(msg) + + @wraps(original_init) + def wrapped_init(self, *args, **kwargs) -> None: + original_init(self, *args, **kwargs) + find_unimplemented_methods(self) + + type.__setattr__(cls, "__init__", wrapped_init) + return cls + + +def align_to(value: int, alignment: int) -> int: + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) + + +def resolve_obj_by_qualname(qualname: str) -> Any: + """ + Resolve an object by its fully qualified name. + """ + module_name, obj_name = qualname.rsplit(".", 1) + module = importlib.import_module(module_name) + return getattr(module, obj_name) + + +# From vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py +def import_pynvml(): + """ + Historical comments: + + libnvml.so is the library behind nvidia-smi, and + pynvml is a Python wrapper around it. We use it to get GPU + status without initializing CUDA context in the current process. + Historically, there are two packages that provide pynvml: + - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official + wrapper. It is a dependency of FastVideo, and is installed when users + install FastVideo. It provides a Python module named `pynvml`. + - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper. + Prior to version 12.0, it also provides a Python module `pynvml`, + and therefore conflicts with the official one which is a standalone Python file. + This causes errors when both of them are installed. + Starting from version 12.0, it migrates to a new module + named `pynvml_utils` to avoid the conflict. + It is so confusing that many packages in the community use the + unofficial one by mistake, and we have to handle this case. + For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial + one, and it will cause errors, see the issue + https://github.com/vllm-project/vllm/issues/12847 for example. + After all the troubles, we decide to copy the official `pynvml` + module to our codebase, and use it directly. + """ + import wan.third_party.pynvml as pynvml + + return pynvml + + +def maybe_download_model( + model_name_or_path: str, local_dir: str | None = None, download: bool = True +) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + # If the path exists locally, return it + if os.path.exists(model_name_or_path): + logger.info("Model already exists locally at %s", model_name_or_path) + return model_name_or_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info( + "Downloading model snapshot from HF Hub for %s...", model_name_or_path + ) + with get_lock(model_name_or_path): + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir, + ) + logger.info("Downloaded model to %s", local_path) + return str(local_path) + except Exception as e: + raise ValueError( + f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" + ) from e + + +def maybe_download_lora( + model_name_or_path: str, local_dir: str | None = None, download: bool = True +) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + local_path = maybe_download_model(model_name_or_path, local_dir, download) + weight_name = _best_guess_weight_name( + model_name_or_path, file_extension=".safetensors" + ) + return os.path.join(local_path, weight_name) + + +def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: + """ + Verify that the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary + """ + + # Check for model_index.json which is required for diffusers models + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + raise ValueError( + f"Model directory {model_path} does not contain model_index.json. " + "Only Hugging Face diffusers format is supported." + ) + + # Check for transformer and vae directories + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + + if not os.path.exists(transformer_dir): + raise ValueError( + f"Model directory {model_path} does not contain a transformer/ directory." + ) + + if not os.path.exists(vae_dir): + raise ValueError( + f"Model directory {model_path} does not contain a vae/ directory." + ) + + # Load the config + with open(config_path) as f: + config = json.load(f) + + # Verify diffusers version exists + if "_diffusers_version" not in config: + raise ValueError("model_index.json does not contain _diffusers_version") + + logger.info("Diffusers version: %s", config["_diffusers_version"]) + return cast(dict[str, Any], config) + + +def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: + """ + Download and extract just the model_index.json for a Hugging Face model. + + Args: + model_name_or_path: Path or HF Hub model ID + + Returns: + The parsed model_index.json as a dictionary + """ + import tempfile + + from huggingface_hub import hf_hub_download + + # If it's a local path, verify it directly + if os.path.exists(model_name_or_path): + return verify_model_config_and_directory(model_name_or_path) + + # For remote models, download just the model_index.json + try: + with tempfile.TemporaryDirectory() as tmp_dir: + # Download just the model_index.json file + model_index_path = hf_hub_download( + repo_id=model_name_or_path, + filename="model_index.json", + local_dir=tmp_dir, + ) + + # Load the model_index.json + with open(model_index_path) as f: + config: dict[str, Any] = json.load(f) + + # Verify it has the required fields + if "_class_name" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _class_name field" + ) + + if "_diffusers_version" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _diffusers_version field" + ) + + # Add the pipeline name for downstream use + config["pipeline_name"] = config["_class_name"] + + logger.info( + "Downloaded model_index.json for %s, pipeline: %s", + model_name_or_path, + config["_class_name"], + ) + return config + + except Exception as e: + raise ValueError( + f"Failed to download or parse model_index.json for {model_name_or_path}: {e}" + ) from e + + +def update_environment_variables(envs: dict[str, str]): + for k, v in envs.items(): + if k in os.environ and os.environ[k] != v: + logger.warning( + "Overwriting environment variable %s " "from '%s' to '%s'", + k, + os.environ[k], + v, + ) + os.environ[k] = v + + +def run_method( + obj: Any, method: str | bytes | Callable, args: tuple[Any], kwargs: dict[str, Any] +) -> Any: + """ + Run a method of an object with the given arguments and keyword arguments. + If the method is string, it will be converted to a method using getattr. + If the method is serialized bytes and will be deserialized using + cloudpickle. + If the method is a callable, it will be called directly. + """ + if isinstance(method, bytes): + func = partial(cloudpickle.loads(method), obj) + elif isinstance(method, str): + try: + func = getattr(obj, method) + except AttributeError: + raise NotImplementedError( + f"Method {method!r} is not" " implemented." + ) from None + else: + func = partial(method, obj) # type: ignore + return func(*args, **kwargs) + + +def shallow_asdict(obj) -> dict[str, Any]: + if not is_dataclass(obj): + raise TypeError("Expected dataclass instance") + return {f.name: getattr(obj, f.name) for f in fields(obj)} + + +# TODO: validate that this is fine +def kill_itself_when_parent_died() -> None: + # if sys.platform == "linux": + # sigkill this process when parent worker manager dies + PR_SET_PDEATHSIG = 1 + import platform + + if platform.system() == "Linux": + libc = ctypes.CDLL("libc.so.6") + libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) + # elif platform.system() == "Darwin": + # libc = ctypes.CDLL("libc.dylib") + # logger.warning("kill_itself_when_parent_died is only supported in linux.") + else: + logger.warning("kill_itself_when_parent_died is only supported in linux.") + + +def get_exception_traceback() -> str: + etype, value, tb = sys.exc_info() + err_str = "".join(traceback.format_exception(etype, value, tb)) + return err_str + + +class TypeBasedDispatcher: + + def __init__(self, mapping: list[tuple[type, Callable]]): + self._mapping = mapping + + def __call__(self, obj: Any): + for ty, fn in self._mapping: + if isinstance(obj, ty): + return fn(obj) + raise ValueError(f"Invalid object: {obj}") + + +# For non-torch.distributed debugging +def remote_breakpoint() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("localhost", 0)) # Let the OS pick an ephemeral port. + port = s.getsockname()[1] + RemotePdb(host="localhost", port=port).set_trace() + + +@dataclass +class MixedPrecisionState: + param_dtype: torch.dtype | None = None + reduce_dtype: torch.dtype | None = None + output_dtype: torch.dtype | None = None + compute_dtype: torch.dtype | None = None + mp_policy: MixedPrecisionPolicy | None = None + + +# Thread-local storage for mixed precision state +_mixed_precision_state = threading.local() + + +def get_mixed_precision_state() -> MixedPrecisionState: + """Get the current mixed precision state.""" + if not hasattr(_mixed_precision_state, "state"): + raise ValueError("Mixed precision state not set") + return cast(MixedPrecisionState, _mixed_precision_state.state) + + +def set_mixed_precision_policy( + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + output_dtype: torch.dtype | None = None, + mp_policy: MixedPrecisionPolicy | None = None, +): + """Set mixed precision policy globally. + + Args: + param_dtype: Parameter dtype used for training + reduce_dtype: Reduction dtype used for gradients + output_dtype: Optional output dtype + """ + state = MixedPrecisionState( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + _mixed_precision_state.state = state + + +def get_compute_dtype() -> torch.dtype: + """Get the current compute dtype from mixed precision policy. + + Returns: + torch.dtype: The compute dtype to use, defaults to get_default_dtype() if no policy set + """ + if not hasattr(_mixed_precision_state, "state"): + return torch.get_default_dtype() + else: + state = get_mixed_precision_state() + return state.param_dtype + + +def dict_to_3d_list( + mask_strategy: dict[str, Any] | None = None, + t_max: int | None = None, + l_max: int | None = None, + h_max: int | None = None, +) -> list[list[list[torch.Tensor | None]]]: + """ + Convert a dictionary of mask indices to a 3D list of tensors. + Args: + mask_strategy: keys are "t_l_h", values are torch.Tensor masks. + t_max, l_max, h_max: if provided (all three), force the output shape to (t_max, l_max, h_max). + If all three are None, infer shape from the data. + """ + # Case 1: no data, but fixed shape requested + if mask_strategy is None: + assert ( + t_max is not None and l_max is not None and h_max is not None + ), "If mask_strategy is None, you must provide t_max, l_max, and h_max" + return [ + [[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max) + ] + + # Parse all keys into integer tuples + indices = [tuple(map(int, key.split("_"))) for key in mask_strategy] + + # Decide on dimensions + if t_max is None and l_max is None and h_max is None: + # fully dynamic: infer from data + max_timesteps_idx = max(t for t, _, _ in indices) + 1 + max_layer_idx = max(l for _, l, _ in indices) + 1 # noqa: E741 + max_head_idx = max(h for _, _, h in indices) + 1 + else: + # require all three to be provided + assert t_max is not None and l_max is not None and h_max is not None, ( + "Either supply none of (t_max, l_max, h_max) to infer dimensions, " + "or supply all three to fix the shape." + ) + max_timesteps_idx = t_max + max_layer_idx = l_max + max_head_idx = h_max + + # Preallocate + result = [ + [[None for _ in range(max_head_idx)] for _ in range(max_layer_idx)] + for _ in range(max_timesteps_idx) + ] + + # Fill in, skipping any out-of-bounds entries + for key, value in mask_strategy.items(): + t, l, h = map(int, key.split("_")) # noqa: E741 + if ( + 0 <= t < max_timesteps_idx + and 0 <= l < max_layer_idx + and 0 <= h < max_head_idx + ): + result[t][l][h] = value + # else: silently ignore any key that doesn't fit + + return result + + +def set_random_seed(seed: int) -> None: + from platforms import current_platform + + current_platform.seed_everything(seed) + + +@lru_cache(maxsize=1) +def is_vsa_available() -> bool: + return importlib.util.find_spec("vsa") is not None diff --git a/worldcompass/LICENSE b/worldcompass/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a7b8b5e6a6dc555ad43165aa5b8f65c120032468 --- /dev/null +++ b/worldcompass/LICENSE @@ -0,0 +1,187 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. \ No newline at end of file diff --git a/worldcompass/README.md b/worldcompass/README.md new file mode 100644 index 0000000000000000000000000000000000000000..fc8ed1567390c69c4b0c206204942fed31b18de8 --- /dev/null +++ b/worldcompass/README.md @@ -0,0 +1,339 @@ +
+ + +

🧭 WorldCompass: Reinforcement Learning for Long-Horizon World Models

+ +[English](README.md) | [简体中文](README_zh.md) + + + +
+ +
+ + + + + + + +
+ +[//]: # ( ) + +[//]: # ( ) + +[//]: # ( PyPI - Downloads) + + +## Introduction to WorldCompass + +WorldCompass uses Group Relative Policy Optimization (GRPO) idea to improve +action-following and visual quality in autoregressive video generation. + +[Research Paper](https://arxiv.org/abs/2602.09022) + +## Training Guidence + +The training pipeline includes: + +1. **Environment Setup** - Install dependencies +2. **Model Download** - Download pretrained checkpoints +3. **Dataset Preparation** - Prepare image-text latents and action trajectories +4. **Training** - Run reinforcement learning training + +--- + +### Step 1: Environment Setup + +--- + +**1.1 Create Conda Environment** + +```bash +# Create and activate environment +conda create -n worldcompass python=3.10 -y +conda activate worldcompass + +# Install dependencies +pip install -r requirements.txt +pip install transformers==4.50.0 + +# Install Flash Attention (recommended for faster training) +pip install flash-attn==2.7.3 --no-build-isolation +``` + +--- + +### Step 2: Model Preparation + +--- + +**2.1 Download Model Checkpoints** + +```bash +# Download all required models +python download_models_worldcompass.py --hf_token --cache_dir +``` + +**Note**: You need access to FLUX.1-Redux-dev for the vision encoder: + +1. Request access: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev +2. Create token: https://huggingface.co/settings/tokens (select "Read" + permission) + +The script will download: + +- **HunyuanVideo-1.5** base model (VAE, scheduler, 480p transformer) +- **HY-WorldPlay** action models (AR model, bidirectional model, distilled + model) +- **Qwen2.5-VL-7B-Instruct** text encoder +- **ByT5** encoders (byt5-small + Glyph-SDXL-v2) +- **SigLIP** vision encoder (from FLUX.1-Redux-dev) +- **DepthAnythingV3** camera pose estimation model - Option 1 +- **Hunyuan-WorldMirror** camera pose estimation model - Option 2 + +--- + +**2.2 Clone DepthAnythingV3 Repository (Optional, Using WorldMirror as Reward by default)** + +```bash +git clone https://github.com/ByteDance-Seed/Depth-Anything-3.git DepthAnythingV3 +mv ./DepthAnythingV3/src/depth_anything_3 ./depth_anything_3 +``` + +--- + +### Step 3: Dataset Preparation + +--- + +**3.1 Prepare Input Data** + +Create a JSON file containing your training data: + +```json +[ + { + "image_path": "/path/to/image1.jpg", + "caption": "A serene park with trees and a bridge over water" + }, + { + "image_path": "/path/to/image2.png", + "caption": "A modern city street at sunset" + } +] +``` + +--- + +**3.2 Extract Image-Text Latents** + +This step encodes images and text into latent features using VAE and text +encoders. + +**Single GPU**: + +```bash +python prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/train.json \ + --output_dir /path/to/train_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +**Multi-GPU (Recommended)**: + +```bash +torchrun --nproc_per_node=8 prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/train.json \ + --output_dir /path/to/train_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +**Output Structure**: + +``` +/path/to/train_latents/ +├── latents/ +│ ├── 0_000000.pt +│ ├── 0_000001.pt +│ └── ... +└── latents.json # Index file for training +``` + +Each `.pt` file contains: + +- `latent`: VAE-encoded image features [1, C, T, H, W] +- `image_cond`: First frame condition [1, C, 1, H, W] +- `prompt_embeds`: Text embeddings [1, L, D] +- `prompt_mask`: Text attention mask [1, L] +- `vision_states`: Visual features [1, N, D] +- `byt5_text_states`: ByT5 text features [1, 256, 1472] +- `byt5_text_mask`: ByT5 attention mask [1, 256] + +--- + +**3.3 Prepare Evaluation Data** + +To support observing model performance on a fixed subset during training, we +recommend building a mini evaluation dataset. The JSON file format is the same +as the training data. The recommended number should be divisible by the expected +number of GPUs, and for efficiency, typically 16 or 32 samples are sufficient. + +```bash +python prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/eval.json \ + --output_dir /path/to/eval_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +--- + +**3.4 Generate Action Trajectories** + +Generate random camera trajectory sequences for training: + +```bash +python prepare_dataset/prepare_custom_action.py +``` + +**Default Output**: `prepare_dataset/harder_random_poses.json` (1000 +trajectories, 128 actions each, default settings favor generating complex +composite actions) + +**Customize**: Edit `prepare_custom_action.py` to adjust number of trajectories +and actions per trajectory, as well as trajectory template synthesis rules. + +--- + +### Step 4: Start Training + +--- + +**4.1 Configure Training Script** + +Modify the TODO configs in `scripts/train_worldcompass.sh`: + +```bash + +export MASTER_ADDR="" # [TODO] Set master address + +# NOTE: Run download_models.py (Step 2) can get the exact paths for the variables below. +CACHE_DIR="" # [TODO] Fill in the overall checkpoint cache directory (as printed by download_models.py) +HUNYUAN_CHECKPOINT="" # [TODO] Fill in the HunYuan checkpoint directory +WORLDPLAY_CHECKPOINT="" # [TODO] Fill in the WorldPlay checkpoint directory + +TRAIN_LATENTS_DIR="" # [TODO] Path to the train latents directory +EVAL_LATENTS_DIR="" # [TODO] Path to the eval latents directory +POSE_PATH="" # [TODO] Path to the custom action pose json +OUTPUT_DIR="" # [TODO] Path to the output directory + +--wandb_key "" # [TODO] Set wandb key +--wandb_entity "" # [TODO] Set wandb entity +``` + +--- + +**4.2 Training!!!** + +Training on 8 GPUs can show initial improvements, but more nodes typically +result in more stable training and better final results. + +**Single Node (8 GPUs)**: + +```bash +bash scripts/train_worldcompass.sh 0 1 +``` + +**Multi-Node Training**: + +```bash +# On each node, run with appropriate rank +# Node 0 (master): +bash scripts/train_worldcompass.sh 0 4 # 4 nodes total + +# Node 1: +bash scripts/train_worldcompass.sh 1 4 + +# Node 2: +bash scripts/train_worldcompass.sh 2 4 + +# Node 3: +bash scripts/train_worldcompass.sh 3 4 +``` + +**Note**: Ensure all nodes can communicate via `MASTER_ADDR` and `MASTER_PORT`. + +**4.3 Monitor Training** + +Training logs and checkpoints will be saved to: + +``` +${CKPT_DIR}/${exp_name}/ +├── checkpoint-{step}/ +│ ├── transformer/ +│ │ └── diffusion_pytorch_model.safetensors +│ └── training_state.pt +└── generated_videos/ # Sample videos during training +``` + +If WandB is configured, metrics will be logged to your WandB dashboard. + +### 4.4 Hardware Requirements + +Our experiments were conducted on GPUs with 96GB memory. If you encounter GPU +OOM issues, consider: + +- Reducing `window_frames` +- Converting VAE to `torch.bf16` (may cause training instability) + +## Results + +We have conducted comprehensive validation on the **WorldPlay** baseline model. Our evaluation demonstrates that after undergoing **WorldCompass** post-training, the model's capabilities undergo a qualitative leap, particularly in both interaction following and visual quality. + + +

+ +

+

+ +

+ + +## 📚 Citation + +```bibtex +@article{hyworld2025, + title={HY-World 1.5: A Systematic Framework for Interactive World Modeling with Real-Time Latency and Geometric Consistency}, + author={Team HunyuanWorld}, + journal={arXiv preprint}, + year={2025} +} + +@article{wang2026worldcompass, + title={WorldCompass: Reinforcement Learning for Long-Horizon World Models}, + author={Wang, Zehan and Wang, Tengfei and Zhang, Haiyu and Zuo, Xuhui and Wu, Junta and Wang, Haoyuan and Sun, Wenqiang and Wang, Zhenwei and Cao, Chenjie and Zhao, Hengshuang and others}, + journal={arXiv preprint}, + year={2026} +} + +@article{worldplay2025, + title={WorldPlay: Towards Long-Term Geometric Consistency for Real-Time Interactive World Model}, + author={Wenqiang Sun and Haiyu Zhang and Haoyuan Wang and Junta Wu and Zehan Wang and Zhenwei Wang and Yunhong Wang and Jun Zhang and Tengfei Wang and Chunchao Guo}, + year={2025}, + journal={arXiv preprint} +} +``` + +## Contact + +Please send emails to tengfeiwang12@gmail.com if there is any question + +## 🙏 Acknowledgements + +We would like to thank +[HunyuanWorld](https://github.com/Tencent-Hunyuan/HunyuanWorld-1.0), +[HunyuanWorld-Mirror](https://github.com/Tencent-Hunyuan/HunyuanWorld-Mirror), +[HunyuanVideo](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5), +[DiffusionNFT](https://github.com/NVlabs/DiffusionNFT) and +[FastVideo](https://github.com/hao-ai-lab/FastVideo) for their great work. diff --git a/worldcompass/README_zh.md b/worldcompass/README_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..a06e10721b76ac228fb3a7cb84b811a20e336849 --- /dev/null +++ b/worldcompass/README_zh.md @@ -0,0 +1,325 @@ +
+ + +

🧭 WorldCompass: Reinforcement Learning for Long-Horizon World Models

+ +[English](README.md) | [简体中文](README_zh.md) + + + +
+ +
+ + + + + + + +
+ +[//]: # ( ) + +[//]: # ( ) + +[//]: # ( PyPI - Downloads) + +## WorldCompass 简介 +[English](README.md) | [简体中文](README_zh.md) + +WorldCompass 借鉴了GRPO的思想，旨在提升自回归视频生成中的动作跟随能力与视觉质量。 + +[研究论文](https://arxiv.org/abs/2602.09022) + +## 训练指南 + +训练流程包括： + +1. **环境配置** - 安装依赖 +2. **模型下载** - 下载预训练检查点 +3. **数据集准备** - 准备图文潜变量及动作轨迹 +4. **训练** - 运行强化学习训练 + +--- + +### 第一步：环境配置 + +--- + +**1.1 创建 Conda 环境** + +```bash +# 创建并激活环境 +conda create -n worldcompass python=3.10 -y +conda activate worldcompass + +# 安装依赖 +pip install -r requirements.txt +pip install transformers==4.50.0 + +# 安装 Flash Attention（推荐，可加速训练） +pip install flash-attn==2.7.3 --no-build-isolation +``` + +--- + +### 第二步：模型准备 + +--- + +**2.1 下载模型检查点** + +```bash +# 下载所有所需模型 +python download_models_worldcompass.py --hf_token --cache_dir +``` + +**注意**：您需要获得 FLUX.1-Redux-dev 的访问权限以使用视觉编码器： + +1. 申请访问权限：https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev +2. 创建 Token：https://huggingface.co/settings/tokens（选择"Read"权限） + +该脚本将下载以下内容： + +- **HunyuanVideo-1.5** 基础模型（VAE、调度器、480p Transformer） +- **HY-WorldPlay** 动作模型（AR 模型、双向模型、蒸馏模型） +- **Qwen2.5-VL-7B-Instruct** 文本编码器 +- **ByT5** 编码器（byt5-small + Glyph-SDXL-v2） +- **SigLIP** 视觉编码器（来自 FLUX.1-Redux-dev） +- **DepthAnythingV3** 相机位姿估计模型 - 方案一 +- **Hunyuan-WorldMirror** 相机位姿估计模型 - 方案二 + +--- + +**2.2 克隆 DepthAnythingV3 仓库（可选，默认使用WorldMirror作为奖励）** + +```bash +git clone https://github.com/ByteDance-Seed/Depth-Anything-3.git DepthAnythingV3 +mv ./DepthAnythingV3/src/depth_anything_3 ./depth_anything_3 +``` + +--- + +### 第三步：数据集准备 + +--- + +**3.1 准备输入数据** + +创建一个包含训练数据的 JSON 文件： + +```json +[ + { + "image_path": "/path/to/image1.jpg", + "caption": "A serene park with trees and a bridge over water" + }, + { + "image_path": "/path/to/image2.png", + "caption": "A modern city street at sunset" + } +] +``` + +--- + +**3.2 提取图文潜变量** + +该步骤使用 VAE 和文本编码器将图像与文本编码为潜变量特征。 + +**单卡**： + +```bash +python prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/train.json \ + --output_dir /path/to/train_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +**多卡（推荐）**： + +```bash +torchrun --nproc_per_node=8 prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/train.json \ + --output_dir /path/to/train_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +**输出结构**： + +``` +/path/to/train_latents/ +├── latents/ +│ ├── 0_000000.pt +│ ├── 0_000001.pt +│ └── ... +└── latents.json # 训练索引文件 +``` + +每个 `.pt` 文件包含： + +- `latent`：VAE 编码的图像特征 [1, C, T, H, W] +- `image_cond`：首帧条件 [1, C, 1, H, W] +- `prompt_embeds`：文本嵌入 [1, L, D] +- `prompt_mask`：文本注意力掩码 [1, L] +- `vision_states`：视觉特征 [1, N, D] +- `byt5_text_states`：ByT5 文本特征 [1, 256, 1472] +- `byt5_text_mask`：ByT5 注意力掩码 [1, 256] + +--- + +**3.3 准备评估数据** + +为了在训练过程中在固定子集上观察模型性能，建议构建一个小型评估数据集。JSON 文件格式与训练数据相同。推荐的样本数量应能被预期的 GPU 数量整除，通常 16 或 32 个样本即可满足效率需求。 + +```bash +python prepare_dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/eval.json \ + --output_dir /path/to/eval_latents \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 +``` + +--- + +**3.4 生成动作轨迹** + +为训练生成随机相机轨迹序列： + +```bash +python prepare_dataset/prepare_custom_action.py +``` + +**默认输出**：`prepare_dataset/harder_random_poses.json`（1000 条轨迹，每条 128 个动作，默认设置倾向于生成复合动作） + +**自定义**：编辑 `prepare_custom_action.py` 以调整轨迹数量、每条轨迹的动作数量以及轨迹模板合成规则。 + +--- + +### 第四步：开始训练 + +--- + +**4.1 配置训练脚本** + +修改 `scripts/train_worldcompass.sh` 中标有 TODO 的配置项： + +```bash + +export MASTER_ADDR="" # [TODO] 设置主节点地址 + +# 注意：运行 download_models.py（第二步）可获取以下变量的精确路径。 +CACHE_DIR="" # [TODO] 填写整体检查点缓存目录（由 download_models.py 输出） +HUNYUAN_CHECKPOINT="" # [TODO] 填写 HunYuan 检查点目录 +WORLDPLAY_CHECKPOINT="" # [TODO] 填写 WorldPlay 检查点目录 + +TRAIN_LATENTS_DIR="" # [TODO] 训练潜变量目录路径 +EVAL_LATENTS_DIR="" # [TODO] 评估潜变量目录路径 +POSE_PATH="" # [TODO] 自定义动作位姿 JSON 文件路径 +OUTPUT_DIR="" # [TODO] 输出目录路径 + +--wandb_key "" # [TODO] 设置 wandb key +--wandb_entity "" # [TODO] 设置 wandb entity +``` + +--- + +**4.2 开始训练！！！** + +在 8 张 GPU 上训练可以初步看到改善效果，但使用更多节点通常能带来更稳定的训练过程和更好的最终结果。 + +**单节点（8 张 GPU）**： + +```bash +bash scripts/train_worldcompass.sh 0 1 +``` + +**多节点训练**： + +```bash +# 在每个节点上以相应的 rank 运行 +# 节点 0（主节点）： +bash scripts/train_worldcompass.sh 0 4 # 共 4 个节点 + +# 节点 1： +bash scripts/train_worldcompass.sh 1 4 + +# 节点 2： +bash scripts/train_worldcompass.sh 2 4 + +# 节点 3： +bash scripts/train_worldcompass.sh 3 4 +``` + +**注意**：请确保所有节点均可通过 `MASTER_ADDR` 和 `MASTER_PORT` 相互通信。 + +**4.3 监控训练** + +训练日志和检查点将保存至： + +``` +${CKPT_DIR}/${exp_name}/ +├── checkpoint-{step}/ +│ ├── transformer/ +│ │ └── diffusion_pytorch_model.safetensors +│ └── training_state.pt +└── generated_videos/ # 训练过程中生成的样例视频 +``` + +如果已配置 WandB，指标将记录到您的 WandB 仪表盘。 + +### 4.4 硬件要求 + +我们的实验在具有 96GB 显存的 GPU 上进行。如果遇到 GPU 显存不足（OOM）问题，可以考虑： + +- 减小 `window_frames` +- 将 VAE 转换为 `torch.bf16`（可能导致训练不稳定） + +## 结果 + +我们在WorldPlay基准模型上进行了全面的验证。评估结果表明，在经过 WorldCompass 后训练后，该模型的能力实现了明显提升，特别是在交互指令遵循和视觉质量这两个方面。 + +

+ +

+

+ +

+ +## 📚 引用 + +```bibtex +@article{hyworld2025, + title={HY-World 1.5: A Systematic Framework for Interactive World Modeling with Real-Time Latency and Geometric Consistency}, + author={Team HunyuanWorld}, + journal={arXiv preprint}, + year={2025} +} + +@article{wang2026worldcompass, + title={WorldCompass: Reinforcement Learning for Long-Horizon World Models}, + author={Wang, Zehan and Wang, Tengfei and Zhang, Haiyu and Zuo, Xuhui and Wu, Junta and Wang, Haoyuan and Sun, Wenqiang and Wang, Zhenwei and Cao, Chenjie and Zhao, Hengshuang and others}, + journal={arXiv preprint}, + year={2026} +} + +@article{worldplay2025, + title={WorldPlay: Towards Long-Term Geometric Consistency for Real-Time Interactive World Model}, + author={Wenqiang Sun and Haiyu Zhang and Haoyuan Wang and Junta Wu and Zehan Wang and Zhenwei Wang and Yunhong Wang and Jun Zhang and Tengfei Wang and Chunchao Guo}, + year={2025}, + journal={arXiv preprint} +} +``` + +## 联系方式 + +如有任何问题，请发送邮件至 tengfeiwang12@gmail.com + +## 🙏 致谢 + +我们衷心感谢 +[HunyuanWorld](https://github.com/Tencent-Hunyuan/HunyuanWorld-1.0)、 +[HunyuanWorld-Mirror](https://github.com/Tencent-Hunyuan/HunyuanWorld-Mirror)、 +[HunyuanVideo](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)、 +[DiffusionNFT](https://github.com/NVlabs/DiffusionNFT) 以及 +[FastVideo](https://github.com/hao-ai-lab/FastVideo) 的出色工作。 diff --git a/worldcompass/assets/image.png b/worldcompass/assets/image.png new file mode 100644 index 0000000000000000000000000000000000000000..f79e94721bbdd8f1c0aa52760647b56925df0b03 --- /dev/null +++ b/worldcompass/assets/image.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:71a6c33ddb81e6f17c23c250cc295ccd2b185694406064cc4a184949c0eeb7ce +size 243430 diff --git a/worldcompass/assets/results.png b/worldcompass/assets/results.png new file mode 100644 index 0000000000000000000000000000000000000000..96bf6bc902429287d146f1fd9f691950602115e1 Binary files /dev/null and b/worldcompass/assets/results.png differ diff --git a/worldcompass/assets/teaser.webp b/worldcompass/assets/teaser.webp new file mode 100644 index 0000000000000000000000000000000000000000..8143f8ccc8600827be089b16daec8eb3e886b5e1 --- /dev/null +++ b/worldcompass/assets/teaser.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c4c21c953b31c829aaf3c28fa719bac4a9303bb88e06d15aa09a72f19604a1c6 +size 501236 diff --git a/worldcompass/download_models_worldcompass.py b/worldcompass/download_models_worldcompass.py new file mode 100644 index 0000000000000000000000000000000000000000..f2b1629376c7b141c9a154526f61df393d15842c --- /dev/null +++ b/worldcompass/download_models_worldcompass.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +"""Download script for HY-WorldPlay models. Downloads all required models from HuggingFace and +ModelScope. + +Usage: + python download_models.py --hf_token + +The HF token is required for downloading the vision encoder from FLUX.1-Redux-dev. +Request access at: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev +""" + +import argparse +import os +import shutil +import sys + + +def check_dependencies(): + """Check and install required dependencies.""" + try: + from huggingface_hub import snapshot_download + except ImportError: + print("Installing huggingface_hub...") + os.system("pip install -U 'huggingface_hub[cli]'") + + try: + import modelscope + except ImportError: + print("Installing modelscope...") + os.system("pip install modelscope") + + +def download_hy_worldplay(cache_dir=None): + """Download only ar_model/diffusion_pytorch_model.safetensors from tencent/HY-WorldPlay.""" + from huggingface_hub import snapshot_download + + print("\n" + "=" * 60) + print( + "[1/8] Downloading ar_model/diffusion_pytorch_model.safetensors from tencent/HY-WorldPlay..." + ) + print("=" * 60) + + # Only download the ar_model/diffusion_pytorch_model.safetensors file + worldplay_path = snapshot_download( + "tencent/HY-WorldPlay", + allow_patterns=["ar_model/diffusion_pytorch_model.safetensors"], + cache_dir=cache_dir, + ) + model_path = os.path.join( + worldplay_path, "ar_model", "diffusion_pytorch_model.safetensors" + ) + # os.makedirs(target_dir, exist_ok=True) + # target_path = os.path.join(target_dir, "diffusion_pytorch_model.safetensors") + # if os.path.abspath(model_path) != os.path.abspath(target_path): + # shutil.move(model_path, target_path) + # model_path = target_path + + print(f"Downloaded file: {model_path}") + + return model_path + + +def download_hunyuan_video(cache_dir=None): + """Download HunyuanVideo-1.5 base models (vae, scheduler, transformer).""" + from huggingface_hub import snapshot_download + + print("\n" + "=" * 60) + print( + "[2/8] Downloading tencent/HunyuanVideo-1.5 (vae, scheduler, transformer)..." + ) + print("=" * 60) + + hunyuan_path = snapshot_download( + "tencent/HunyuanVideo-1.5", + allow_patterns=["vae/*", "scheduler/*", "transformer/480p_i2v/*"], + cache_dir=cache_dir, + ) + print(f"Downloaded to: {hunyuan_path}") + return hunyuan_path + + +def download_llm_text_encoder(hunyuan_path, cache_dir=None): + """Download Qwen2.5-VL-7B-Instruct as the LLM text encoder.""" + from huggingface_hub import snapshot_download + + print("\n" + "=" * 60) + print("[3/8] Downloading LLM text encoder (Qwen2.5-VL-7B-Instruct)...") + print("=" * 60) + + text_encoder_base = os.path.join(hunyuan_path, "text_encoder") + os.makedirs(text_encoder_base, exist_ok=True) + + llm_target = os.path.join(text_encoder_base, "llm") + + if ( + os.path.exists(llm_target) + and os.path.isdir(llm_target) + and len(os.listdir(llm_target)) > 5 + ): + print(f"LLM text encoder already exists at: {llm_target}") + return + + # Clean up old/broken downloads + if os.path.islink(llm_target): + os.unlink(llm_target) + elif os.path.exists(llm_target): + shutil.rmtree(llm_target) + + print("Downloading Qwen/Qwen2.5-VL-7B-Instruct (~15GB)...") + qwen_cache = snapshot_download( + "Qwen/Qwen2.5-VL-7B-Instruct", cache_dir=cache_dir + ) + + # Copy files (resolve symlinks) + os.makedirs(llm_target, exist_ok=True) + for item in os.listdir(qwen_cache): + src = os.path.realpath(os.path.join(qwen_cache, item)) + dst = os.path.join(llm_target, item) + if os.path.isdir(src): + shutil.copytree(src, dst, dirs_exist_ok=True) + else: + shutil.copy2(src, dst) + + print(f"Copied to: {llm_target}") + + +def download_byt5_encoders(hunyuan_path, cache_dir=None): + """Download ByT5 text encoders (byt5-small and Glyph-SDXL-v2).""" + from huggingface_hub import snapshot_download + from modelscope import snapshot_download as ms_snapshot_download + + print("\n" + "=" * 60) + print("[4/8] Downloading ByT5 text encoders...") + print("=" * 60) + + text_encoder_base = os.path.join(hunyuan_path, "text_encoder") + os.makedirs(text_encoder_base, exist_ok=True) + + # 1. Download google/byt5-small + byt5_target = os.path.join(text_encoder_base, "byt5-small") + if ( + os.path.exists(byt5_target) + and os.path.isdir(byt5_target) + and len(os.listdir(byt5_target)) > 3 + ): + print(f"byt5-small already exists at: {byt5_target}") + else: + if os.path.islink(byt5_target): + os.unlink(byt5_target) + elif os.path.exists(byt5_target): + shutil.rmtree(byt5_target) + + print("Downloading google/byt5-small...") + byt5_cache = snapshot_download("google/byt5-small", cache_dir=cache_dir) + + os.makedirs(byt5_target, exist_ok=True) + for item in os.listdir(byt5_cache): + src = os.path.realpath(os.path.join(byt5_cache, item)) + dst = os.path.join(byt5_target, item) + if os.path.isdir(src): + shutil.copytree(src, dst, dirs_exist_ok=True) + else: + shutil.copy2(src, dst) + print(f"Copied to: {byt5_target}") + + # 2. Download Glyph-SDXL-v2 from ModelScope + glyph_target = os.path.join(text_encoder_base, "Glyph-SDXL-v2") + if os.path.exists(glyph_target) and os.path.exists( + os.path.join(glyph_target, "checkpoints", "byt5_model.pt") + ): + print(f"Glyph-SDXL-v2 already exists at: {glyph_target}") + else: + if os.path.exists(glyph_target): + shutil.rmtree(glyph_target) + + print("Downloading AI-ModelScope/Glyph-SDXL-v2 from ModelScope...") + # Use custom cache_dir if provided, otherwise use /tmp/glyph_cache + glyph_cache_dir = cache_dir if cache_dir else "/tmp/glyph_cache" + glyph_cache = ms_snapshot_download( + "AI-ModelScope/Glyph-SDXL-v2", cache_dir=glyph_cache_dir + ) + + os.makedirs(glyph_target, exist_ok=True) + for item in os.listdir(glyph_cache): + src = os.path.join(glyph_cache, item) + dst = os.path.join(glyph_target, item) + if os.path.isdir(src): + if os.path.exists(dst): + shutil.rmtree(dst) + shutil.copytree(src, dst) + else: + shutil.copy2(src, dst) + print(f"Copied to: {glyph_target}") + + +def download_vision_encoder(hunyuan_path, hf_token, cache_dir=None): + """Download SigLIP vision encoder from FLUX.1-Redux-dev.""" + from huggingface_hub import snapshot_download + + print("\n" + "=" * 60) + print("[5/8] Downloading Vision Encoder (SigLIP from FLUX.1-Redux-dev)...") + print("=" * 60) + + if not hf_token: + print("WARNING: No HF token provided!") + print( + "The vision encoder requires access to: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev" + ) + print("Skipping vision encoder download.") + print("\nYou can download it manually later.") + return + + vision_encoder_base = os.path.join(hunyuan_path, "vision_encoder") + os.makedirs(vision_encoder_base, exist_ok=True) + + siglip_target = os.path.join(vision_encoder_base, "siglip") + + if ( + os.path.exists(siglip_target) + and os.path.isdir(siglip_target) + and len(os.listdir(siglip_target)) > 3 + ): + print(f"siglip already exists at: {siglip_target}") + return + + # Clean up old/broken downloads + if os.path.islink(siglip_target): + os.unlink(siglip_target) + elif os.path.exists(siglip_target): + shutil.rmtree(siglip_target) + + print("Downloading black-forest-labs/FLUX.1-Redux-dev...") + try: + flux_cache = snapshot_download( + "black-forest-labs/FLUX.1-Redux-dev", + token=hf_token, + cache_dir=cache_dir, + ) + + # Copy files (resolve symlinks) + os.makedirs(siglip_target, exist_ok=True) + for item in os.listdir(flux_cache): + src = os.path.realpath(os.path.join(flux_cache, item)) + dst = os.path.join(siglip_target, item) + if os.path.isdir(src): + shutil.copytree(src, dst, dirs_exist_ok=True) + else: + shutil.copy2(src, dst) + print(f"Copied to: {siglip_target}") + except Exception as e: + print(f"ERROR: Failed to download vision encoder: {e}") + print( + "Make sure you have requested access to FLUX.1-Redux-dev and your token is valid." + ) + + +def download_worldmirror(cache_dir=None): + """Download HunyuanWorld-Mirror for camera pose estimation.""" + print("\n" + "=" * 60) + print("[6/8] Downloading WorldMirror (tencent/HunyuanWorld-Mirror)...") + print("=" * 60) + + try: + from huggingface_hub import snapshot_download + + print("Downloading tencent/HunyuanWorld-Mirror...") + worldmirror_path = snapshot_download( + "tencent/HunyuanWorld-Mirror", cache_dir=cache_dir + ) + print(f"Downloaded to: {worldmirror_path}") + + # Try to load with WorldMirror if available + try: + from reward_function.HunyuanWorldMirror import WorldMirror + + worldmirror_model = WorldMirror.from_pretrained(worldmirror_path) + print("✓ WorldMirror model loaded successfully") + except ImportError: + print( + "✓ WorldMirror model downloaded (module will load on first use)" + ) + + return worldmirror_path + except ImportError: + print("WARNING: WorldMirror module not found.") + print( + "Skipping WorldMirror download. The model will be downloaded on first use." + ) + return None + except Exception as e: + print(f"WARNING: Failed to download WorldMirror: {e}") + print("The model will be downloaded on first use.") + return None + + +def download_depth_anything_3(cache_dir=None): + """Download DepthAnything3 model for camera pose estimation (optional).""" + print("\n" + "=" * 60) + print("[7/8] Downloading DepthAnything3 (depth-anything/DA3-GIANT-1.1)...") + print("=" * 60) + + from huggingface_hub import snapshot_download + + print("Downloading depth-anything/DA3-GIANT-1.1...") + da3_path = snapshot_download( + "depth-anything/DA3-GIANT-1.1", cache_dir=cache_dir + ) + print(f"Downloaded to: {da3_path}") + + return None + + +def print_paths(cache_dir=None): + """Print the model paths for run.sh configuration.""" + from huggingface_hub import snapshot_download + + print("\n" + "=" * 60) + print("[8/8] Verifying downloads...") + print("=" * 60) + + hunyuan_path = snapshot_download( + "tencent/HunyuanVideo-1.5", local_files_only=True, cache_dir=cache_dir + ) + worldplay_path = snapshot_download( + "tencent/HY-WorldPlay", local_files_only=True, cache_dir=cache_dir + ) + + # Try to get WorldMirror path + worldmirror_path = snapshot_download( + "tencent/HunyuanWorld-Mirror", + local_files_only=True, + cache_dir=cache_dir, + ) + da3_path = snapshot_download( + "depth-anything/DA3-GIANT-1.1", + local_files_only=True, + cache_dir=cache_dir, + ) + + print("\n" + "=" * 60) + print("ALL DOWNLOADS COMPLETE!") + print("=" * 60) + print( + "\nADD these paths to your prepare_dataset/prepare_image_text_latent_simple.py:\n" + ) + print(f"--hunyuan_checkpoint_path {hunyuan_path}") + + print("\nModify these paths in your scripts/full_hy_nft.sh:\n") + print(f"CACHE_DIR={cache_dir}") + print(f"HUNYUAN_CHECKPOINT={hunyuan_path}") + print(f"WORLDPLAY_CHECKPOINT={worldplay_path}") + + print( + "\nYou can now run: bash prepare_dataset/extract_latents.sh to prepare your dataset" + ) + print("\nAnd then run: bash scripts/full_hy_nft.sh to start training") + + +def main(): + parser = argparse.ArgumentParser( + description="Download all required models for HY-WorldPlay", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Download to default cache directory (~/.cache/huggingface/hub/) + python download_models.py --hf_token hf_xxxxxxxxxxxxx + + # Download to custom directory + python download_models.py --hf_token hf_xxxxxxxxxxxxx --cache_dir /mnt/data/models + +Note: + The HuggingFace token is required for downloading the vision encoder + from black-forest-labs/FLUX.1-Redux-dev. You need to: + 1. Request access at: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev + 2. Wait for approval (usually instant) + 3. Create a token at: https://huggingface.co/settings/tokens (select "Read" permission) + """, + ) + parser.add_argument( + "--hf_token", + type=str, + default=None, + help="HuggingFace token for downloading gated models (required for vision encoder)", + ) + parser.add_argument( + "--skip_vision_encoder", + action="store_true", + help="Skip downloading the vision encoder (if you don't have FLUX access yet)", + ) + parser.add_argument( + "--skip_worldmirror", + action="store_true", + help="Skip downloading WorldMirror model (optional, used for camera pose estimation)", + ) + parser.add_argument( + "--skip_depth_anything_3", + action="store_true", + help="Skip downloading DepthAnything3 model (optional, alternative to WorldMirror)", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="Directory to cache downloaded models (default: ~/.cache/huggingface/hub/)", + ) + + args = parser.parse_args() + + # Validate and prepare cache_dir + cache_dir = None + if args.cache_dir: + cache_dir = os.path.abspath(os.path.expanduser(args.cache_dir)) + os.makedirs(cache_dir, exist_ok=True) + print(f"\n{'=' * 60}") + print(f"Using custom cache directory: {cache_dir}") + print(f"{'=' * 60}\n") + + print("=" * 60) + print("HY-WorldPlay Model Download Script") + print("=" * 60) + + # Check dependencies + check_dependencies() + + # Download models + worldplay_path = download_hy_worldplay(cache_dir=cache_dir) + hunyuan_path = download_hunyuan_video(cache_dir=cache_dir) + download_llm_text_encoder(hunyuan_path, cache_dir=cache_dir) + download_byt5_encoders(hunyuan_path, cache_dir=cache_dir) + + if not args.skip_vision_encoder: + download_vision_encoder( + hunyuan_path, args.hf_token, cache_dir=cache_dir + ) + else: + print( + "\n[5/8] Skipping vision encoder download (--skip_vision_encoder flag)" + ) + + download_worldmirror(cache_dir=cache_dir) + download_depth_anything_3(cache_dir=cache_dir) + + # Print final paths + print_paths(cache_dir=cache_dir) + + +if __name__ == "__main__": + main() diff --git a/worldcompass/fastvideo/STA_configuration.py b/worldcompass/fastvideo/STA_configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..0f6d130091a048d1f99b1f206caac6d17adc966e --- /dev/null +++ b/worldcompass/fastvideo/STA_configuration.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from collections import defaultdict +from typing import Any + +import numpy as np + +from fastvideo.utils import dict_to_3d_list + + +def configure_sta( + mode: str = "STA_searching", + layer_num: int = 40, + time_step_num: int = 50, + head_num: int = 40, + **kwargs, +) -> list[list[list[Any]]]: + """Configure Sliding Tile Attention (STA) parameters based on the specified mode. + + Parameters: + ---------- + mode : str + The STA mode to use. Options are: + - 'STA_searching': Generate a set of mask candidates for initial search + - 'STA_tuning': Select best mask strategy based on previously saved results + - 'STA_inference': Load and use a previously tuned mask strategy + layer_num: int, number of layers + time_step_num: int, number of timesteps + head_num: int, number of heads + + **kwargs : dict + Mode-specific parameters: + + For 'STA_searching': + - mask_candidates: list of str, optional, mask candidates to use + - mask_selected: list of int, optional, indices of selected masks + + For 'STA_tuning': + - mask_search_files_path: str, required, path to mask search results + - mask_candidates: list of str, optional, mask candidates to use + - mask_selected: list of int, optional, indices of selected masks + - skip_time_steps: int, optional, number of time steps to use full attention (default 12) + - save_dir: str, optional, directory to save mask strategy (default "mask_candidates") + + For 'STA_inference': + - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json") + """ + valid_modes = [ + "STA_searching", + "STA_tuning", + "STA_inference", + "STA_tuning_cfg", + ] + if mode not in valid_modes: + raise ValueError(f"Mode must be one of {valid_modes}, got {mode}") + + if mode == "STA_searching": + # Get parameters with defaults + mask_candidates: list[str] | None = kwargs.get("mask_candidates") + if mask_candidates is None: + raise ValueError( + "mask_candidates is required for STA_searching mode" + ) + mask_selected: list[int] = kwargs.get( + "mask_selected", list(range(len(mask_candidates))) + ) + + # Parse selected masks + selected_masks: list[list[int]] = [] + for index in mask_selected: + mask = mask_candidates[index] + masks_list = [int(x) for x in mask.split(",")] + selected_masks.append(masks_list) + + # Create 3D mask structure with fixed dimensions (t=50, l=60) + masks_3d: list[list[list[list[int]]]] = [] + for i in range(time_step_num): # Fixed t dimension = 50 + row = [] + for j in range(layer_num): # Fixed l dimension = 60 + row.append(selected_masks) # Add all masks at each position + masks_3d.append(row) + + return masks_3d + + elif mode == "STA_tuning": + # Get required parameters + mask_search_files_path: str | None = kwargs.get( + "mask_search_files_path" + ) + if not mask_search_files_path: + raise ValueError( + "mask_search_files_path is required for STA_tuning mode" + ) + + # Get optional parameters with defaults + mask_candidates_tuning: list[str] | None = kwargs.get("mask_candidates") + if mask_candidates_tuning is None: + raise ValueError("mask_candidates is required for STA_tuning mode") + mask_selected_tuning: list[int] = kwargs.get( + "mask_selected", list(range(len(mask_candidates_tuning))) + ) + skip_time_steps_tuning: int | None = kwargs.get("skip_time_steps") + save_dir_tuning: str | None = kwargs.get("save_dir", "mask_candidates") + + # Parse selected masks + selected_masks_tuning: list[list[int]] = [] + for index in mask_selected_tuning: + mask = mask_candidates_tuning[index] + masks_list = [int(x) for x in mask.split(",")] + selected_masks_tuning.append(masks_list) + + # Read JSON results + results = read_specific_json_files(mask_search_files_path) + averaged_results = average_head_losses(results, selected_masks_tuning) + + # Add full attention mask for specific cases + full_attention_mask_tuning: list[int] | None = kwargs.get( + "full_attention_mask" + ) + if full_attention_mask_tuning is not None: + selected_masks_tuning.append(full_attention_mask_tuning) + + # Select best mask strategy + timesteps_tuning: int = kwargs.get("timesteps", time_step_num) + if skip_time_steps_tuning is None: + skip_time_steps_tuning = 12 + mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( + averaged_results, + selected_masks_tuning, + skip_time_steps_tuning, + timesteps_tuning, + head_num, + ) + + # Save mask strategy + if save_dir_tuning is not None: + os.makedirs(save_dir_tuning, exist_ok=True) + file_path = os.path.join( + save_dir_tuning, f"mask_strategy_s{skip_time_steps_tuning}.json" + ) + with open(file_path, "w") as f: + json.dump(mask_strategy, f, indent=4) + print(f"Successfully saved mask_strategy to {file_path}") + + # Print sparsity and strategy counts for information + print(f"Overall sparsity: {sparsity:.4f}") + print("\nStrategy usage counts:") + total_heads = time_step_num * layer_num * head_num # Fixed dimensions + for strategy, count in strategy_counts.items(): + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list( + mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num + ) + + return mask_strategy_3d + elif mode == "STA_tuning_cfg": + # Get required parameters for both positive and negative paths + mask_search_files_path_pos: str | None = kwargs.get( + "mask_search_files_path_pos" + ) + mask_search_files_path_neg: str | None = kwargs.get( + "mask_search_files_path_neg" + ) + save_dir_cfg: str | None = kwargs.get("save_dir") + + if ( + not mask_search_files_path_pos + or not mask_search_files_path_neg + or not save_dir_cfg + ): + raise ValueError( + "mask_search_files_path_pos, mask_search_files_path_neg," + " and save_dir are required for STA_tuning_cfg mode" + ) + + # Get optional parameters with defaults + mask_candidates_cfg: list[str] | None = kwargs.get("mask_candidates") + if mask_candidates_cfg is None: + raise ValueError( + "mask_candidates is required for STA_tuning_cfg mode" + ) + mask_selected_cfg: list[int] = kwargs.get( + "mask_selected", list(range(len(mask_candidates_cfg))) + ) + skip_time_steps_cfg: int | None = kwargs.get("skip_time_steps") + + # Parse selected masks + selected_masks_cfg: list[list[int]] = [] + for index in mask_selected_cfg: + mask = mask_candidates_cfg[index] + masks_list = [int(x) for x in mask.split(",")] + selected_masks_cfg.append(masks_list) + + # Read JSON results for both positive and negative paths + pos_results = read_specific_json_files(mask_search_files_path_pos) + neg_results = read_specific_json_files(mask_search_files_path_neg) + # Combine positive and negative results into one list + combined_results = pos_results + neg_results + + # Average the combined results + averaged_results = average_head_losses( + combined_results, selected_masks_cfg + ) + + # Add full attention mask for specific cases + full_attention_mask_cfg: list[int] | None = kwargs.get( + "full_attention_mask" + ) + if full_attention_mask_cfg is not None: + selected_masks_cfg.append(full_attention_mask_cfg) + + timesteps_cfg: int = kwargs.get("timesteps", time_step_num) + if skip_time_steps_cfg is None: + skip_time_steps_cfg = 12 + # Select best mask strategy using combined results + mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( + averaged_results, + selected_masks_cfg, + skip_time_steps_cfg, + timesteps_cfg, + head_num, + ) + + # Save mask strategy + os.makedirs(save_dir_cfg, exist_ok=True) + file_path = os.path.join( + save_dir_cfg, f"mask_strategy_s{skip_time_steps_cfg}.json" + ) + with open(file_path, "w") as f: + json.dump(mask_strategy, f, indent=4) + print(f"Successfully saved mask_strategy to {file_path}") + + # Print sparsity and strategy counts for information + print(f"Overall sparsity: {sparsity:.4f}") + print("\nStrategy usage counts:") + total_heads = time_step_num * layer_num * head_num # Fixed dimensions + for strategy, count in strategy_counts.items(): + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list( + mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num + ) + + return mask_strategy_3d + + else: # STA_inference + # Get parameters with defaults + load_path: str | None = kwargs.get( + "load_path", "mask_candidates/mask_strategy.json" + ) + if load_path is None: + raise ValueError("load_path is required for STA_inference mode") + + # Load previously saved mask strategy + with open(load_path) as f: + mask_strategy = json.load(f) + + # Convert dictionary to 3D list with fixed dimensions + mask_strategy_3d = dict_to_3d_list( + mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num + ) + + return mask_strategy_3d + + +# Helper functions + + +def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]: + """Read and parse JSON files containing mask search results.""" + json_contents: list[dict[str, Any]] = [] + + # List files only in the current directory (no walk) + files = os.listdir(folder_path) + # Filter files + matching_files = [f for f in files if "mask" in f and f.endswith(".json")] + print(f"Found {len(matching_files)} matching files: {matching_files}") + + for file_name in matching_files: + file_path = os.path.join(folder_path, file_name) + with open(file_path) as file: + data = json.load(file) + json_contents.append(data) + + return json_contents + + +def average_head_losses( + results: list[dict[str, Any]], selected_masks: list[list[int]] +) -> dict[str, dict[str, np.ndarray]]: + """Average losses across all prompts for each mask strategy.""" + # Initialize a dictionary to store the averaged results + averaged_losses: dict[str, dict[str, np.ndarray]] = {} + loss_type = "L2_loss" + # Get all loss types (e.g., 'L2_loss') + averaged_losses[loss_type] = {} + + for mask in selected_masks: + mask_str = str(mask) + data_shape = np.array(results[0][loss_type][mask_str]).shape + accumulated_data = np.zeros(data_shape) + + # Sum across all prompts + for prompt_result in results: + accumulated_data += np.array(prompt_result[loss_type][mask_str]) + + # Average by dividing by number of prompts + averaged_data = accumulated_data / len(results) + averaged_losses[loss_type][mask_str] = averaged_data + + return averaged_losses + + +def select_best_mask_strategy( + averaged_results: dict[str, dict[str, np.ndarray]], + selected_masks: list[list[int]], + skip_time_steps: int = 12, + timesteps: int = 50, + head_num: int = 40, +) -> tuple[dict[str, list[int]], float, dict[str, int]]: + """Select the best mask strategy for each head based on loss minimization.""" + best_mask_strategy: dict[str, list[int]] = {} + loss_type = "L2_loss" + # Get the shape of time steps and layers + layers = len(averaged_results[loss_type][str(selected_masks[0])][0]) + + # Counter for sparsity calculation + total_tokens = 0 # total number of masked tokens + total_length = 0 # total sequence length + + strategy_counts: dict[str, int] = { + str(strategy): 0 for strategy in selected_masks + } + full_attn_strategy = selected_masks[-1] # Last strategy is full attention + print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ") + + for t in range(timesteps): + for layer_idx in range(layers): + for h in range(head_num): + if t < skip_time_steps: # First steps use full attention + strategy = full_attn_strategy + else: + # Get losses for this head across all strategies + head_losses = [] + for strategy in selected_masks[ + :-1 + ]: # Exclude full attention + head_losses.append( + averaged_results[loss_type][str(strategy)][t][ + layer_idx + ][h] + ) + + # Find which strategy gives minimum loss + best_strategy_idx = np.argmin(head_losses) + strategy = selected_masks[best_strategy_idx] + + best_mask_strategy[f"{t}_{layer_idx}_{h}"] = strategy + + # Calculate sparsity + nums = strategy # strategy is already a list of numbers + total_tokens += ( + nums[0] * nums[1] * nums[2] + ) # masked tokens for chosen strategy + total_length += ( + full_attn_strategy[0] + * full_attn_strategy[1] + * full_attn_strategy[2] + ) + + # Count strategy usage + strategy_counts[str(strategy)] += 1 + + overall_sparsity = 1 - total_tokens / total_length + + return best_mask_strategy, overall_sparsity, strategy_counts + + +def save_mask_search_results( + mask_search_final_result: list[dict[str, list[float]]], + prompt: str, + mask_strategies: list[str], + output_dir: str = "output/mask_search_result/", +) -> str | None: + if not mask_search_final_result: + print("No mask search results to save") + return None + + # Create result dictionary with defaultdict for nested lists + mask_search_dict: dict[str, dict[str, list[list[float]]]] = { + "L2_loss": defaultdict(list), + "L1_loss": defaultdict(list), + } + + mask_selected = list(range(len(mask_strategies))) + selected_masks: list[list[int]] = [] + for index in mask_selected: + mask = mask_strategies[index] + masks_list = [int(x) for x in mask.split(",")] + selected_masks.append(masks_list) + + # Process each mask strategy + for i, mask_strategy in enumerate(selected_masks): + mask_strategy_str = str(mask_strategy) + # Process L2 loss + step_results: list[list[float]] = [] + for step_data in mask_search_final_result: + if isinstance(step_data, dict) and "L2_loss" in step_data: + layer_losses = [float(loss) for loss in step_data["L2_loss"]] + step_results.append(layer_losses) + mask_search_dict["L2_loss"][mask_strategy_str] = step_results + + step_results = [] + for step_data in mask_search_final_result: + if isinstance(step_data, dict) and "L1_loss" in step_data: + layer_losses = [float(loss) for loss in step_data["L1_loss"]] + step_results.append(layer_losses) + mask_search_dict["L1_loss"][mask_strategy_str] = step_results + + # Create the output directory if it doesn't exist + os.makedirs(output_dir, exist_ok=True) + + # Create a filename based on the first 20 characters of the prompt + filename = prompt[:50].replace(" ", "_") + filepath = os.path.join(output_dir, f"mask_search_{filename}.json") + + # Save the results to a JSON file + with open(filepath, "w") as f: + json.dump(mask_search_dict, f, indent=4) + + print(f"Successfully saved mask research results to {filepath}") + + return filepath diff --git a/worldcompass/fastvideo/__init__.py b/worldcompass/fastvideo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9df5c6874cbf08ffc8f441183fdda217e7c07391 --- /dev/null +++ b/worldcompass/fastvideo/__init__.py @@ -0,0 +1,15 @@ +try: + from fastvideo.configs.pipelines import PipelineConfig + from fastvideo.configs.sample import SamplingParam + from fastvideo.entrypoints.video_generator import VideoGenerator + from fastvideo.version import __version__ + + __all__ = [ + "VideoGenerator", + "PipelineConfig", + "SamplingParam", + "__version__", + ] + +except Exception as e: + __all__ = [] diff --git a/worldcompass/fastvideo/attention/__init__.py b/worldcompass/fastvideo/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6d9cfa0c560a72379a44be43d5f53583a270ee2 --- /dev/null +++ b/worldcompass/fastvideo/attention/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.attention.backends.abstract import ( + AttentionBackend, + AttentionMetadata, + AttentionMetadataBuilder, +) +from fastvideo.attention.layer import ( + DistributedAttention, + DistributedAttention_VSA, + LocalAttention, +) +from fastvideo.attention.selector import get_attn_backend + +__all__ = [ + "DistributedAttention", + "LocalAttention", + "DistributedAttention_VSA", + "AttentionBackend", + "AttentionMetadata", + "AttentionMetadataBuilder", + # "AttentionState", + "get_attn_backend", +] diff --git a/worldcompass/fastvideo/attention/backends/__init__.py b/worldcompass/fastvideo/attention/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/attention/backends/abstract.py b/worldcompass/fastvideo/attention/backends/abstract.py new file mode 100644 index 0000000000000000000000000000000000000000..8145ec4337802b69255a05949a0c1934dfa46ce4 --- /dev/null +++ b/worldcompass/fastvideo/attention/backends/abstract.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py + +from abc import ABC, abstractmethod +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar + +if TYPE_CHECKING: + pass + +import torch + + +class AttentionBackend(ABC): + """Abstract class for attention backends.""" + + # For some attention backends, we allocate an output tensor before + # calling the custom op. When piecewise cudagraph is enabled, this + # makes sure the output tensor is allocated inside the cudagraph. + accept_output_buffer: bool = False + + @staticmethod + @abstractmethod + def get_name() -> str: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_impl_cls() -> type["AttentionImpl"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + # @staticmethod + # @abstractmethod + # def get_state_cls() -> Type["AttentionState"]: + # raise NotImplementedError + + # @classmethod + # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": + # return cls.get_metadata_cls()(*args, **kwargs) + + @staticmethod + @abstractmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + +@dataclass +class AttentionMetadata: + """Attention metadata for prefill and decode batched together.""" + + # Current step of diffusion process + current_timestep: int + + def asdict_zerocopy( + self, skip_fields: set[str] | None = None + ) -> dict[str, Any]: + """Similar to dataclasses.asdict, but avoids deepcopying.""" + if skip_fields is None: + skip_fields = set() + # Note that if we add dataclasses as fields, they will need + # similar handling. + return { + field.name: getattr(self, field.name) + for field in fields(self) + if field.name not in skip_fields + } + + +T = TypeVar("T", bound=AttentionMetadata) + + +class AttentionMetadataBuilder(ABC, Generic[T]): + """Abstract class for attention metadata builders.""" + + @abstractmethod + def __init__(self) -> None: + """Create the builder, remember some configuration and parameters.""" + raise NotImplementedError + + @abstractmethod + def prepare(self) -> None: + """Prepare for one batch.""" + raise NotImplementedError + + @abstractmethod + def build( + self, + **kwargs: dict[str, Any], + ) -> AttentionMetadata: + """Build attention metadata with on-device tensors.""" + raise NotImplementedError + + +class AttentionLayer(Protocol): + _k_scale: torch.Tensor + _v_scale: torch.Tensor + _k_scale_float: float + _v_scale_float: float + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: ... + + +class AttentionImpl(ABC, Generic[T]): + @abstractmethod + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + raise NotImplementedError + + def preprocess_qkv( + self, qkv: torch.Tensor, attn_metadata: T + ) -> torch.Tensor: + """Preprocess QKV tensor before performing attention operation. + + Default implementation returns the tensor unchanged. + Subclasses can override this to implement custom preprocessing + like reshaping, tiling, scaling, or other transformations. + + Called AFTER all_to_all for distributed attention + + Args: + qkv: The query-key-value tensor + attn_metadata: Metadata for the attention operation + + Returns: + Processed QKV tensor + """ + return qkv + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: T, + ) -> torch.Tensor: + """Postprocess the output tensor after the attention operation. + + Default implementation returns the tensor unchanged. + Subclasses can override this to implement custom postprocessing + like untiling, scaling, or other transformations. + + Called BEFORE all_to_all for distributed attention + + Args: + output: The output tensor from the attention operation + attn_metadata: Metadata for the attention operation + + Returns: + Postprocessed output tensor + """ + + return output + + @abstractmethod + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T, + ) -> torch.Tensor: + raise NotImplementedError diff --git a/worldcompass/fastvideo/attention/backends/flash_attn.py b/worldcompass/fastvideo/attention/backends/flash_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..b06299e982c6e098eae17d08f53fbb9e12ffc0cb --- /dev/null +++ b/worldcompass/fastvideo/attention/backends/flash_attn.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from flash_attn import flash_attn_func as flash_attn_2_func + +try: + from flash_attn_interface import flash_attn_func as flash_attn_3_func + + # flash_attn 3 no longer have a different API, see following commit: + # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7 + flash_attn_func = flash_attn_3_func +except ImportError: + flash_attn_func = flash_attn_2_func + +from fastvideo.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class FlashAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN" + + @staticmethod + def get_impl_cls() -> type["FlashAttentionImpl"]: + return FlashAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + +class FlashAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ): + output = flash_attn_func( + query, # type: ignore[no-untyped-call] + key, + value, + softmax_scale=self.softmax_scale, + causal=self.causal, + ) + return output diff --git a/worldcompass/fastvideo/attention/backends/sdpa.py b/worldcompass/fastvideo/attention/backends/sdpa.py new file mode 100644 index 0000000000000000000000000000000000000000..f35dde820423b6aa86a8f674df9541e02936a1ef --- /dev/null +++ b/worldcompass/fastvideo/attention/backends/sdpa.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from fastvideo.attention.backends.abstract import ( # FlashAttentionMetadata, + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class SDPABackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "SDPA" + + @staticmethod + def get_impl_cls() -> type["SDPAImpl"]: + return SDPAImpl + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class SDPAImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + self.dropout = extra_impl_args.get("dropout_p", 0.0) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + # transpose to bs, heads, seq_len, head_dim + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn_kwargs = { + "attn_mask": None, + "dropout_p": self.dropout, + "is_causal": self.causal, + "scale": self.softmax_scale, + } + if query.shape[1] != key.shape[1]: + attn_kwargs["enable_gqa"] = True + output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, **attn_kwargs + ) + output = output.transpose(1, 2) + return output diff --git a/worldcompass/fastvideo/attention/backends/sliding_tile_attn.py b/worldcompass/fastvideo/attention/backends/sliding_tile_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..a32aacbeedb8071323adad6d4a3d9124b27b97c4 --- /dev/null +++ b/worldcompass/fastvideo/attention/backends/sliding_tile_attn.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +from dataclasses import dataclass +from typing import Any + +import torch +from einops import rearrange +from st_attn import sliding_tile_attention + +import fastvideo.envs as envs +from fastvideo.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from fastvideo.distributed import get_sp_group +from fastvideo.forward_context import ForwardContext, get_forward_context +from fastvideo.logger import init_logger +from fastvideo.utils import dict_to_3d_list + +logger = init_logger(__name__) + + +class RangeDict(dict): + def __getitem__(self, item: int) -> str: + for key in self.keys(): + if isinstance(key, tuple): + low, high = key + if low <= item <= high: + return str(super().__getitem__(key)) + elif key == item: + return str(super().__getitem__(key)) + raise KeyError(f"seq_len {item} not supported for STA") + + +class SlidingTileAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + # TODO(will-refactor): check this + return [32, 64, 96, 128, 160, 192, 224, 256] + + @staticmethod + def get_name() -> str: + return "SLIDING_TILE_ATTN" + + @staticmethod + def get_impl_cls() -> type["SlidingTileAttentionImpl"]: + return SlidingTileAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]: + return SlidingTileAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]: + return SlidingTileAttentionMetadataBuilder + + +@dataclass +class SlidingTileAttentionMetadata(AttentionMetadata): + current_timestep: int + STA_param: list[ + list[Any] + ] # each timestep with one metadata, shape [num_layers, num_heads] + + +class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): + def __init__(self): + pass + + def prepare(self): + pass + + def build( # type: ignore + self, + STA_param: list[list[Any]], + current_timestep: int, + **kwargs: dict[str, Any], + ) -> SlidingTileAttentionMetadata: + param = STA_param + if param is None: + return SlidingTileAttentionMetadata( + current_timestep=current_timestep, STA_param=[] + ) + return SlidingTileAttentionMetadata( + current_timestep=current_timestep, STA_param=param[current_timestep] + ) + + +class SlidingTileAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + # TODO(will-refactor): for now this is the mask strategy, but maybe we should + # have a more general config for STA? + config_file = envs.FASTVIDEO_ATTENTION_CONFIG + if config_file is None: + raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set") + + # TODO(kevin): get mask strategy for different STA modes + with open(config_file) as f: + mask_strategy = json.load(f) + self.mask_strategy = dict_to_3d_list(mask_strategy) + + self.prefix = prefix + sp_group = get_sp_group() + self.sp_size = sp_group.world_size + # STA config + self.STA_base_tile_size = [6, 8, 8] + self.dit_seq_shape_mapping = RangeDict( + { + (115200, 115456): "30x48x80", + 82944: "36x48x48", + 69120: "18x48x80", + } + ) + self.full_window_mapping = { + "30x48x80": [5, 6, 10], + "36x48x48": [6, 6, 6], + "18x48x80": [3, 6, 10], + } + + def tile(self, x: torch.Tensor) -> torch.Tensor: + return rearrange( + x, + "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", + n_t=self.full_window_size[0], + n_h=self.full_window_size[1], + n_w=self.full_window_size[2], + ts_t=self.STA_base_tile_size[0], + ts_h=self.STA_base_tile_size[1], + ts_w=self.STA_base_tile_size[2], + ) + + def untile(self, x: torch.Tensor) -> torch.Tensor: + x = rearrange( + x, + "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", + n_t=self.full_window_size[0], + n_h=self.full_window_size[1], + n_w=self.full_window_size[2], + ts_t=self.STA_base_tile_size[0], + ts_h=self.STA_base_tile_size[1], + ts_w=self.STA_base_tile_size[2], + ) + return x + + def preprocess_qkv( + self, + qkv: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + img_sequence_length = qkv.shape[1] + self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length] + self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str] + self.dit_seq_shape_int = list( + map(int, self.dit_seq_shape_str.split("x")) + ) + self.img_seq_length = ( + self.dit_seq_shape_int[0] + * self.dit_seq_shape_int[1] + * self.dit_seq_shape_int[2] + ) + return self.tile(qkv) + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: SlidingTileAttentionMetadata, + ) -> torch.Tensor: + return self.untile(output) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_metadata: SlidingTileAttentionMetadata, + ) -> torch.Tensor: + if self.mask_strategy is None: + raise ValueError( + "mask_strategy cannot be None for SlidingTileAttention" + ) + if self.mask_strategy[0] is None: + raise ValueError( + "mask_strategy[0] cannot be None for SlidingTileAttention" + ) + + timestep = attn_metadata.current_timestep + forward_context: ForwardContext = get_forward_context() + forward_batch = forward_context.forward_batch + if forward_batch is None: + raise ValueError("forward_batch cannot be None") + # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl' + layer_idx = int(self.prefix.split(".")[-3]) + if ( + attn_metadata.STA_param is None + or len(attn_metadata.STA_param) <= layer_idx + ): + raise ValueError("Invalid STA_param") + STA_param = attn_metadata.STA_param[layer_idx] + + text_length = q.shape[1] - self.img_seq_length + has_text = text_length > 0 + + query = q.transpose(1, 2).contiguous() + key = k.transpose(1, 2).contiguous() + value = v.transpose(1, 2).contiguous() + + head_num = query.size(1) + sp_group = get_sp_group() + current_rank = sp_group.rank_in_group + start_head = current_rank * head_num + + # searching or tuning mode + if len(STA_param) < head_num * sp_group.world_size: + sparse_attn_hidden_states_all = [] + full_mask_window = STA_param[-1] + for window_size in STA_param[:-1]: + sparse_hidden_states = sliding_tile_attention( + query, + key, + value, + [window_size] * head_num, + text_length, + has_text, + self.dit_seq_shape_str, + ).transpose(1, 2) + sparse_attn_hidden_states_all.append(sparse_hidden_states) + + hidden_states = sliding_tile_attention( + query, + key, + value, + [full_mask_window] * head_num, + text_length, + has_text, + self.dit_seq_shape_str, + ).transpose(1, 2) + + attn_L2_loss = [] + attn_L1_loss = [] + # average loss across all heads + for sparse_attn_hidden_states in sparse_attn_hidden_states_all: + # L2 loss + attn_L2_loss_ = ( + torch.mean( + ( + sparse_attn_hidden_states.float() + - hidden_states.float() + ) + ** 2, + dim=[0, 1, 3], + ) + .cpu() + .numpy() + ) + attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_] + attn_L2_loss.append(attn_L2_loss_) + # L1 loss + attn_L1_loss_ = ( + torch.mean( + torch.abs( + sparse_attn_hidden_states.float() + - hidden_states.float() + ), + dim=[0, 1, 3], + ) + .cpu() + .numpy() + ) + attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_] + attn_L1_loss.append(attn_L1_loss_) + + layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss} + + if forward_batch.is_cfg_negative: + if forward_batch.mask_search_final_result_neg is not None: + forward_batch.mask_search_final_result_neg[timestep].append( + layer_loss_save + ) + else: + if forward_batch.mask_search_final_result_pos is not None: + forward_batch.mask_search_final_result_pos[timestep].append( + layer_loss_save + ) + else: + windows = [ + STA_param[head_idx + start_head] for head_idx in range(head_num) + ] + + hidden_states = sliding_tile_attention( + query, + key, + value, + windows, + text_length, + has_text, + self.dit_seq_shape_str, + ).transpose(1, 2) + + return hidden_states diff --git a/worldcompass/fastvideo/attention/backends/video_sparse_attn.py b/worldcompass/fastvideo/attention/backends/video_sparse_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..c3a1212812cba008a140856b76019016b2ac1a36 --- /dev/null +++ b/worldcompass/fastvideo/attention/backends/video_sparse_attn.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +import functools +import math +from dataclasses import dataclass + +import torch + +try: + from vsa import video_sparse_attn +except ImportError: + video_sparse_attn = None + +from typing import Any + +from fastvideo.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from fastvideo.distributed import get_sp_group +from fastvideo.logger import init_logger + +logger = init_logger(__name__) +VSA_TILE_SIZE = (4, 4, 4) + + +@functools.lru_cache(maxsize=10) +def get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + T, H, W = dit_seq_shape + ts, hs, ws = tile_size + indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape( + T, H, W + ) + ls = [] + for t in range(math.ceil(T / ts)): + for h in range(math.ceil(H / hs)): + for w in range(math.ceil(W / ws)): + ls.append( + indices[ + t * ts : min(t * ts + ts, T), + h * hs : min(h * hs + hs, H), + w * ws : min(w * ws + ws, W), + ].flatten() + ) + index = torch.cat(ls, dim=0) + return index + + +@functools.lru_cache(maxsize=10) +def get_reverse_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + return torch.argsort( + get_tile_partition_indices(dit_seq_shape, tile_size, device) + ) + + +@functools.lru_cache(maxsize=10) +def construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + num_tiles: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + """Compute the number of valid (non‑padded) tokens inside every (ts_t × ts_h × ts_w) tile after + padding ‑‑ flattened in the order (t‑tile, h‑tile, w‑tile) that `rearrange` uses. + + Returns + ------- + torch.LongTensor # shape: [∏ full_window_size] + """ + # unpack + t, h, w = dit_seq_shape + ts_t, ts_h, ts_w = VSA_TILE_SIZE + n_t, n_h, n_w = num_tiles + + def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor: + """Vector with the size of each tile along one dimension.""" + sizes = torch.full((n_tiles,), tile, dtype=torch.int, device=device) + # size of last (possibly partial) tile + remainder = dim_len - (n_tiles - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(t, ts_t, n_t) # [n_t] + h_sizes = _sizes(h, ts_h, n_h) # [n_h] + w_sizes = _sizes(w, ts_w, n_w) # [n_w] + + # broadcast‑multiply to get voxels per tile, then flatten + block_sizes = ( + t_sizes[:, None, None] # [n_t, 1, 1] + * h_sizes[None, :, None] # [1, n_h, 1] + * w_sizes[None, None, :] # [1, 1, n_w] + ).reshape( + -1 + ) # [n_t * n_h * n_w] + + return block_sizes + + +@functools.lru_cache(maxsize=10) +def get_non_pad_index( + variable_block_sizes: torch.LongTensor, + max_block_size: int, +): + n_win = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts_pad = torch.arange(n_win, device=device) * max_block_size + index_pad = ( + starts_pad[:, None] + + torch.arange(max_block_size, device=device)[None, :] + ) + index_mask = ( + torch.arange(max_block_size, device=device)[None, :] + < variable_block_sizes[:, None] + ) + return index_pad[index_mask] + + +class VideoSparseAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_name() -> str: + return "VIDEO_SPARSE_ATTN" + + @staticmethod + def get_impl_cls() -> type["VideoSparseAttentionImpl"]: + return VideoSparseAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]: + return VideoSparseAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]: + return VideoSparseAttentionMetadataBuilder + + +@dataclass +class VideoSparseAttentionMetadata(AttentionMetadata): + current_timestep: int + dit_seq_shape: list[int] + VSA_sparsity: float + num_tiles: list[int] + total_seq_length: int + tile_partition_indices: torch.LongTensor + reverse_tile_partition_indices: torch.LongTensor + variable_block_sizes: torch.LongTensor + non_pad_index: torch.LongTensor + + +class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): + def __init__(self): + pass + + def prepare(self): + pass + + def build( # type: ignore + self, + current_timestep: int, + raw_latent_shape: tuple[int, int, int], + patch_size: tuple[int, int, int], + VSA_sparsity: float, + device: torch.device, + **kwargs: dict[str, Any], + ) -> VideoSparseAttentionMetadata: + patch_size = patch_size + dit_seq_shape = ( + raw_latent_shape[0] // patch_size[0], + raw_latent_shape[1] // patch_size[1], + raw_latent_shape[2] // patch_size[2], + ) + + num_tiles = ( + math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), + math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]), + math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]), + ) + total_seq_length = math.prod(dit_seq_shape) + + tile_partition_indices = get_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device + ) + reverse_tile_partition_indices = get_reverse_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device + ) + variable_block_sizes = construct_variable_block_sizes( + dit_seq_shape, num_tiles, device + ) + non_pad_index = get_non_pad_index( + variable_block_sizes, math.prod(VSA_TILE_SIZE) + ) + + return VideoSparseAttentionMetadata( + current_timestep=current_timestep, + dit_seq_shape=dit_seq_shape, # type: ignore + VSA_sparsity=VSA_sparsity, # type: ignore + num_tiles=num_tiles, # type: ignore + total_seq_length=total_seq_length, # type: ignore + tile_partition_indices=tile_partition_indices, # type: ignore + reverse_tile_partition_indices=reverse_tile_partition_indices, + variable_block_sizes=variable_block_sizes, + non_pad_index=non_pad_index, + ) + + +class VideoSparseAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.prefix = prefix + sp_group = get_sp_group() + self.sp_size = sp_group.world_size + + def tile( + self, + x: torch.Tensor, + num_tiles: list[int], + tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor, + ) -> torch.Tensor: + t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0] + h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1] + w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2] + + x_padded = torch.zeros( + ( + x.shape[0], + t_padded_size * h_padded_size * w_padded_size, + x.shape[-2], + x.shape[-1], + ), + device=x.device, + dtype=x.dtype, + ) + x_padded[:, non_pad_index] = x[:, tile_partition_indices] + return x_padded + + def untile( + self, + x: torch.Tensor, + reverse_tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor, + ) -> torch.Tensor: + x = x[:, non_pad_index][:, reverse_tile_partition_indices] + return x + + def preprocess_qkv( + self, + qkv: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + return self.tile( + qkv, + attn_metadata.num_tiles, + attn_metadata.tile_partition_indices, + attn_metadata.non_pad_index, + ) + + def postprocess_output( + self, + output: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + return self.untile( + output, + attn_metadata.reverse_tile_partition_indices, + attn_metadata.non_pad_index, + ) + + def forward( # type: ignore[override] + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + gate_compress: torch.Tensor, + attn_metadata: VideoSparseAttentionMetadata, + ) -> torch.Tensor: + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + gate_compress = gate_compress.transpose(1, 2).contiguous() + + VSA_sparsity = attn_metadata.VSA_sparsity + + cur_topk = math.ceil( + (1 - VSA_sparsity) + * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)) + ) + + if video_sparse_attn is None: + raise NotImplementedError("video_sparse_attn is not installed") + hidden_states = video_sparse_attn( + query, + key, + value, + variable_block_sizes=attn_metadata.variable_block_sizes, + topk=cur_topk, + block_size=VSA_TILE_SIZE, + compress_attn_weight=gate_compress, + ).transpose(1, 2) + + return hidden_states diff --git a/worldcompass/fastvideo/attention/layer.py b/worldcompass/fastvideo/attention/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..b3491b8a2e34b31eb79c89954a7f99e590a3b8bf --- /dev/null +++ b/worldcompass/fastvideo/attention/layer.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn + +from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather, + sequence_model_parallel_all_to_all_4D, +) +from fastvideo.distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_world_size, +) +from fastvideo.forward_context import ForwardContext, get_forward_context +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.utils import get_compute_dtype + + +class DistributedAttention(nn.Module): + """Distributed attention layer.""" + + def __init__( + self, + num_heads: int, + head_size: int, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: ( + tuple[AttentionBackendEnum, ...] | None + ) = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + super().__init__() + if softmax_scale is None: + self.softmax_scale = head_size**-0.5 + else: + self.softmax_scale = softmax_scale + + if num_kv_heads is None: + num_kv_heads = num_heads + + dtype = get_compute_dtype() + attn_backend = get_attn_backend( + head_size, + dtype, + supported_attention_backends=supported_attention_backends, + ) + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + causal=causal, + softmax_scale=self.softmax_scale, + num_kv_heads=num_kv_heads, + prefix=f"{prefix}.impl", + **extra_impl_args, + ) + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + self.backend = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: torch.Tensor | None = None, + replicated_k: torch.Tensor | None = None, + replicated_v: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass for distributed attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens + replicated_k (Optional[torch.Tensor]): Replicated key tensor + replicated_v (Optional[torch.Tensor]): Replicated value tensor + + Returns: + Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: + - o (torch.Tensor): Output tensor after attention for the main sequence + - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided + """ + # Check input shapes + assert ( + q.dim() == 4 and k.dim() == 4 and v.dim() == 4 + ), "Expected 4D tensors" + batch_size, seq_len, num_heads, head_dim = q.shape + local_rank = get_sp_parallel_rank() + world_size = get_sp_world_size() + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Stack QKV + qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] + + # Redistribute heads across sequence dimension + qkv = sequence_model_parallel_all_to_all_4D( + qkv, scatter_dim=2, gather_dim=1 + ) + # Apply backend-specific preprocess_qkv + qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) + + # Concatenate with replicated QKV if provided + if replicated_q is not None: + assert replicated_k is not None and replicated_v is not None + replicated_qkv = torch.cat( + [replicated_q, replicated_k, replicated_v], dim=0 + ) # [3, seq_len, num_heads, head_dim] + heads_per_rank = num_heads // world_size + replicated_qkv = replicated_qkv[ + :, + :, + local_rank * heads_per_rank : (local_rank + 1) * heads_per_rank, + ] + qkv = torch.cat([qkv, replicated_qkv], dim=1) + + q, k, v = qkv.chunk(3, dim=0) + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + + # Redistribute back if using sequence parallelism + replicated_output = None + if replicated_q is not None: + replicated_output = output[:, seq_len * world_size :] + output = output[:, : seq_len * world_size] + # TODO: make this asynchronous + replicated_output = sequence_model_parallel_all_gather( + replicated_output.contiguous(), dim=2 + ) + # Apply backend-specific postprocess_output + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + output = sequence_model_parallel_all_to_all_4D( + output, scatter_dim=1, gather_dim=2 + ) + return output, replicated_output + + +class DistributedAttention_VSA(DistributedAttention): + """Distributed attention layer with VSA support.""" + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: torch.Tensor | None = None, + replicated_k: torch.Tensor | None = None, + replicated_v: torch.Tensor | None = None, + gate_compress: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass for distributed attention. + + Args: + q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] + gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens + replicated_k (Optional[torch.Tensor]): Replicated key tensor + replicated_v (Optional[torch.Tensor]): Replicated value tensor + + Returns: + Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: + - o (torch.Tensor): Output tensor after attention for the main sequence + - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided + """ + # Check text tokens are not supported for VSA now + assert ( + replicated_q is None + and replicated_k is None + and replicated_v is None + ), "Replicated QKV is not supported for VSA now" + # Check input shapes + assert ( + q.dim() == 4 and k.dim() == 4 and v.dim() == 4 + ), "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Stack QKV + qkvg = torch.cat( + [q, k, v, gate_compress], dim=0 + ) # [3, seq_len, num_heads, head_dim] + + # Redistribute heads across sequence dimension + qkvg = sequence_model_parallel_all_to_all_4D( + qkvg, scatter_dim=2, gather_dim=1 + ) + + qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata) + + q, k, v, gate_compress = qkvg.chunk(4, dim=0) + output = self.attn_impl.forward(q, k, v, gate_compress, ctx_attn_metadata) # type: ignore[call-arg] + + # Redistribute back if using sequence parallelism + replicated_output = None + + # Apply backend-specific postprocess_output + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + output = sequence_model_parallel_all_to_all_4D( + output, scatter_dim=1, gather_dim=2 + ) + return output, replicated_output + + +class LocalAttention(nn.Module): + """Attention layer.""" + + def __init__( + self, + num_heads: int, + head_size: int, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: ( + tuple[AttentionBackendEnum, ...] | None + ) = None, + **extra_impl_args, + ) -> None: + super().__init__() + if softmax_scale is None: + self.softmax_scale = head_size**-0.5 + else: + self.softmax_scale = softmax_scale + if num_kv_heads is None: + num_kv_heads = num_heads + + dtype = get_compute_dtype() + attn_backend = get_attn_backend( + head_size, + dtype, + supported_attention_backends=supported_attention_backends, + ) + impl_cls = attn_backend.get_impl_cls() + self.attn_impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + softmax_scale=self.softmax_scale, + num_kv_heads=num_kv_heads, + causal=causal, + **extra_impl_args, + ) + self.num_heads = num_heads + self.head_size = head_size + self.num_kv_heads = num_kv_heads + self.backend = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + """Apply local attention between query, key and value tensors. + + Args: + q (torch.Tensor): Query tensor of shape [batch_size, seq_len, num_heads, head_dim] + k (torch.Tensor): Key tensor of shape [batch_size, seq_len, num_heads, head_dim] + v (torch.Tensor): Value tensor of shape [batch_size, seq_len, num_heads, head_dim] + + Returns: + torch.Tensor: Output tensor after local attention + """ + # Check input shapes + assert ( + q.dim() == 4 and k.dim() == 4 and v.dim() == 4 + ), "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + return output diff --git a/worldcompass/fastvideo/attention/selector.py b/worldcompass/fastvideo/attention/selector.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4cc906cb4d0f503dfa608ad406a5984faf2621 --- /dev/null +++ b/worldcompass/fastvideo/attention/selector.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py + +import os +from collections.abc import Generator +from contextlib import contextmanager +from functools import cache +from typing import cast + +import torch + +import fastvideo.envs as envs +from fastvideo.attention.backends.abstract import AttentionBackend +from fastvideo.logger import init_logger +from fastvideo.platforms import AttentionBackendEnum, current_platform +from fastvideo.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname + +logger = init_logger(__name__) + + +def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None: + """Convert a string backend name to a _Backend enum value. + + Returns: + * _Backend: enum value if backend_name is a valid in-tree type + * None: otherwise it's an invalid in-tree type or an out-of-tree platform is + loaded. + """ + assert backend_name is not None + return ( + AttentionBackendEnum[backend_name] + if backend_name in AttentionBackendEnum.__members__ + else None + ) + + +def get_env_variable_attn_backend() -> AttentionBackendEnum | None: + """Get the backend override specified by the FastVideo attention backend environment variable, + if one is specified. + + Returns: + + * _Backend enum value if an override is specified + * None otherwise + """ + backend_name = os.environ.get(STR_BACKEND_ENV_VAR) + return None if backend_name is None else backend_name_to_enum(backend_name) + + +# Global state allows a particular choice of backend +# to be forced, overriding the logic which auto-selects +# a backend based on system & workload configuration +# (default behavior if this variable is None) +# +# THIS SELECTION TAKES PRECEDENCE OVER THE +# FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE +forced_attn_backend: AttentionBackendEnum | None = None + + +def global_force_attn_backend( + attn_backend: AttentionBackendEnum | None, +) -> None: + """Force all attention operations to use a specified backend. + + Passing `None` for the argument re-enables automatic + backend selection., + + Arguments: + + * attn_backend: backend selection (None to revert to auto) + """ + global forced_attn_backend + forced_attn_backend = attn_backend + + +def get_global_forced_attn_backend() -> AttentionBackendEnum | None: + """Get the currently-forced choice of attention backend, or None if auto-selection is currently + enabled.""" + return forced_attn_backend + + +def get_attn_backend( + head_size: int, + dtype: torch.dtype, + supported_attention_backends: ( + tuple[AttentionBackendEnum, ...] | None + ) = None, +) -> type[AttentionBackend]: + return _cached_get_attn_backend( + head_size, dtype, supported_attention_backends + ) + + +@cache +def _cached_get_attn_backend( + head_size: int, + dtype: torch.dtype, + supported_attention_backends: ( + tuple[AttentionBackendEnum, ...] | None + ) = None, +) -> type[AttentionBackend]: + # Check whether a particular choice of backend was + # previously forced. + # + # THIS SELECTION OVERRIDES THE FASTVIDEO_ATTENTION_BACKEND + # ENVIRONMENT VARIABLE. + if not supported_attention_backends: + raise ValueError("supported_attention_backends is empty") + selected_backend = None + backend_by_global_setting: AttentionBackendEnum | None = ( + get_global_forced_attn_backend() + ) + if backend_by_global_setting is not None: + selected_backend = backend_by_global_setting + else: + # Check the environment variable and override if specified + backend_by_env_var: str | None = envs.FASTVIDEO_ATTENTION_BACKEND + if backend_by_env_var is not None: + selected_backend = backend_name_to_enum(backend_by_env_var) + + # get device-specific attn_backend + if selected_backend not in supported_attention_backends: + selected_backend = None + attention_cls = current_platform.get_attn_backend_cls( + selected_backend, head_size, dtype + ) + if not attention_cls: + raise ValueError( + f"Invalid attention backend for {current_platform.device_name}" + ) + return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls)) + + +@contextmanager +def global_force_attn_backend_context_manager( + attn_backend: AttentionBackendEnum, +) -> Generator[None, None, None]: + """Globally force a FastVideo attention backend override within a context manager, reverting the + global attention backend override to its prior state upon exiting the context manager. + + Arguments: + + * attn_backend: attention backend to force + + Returns: + + * Generator + """ + + # Save the current state of the global backend override (if any) + original_value = get_global_forced_attn_backend() + + # Globally force the new backend override + global_force_attn_backend(attn_backend) + + # Yield control back to the enclosed code block + try: + yield + finally: + # Revert the original global backend override, if any + global_force_attn_backend(original_value) diff --git a/worldcompass/fastvideo/configs/__init__.py b/worldcompass/fastvideo/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/configs/configs.py b/worldcompass/fastvideo/configs/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..7f01fec2fd7e18926b435085dce723e7c9c3a7f7 --- /dev/null +++ b/worldcompass/fastvideo/configs/configs.py @@ -0,0 +1,180 @@ +import dataclasses +from typing import Any, Optional + +from fastvideo.configs.utils import update_config_from_args +from fastvideo.utils import FlexibleArgumentParser, StoreBoolean + + +@dataclasses.dataclass +class PreprocessConfig: + """Configuration for preprocessing operations.""" + + # Model and dataset configuration + model_path: str = "" + dataset_path: str = "" + dataset_output_dir: str = "./output" + + # Dataloader configuration + dataloader_num_workers: int = 1 + preprocess_video_batch_size: int = 2 + + # Saver configuration + samples_per_file: int = 64 + flush_frequency: int = 256 + + # Video processing parameters + max_height: int = 480 + max_width: int = 848 + num_frames: int = 163 + video_length_tolerance_range: float = 2.0 + train_fps: int = 30 + speed_factor: float = 1.0 + drop_short_ratio: float = 1.0 + do_temporal_sample: bool = False + + # Model configuration + training_cfg_rate: float = 0.0 + + seed: int = 42 + + @staticmethod + def add_cli_args( + parser: FlexibleArgumentParser, prefix: str = "preprocess" + ) -> FlexibleArgumentParser: + """Add preprocessing configuration arguments to the parser.""" + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + + preprocess_args = parser.add_argument_group("Preprocessing Arguments") + # Model & Dataset + preprocess_args.add_argument( + f"--{prefix_with_dot}model-path", + type=str, + default=PreprocessConfig.model_path, + help="Path to the model for preprocessing", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-path", + type=str, + default=PreprocessConfig.dataset_path, + help="Path to the dataset directory for preprocessing", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-output-dir", + type=str, + default=PreprocessConfig.dataset_output_dir, + help="The output directory where the dataset will be written.", + ) + + # Dataloader + preprocess_args.add_argument( + f"--{prefix_with_dot}dataloader-num-workers", + type=int, + default=PreprocessConfig.dataloader_num_workers, + help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}preprocess-video-batch-size", + type=int, + default=PreprocessConfig.preprocess_video_batch_size, + help="Batch size (per device) for the training dataloader.", + ) + + # Saver + preprocess_args.add_argument( + f"--{prefix_with_dot}samples-per-file", + type=int, + default=PreprocessConfig.samples_per_file, + help="Number of samples per output file", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}flush-frequency", + type=int, + default=PreprocessConfig.flush_frequency, + help="How often to save to parquet files", + ) + + # Video processing parameters + preprocess_args.add_argument( + f"--{prefix_with_dot}max-height", + type=int, + default=PreprocessConfig.max_height, + help="Maximum height for video processing", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}max-width", + type=int, + default=PreprocessConfig.max_width, + help="Maximum width for video processing", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}num-frames", + type=int, + default=PreprocessConfig.num_frames, + help="Number of frames to process", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}video-length-tolerance-range", + type=float, + default=PreprocessConfig.video_length_tolerance_range, + help="Video length tolerance range", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}train-fps", + type=int, + default=PreprocessConfig.train_fps, + help="Training FPS", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}speed-factor", + type=float, + default=PreprocessConfig.speed_factor, + help="Speed factor for video processing", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}drop-short-ratio", + type=float, + default=PreprocessConfig.drop_short_ratio, + help="Ratio for dropping short videos", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}do-temporal-sample", + action=StoreBoolean, + default=PreprocessConfig.do_temporal_sample, + help="Whether to do temporal sampling", + ) + + # Model Training configuration + preprocess_args.add_argument( + f"--{prefix_with_dot}training-cfg-rate", + type=float, + default=PreprocessConfig.training_cfg_rate, + help="Training CFG rate", + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}seed", + type=int, + default=PreprocessConfig.seed, + help="Seed for random number generator", + ) + + return parser + + @classmethod + def from_kwargs( + cls, kwargs: dict[str, Any] + ) -> Optional["PreprocessConfig"]: + """Create PreprocessConfig from keyword arguments.""" + preprocess_config = cls() + if not update_config_from_args( + preprocess_config, kwargs, prefix="preprocess", pop_args=True + ): + return None + return preprocess_config + + def check_preprocess_config(self) -> None: + if self.dataset_path == "": + raise ValueError("dataset_path must be set for preprocess mode") + if self.samples_per_file <= 0: + raise ValueError("samples_per_file must be greater than 0") + if self.flush_frequency <= 0: + raise ValueError("flush_frequency must be greater than 0") diff --git a/worldcompass/fastvideo/configs/fasthunyuan_t2v.json b/worldcompass/fastvideo/configs/fasthunyuan_t2v.json new file mode 100644 index 0000000000000000000000000000000000000000..56fefcf36dd5c84562938667aba052959749eb30 --- /dev/null +++ b/worldcompass/fastvideo/configs/fasthunyuan_t2v.json @@ -0,0 +1,45 @@ +{ + "embedded_cfg_scale": 6, + "flow_shift": 17, + "dit_cpu_offload": false, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": true, + "vae_sp": true, + "vae_config": { + "load_encoder": false, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 4, + "use_tiling": true, + "use_temporal_tiling": true, + "use_parallel_tiling": true + }, + "dit_config": { + "prefix": "Hunyuan", + "quant_config": null + }, + "text_encoder_precisions": ["fp16", "fp16"], + "text_encoder_configs": [ + { + "prefix": "llama", + "quant_config": null, + "lora_config": null + }, + { + "prefix": "clip", + "quant_config": null, + "lora_config": null, + "num_hidden_layers_override": null, + "require_post_norm": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false +} diff --git a/worldcompass/fastvideo/configs/models/__init__.py b/worldcompass/fastvideo/configs/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..286c6b9ddcbf73f10e27e98bebe21fe73cb0a3c2 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/__init__.py @@ -0,0 +1,6 @@ +from fastvideo.configs.models.base import ModelConfig +from fastvideo.configs.models.dits.base import DiTConfig +from fastvideo.configs.models.encoders.base import EncoderConfig +from fastvideo.configs.models.vaes.base import VAEConfig + +__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"] diff --git a/worldcompass/fastvideo/configs/models/base.py b/worldcompass/fastvideo/configs/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..989e07bacec99775191b4368198b25a126ac380f --- /dev/null +++ b/worldcompass/fastvideo/configs/models/base.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field, fields +from typing import Any + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +# 1. ArchConfig contains all fields from diffuser's/transformer's config.json +# (i.e. all fields related to the architecture of the model) +# 2. ArchConfig should be inherited & overridden by each model arch_config +# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users +@dataclass +class ArchConfig: + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=list + ) # mapping from huggingface weight names to custom names + + +@dataclass +class ModelConfig: + # Every model config parameter can be categorized into either ArchConfig or everything else + # Diffuser/Transformer parameters + arch_config: ArchConfig = field(default_factory=ArchConfig) + + # FastVideo-specific parameters here + # i.e. STA, quantization, teacache + + def __getattr__(self, name): + # Only called if 'name' is not found in ModelConfig directly + if hasattr(self.arch_config, name): + return getattr(self.arch_config, name) + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + + def __getstate__(self): + # Return a dictionary of attributes to pickle + # Convert to dict and exclude any problematic attributes + state = self.__dict__.copy() + return state + + def __setstate__(self, state): + # Restore instance attributes from the unpickled state + self.__dict__.update(state) + + # This should be used only when loading from transformers/diffusers + def update_model_arch(self, source_model_dict: dict[str, Any]) -> None: + arch_config = self.arch_config + valid_fields = {f.name for f in fields(arch_config)} + + for key, value in source_model_dict.items(): + if key in valid_fields: + setattr(arch_config, key, value) + else: + raise AttributeError( + f"{type(arch_config).__name__} has no field '{key}'" + ) + + if hasattr(arch_config, "__post_init__"): + arch_config.__post_init__() + + def update_model_config(self, source_model_dict: dict[str, Any]) -> None: + assert ( + "arch_config" not in source_model_dict + ), "Source model config shouldn't contain arch_config." + + valid_fields = {f.name for f in fields(self)} + + for key, value in source_model_dict.items(): + if key in valid_fields: + setattr(self, key, value) + else: + logger.warning( + "%s does not contain field '%s'!", type(self).__name__, key + ) + raise AttributeError(f"Invalid field: {key}") + + if hasattr(self, "__post_init__"): + self.__post_init__() diff --git a/worldcompass/fastvideo/configs/models/dits/__init__.py b/worldcompass/fastvideo/configs/models/dits/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72271a525168044cbc488b0f34ef386ce01a598b --- /dev/null +++ b/worldcompass/fastvideo/configs/models/dits/__init__.py @@ -0,0 +1,5 @@ +from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig +from fastvideo.configs.models.dits.stepvideo import StepVideoConfig +from fastvideo.configs.models.dits.wanvideo import WanVideoConfig + +__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"] diff --git a/worldcompass/fastvideo/configs/models/dits/base.py b/worldcompass/fastvideo/configs/models/dits/base.py new file mode 100644 index 0000000000000000000000000000000000000000..12fb603f43dbe0deaf8dd72e8ccbb4a51df4ae58 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/dits/base.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.base import ArchConfig, ModelConfig +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.platforms import AttentionBackendEnum + + +@dataclass +class DiTArchConfig(ArchConfig): + _fsdp_shard_conditions: list = field(default_factory=list) + _compile_conditions: list = field(default_factory=list) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.SAGE_ATTN, + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + ) + + hidden_size: int = 0 + num_attention_heads: int = 0 + num_channels_latents: int = 0 + exclude_lora_layers: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if not self._compile_conditions: + self._compile_conditions = self._fsdp_shard_conditions.copy() + + +@dataclass +class DiTConfig(ModelConfig): + arch_config: DiTArchConfig = field(default_factory=DiTArchConfig) + + # FastVideoDiT-specific parameters + prefix: str = "" + quant_config: QuantizationConfig | None = None + + @staticmethod + def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any: + """Add CLI arguments for DiTConfig fields.""" + parser.add_argument( + f"--{prefix}.prefix", + type=str, + dest=f"{prefix.replace('-', '_')}.prefix", + default=DiTConfig.prefix, + help="Prefix for the DiT model", + ) + + parser.add_argument( + f"--{prefix}.quant-config", + type=str, + dest=f"{prefix.replace('-', '_')}.quant_config", + default=None, + help="Quantization configuration for the DiT model", + ) + + return parser diff --git a/worldcompass/fastvideo/configs/models/dits/hunyuanvideo.py b/worldcompass/fastvideo/configs/models/dits/hunyuanvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..057dbe2311654fac47b99f8627e285457168d3bb --- /dev/null +++ b/worldcompass/fastvideo/configs/models/dits/hunyuanvideo.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_double_block(n: str, m) -> bool: + return "double" in n and str.isdigit(n.split(".")[-1]) + + +def is_single_block(n: str, m) -> bool: + return "single" in n and str.isdigit(n.split(".")[-1]) + + +def is_refiner_block(n: str, m) -> bool: + return "refiner" in n and str.isdigit(n.split(".")[-1]) + + +def is_txt_in(n: str, m) -> bool: + return n.split(".")[-1] == "txt_in" + + +@dataclass +class HunyuanVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [ + is_double_block, + is_single_block, + is_refiner_block, + ] + ) + + _compile_conditions: list = field( + default_factory=lambda: [is_double_block, is_single_block, is_txt_in] + ) + + param_names_mapping: dict = field( + default_factory=lambda: { + # 1. context_embedder.time_text_embed submodules (specific rules, applied first): + r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1", + r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1", + r"^context_embedder\.proj_in\.(.*)$": r"txt_in.input_embedder.\1", + r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1", + r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$": r"txt_in.refiner_blocks.\1.norm1.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$": r"txt_in.refiner_blocks.\1.norm2.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( + r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", + 0, + 3, + ), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( + r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", + 1, + 3, + ), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( + r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", + 2, + 3, + ), + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"txt_in.refiner_blocks.\1.self_attn_proj.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2", + r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", + # 3. x_embedder mapping: + r"^x_embedder\.proj\.(.*)$": r"img_in.proj.\1", + # 4. Top-level time_text_embed mappings: + r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"time_in.mlp.fc_in.\1", + r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"time_in.mlp.fc_out.\1", + r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$": r"guidance_in.mlp.fc_in.\1", + r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$": r"guidance_in.mlp.fc_out.\1", + r"^time_text_embed\.text_embedder\.linear_1\.(.*)$": r"vector_in.fc_in.\1", + r"^time_text_embed\.text_embedder\.linear_2\.(.*)$": r"vector_in.fc_out.\1", + # 5. transformer_blocks mapping: + r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$": r"double_blocks.\1.img_mod.linear.\2", + r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$": r"double_blocks.\1.txt_mod.linear.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( + r"double_blocks.\1.img_attn_qkv.\2", + 0, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( + r"double_blocks.\1.img_attn_qkv.\2", + 1, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( + r"double_blocks.\1.img_attn_qkv.\2", + 2, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": ( + r"double_blocks.\1.txt_attn_qkv.\2", + 0, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": ( + r"double_blocks.\1.txt_attn_qkv.\2", + 1, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": ( + r"double_blocks.\1.txt_attn_qkv.\2", + 2, + 3, + ), + r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"double_blocks.\1.img_attn_proj.\2", + # Corrected: merge attn.to_add_out into the main projection. + r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2", + r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2", + r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2", + # 6. single_transformer_blocks mapping: + r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"single_blocks.\1.q_norm.\2", + r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"single_blocks.\1.k_norm.\2", + r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( + r"single_blocks.\1.linear1.\2", + 0, + 4, + ), + r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( + r"single_blocks.\1.linear1.\2", + 1, + 4, + ), + r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( + r"single_blocks.\1.linear1.\2", + 2, + 4, + ), + r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": ( + r"single_blocks.\1.linear1.\2", + 3, + 4, + ), + # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch. + r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$": r"single_blocks.\1.linear2.\2", + r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$": r"single_blocks.\1.modulation.linear.\2", + # 7. Final layers mapping: + r"^norm_out\.linear\.(.*)$": r"final_layer.adaLN_modulation.linear.\1", + r"^proj_out\.(.*)$": r"final_layer.linear.\1", + } + ) + + # Reverse mapping for saving checkpoints: custom -> hf + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + patch_size: int = 2 + patch_size_t: int = 1 + in_channels: int = 16 + out_channels: int = 16 + num_attention_heads: int = 24 + attention_head_dim: int = 128 + mlp_ratio: float = 4.0 + num_layers: int = 20 + num_single_layers: int = 40 + num_refiner_layers: int = 2 + rope_axes_dim: tuple[int, int, int] = (16, 56, 56) + guidance_embeds: bool = False + dtype: torch.dtype | None = None + text_embed_dim: int = 4096 + pooled_projection_dim: int = 768 + rope_theta: int = 256 + qk_norm: str = "rms_norm" + exclude_lora_layers: list[str] = field( + default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"] + ) + + def __post_init__(self): + super().__post_init__() + self.hidden_size: int = ( + self.attention_head_dim * self.num_attention_heads + ) + self.num_channels_latents: int = self.in_channels + + +@dataclass +class HunyuanVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig) + + prefix: str = "Hunyuan" diff --git a/worldcompass/fastvideo/configs/models/dits/stepvideo.py b/worldcompass/fastvideo/configs/models/dits/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..92815ec5936a7f6bc1a788a1a6d95289ae1a0df5 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/dits/stepvideo.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class StepVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [ + lambda n, m: "transformer_blocks" in n + and n.split(".")[-1].isdigit() + ] + ) + + param_names_mapping: dict = field( + default_factory=lambda: { + # transformer block + r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$": r"transformer_blocks.\1.norm1.norm.\2", + r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$": r"transformer_blocks.\1.norm2.norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$": r"transformer_blocks.\1.ff.fc_in.weight", + r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$": r"transformer_blocks.\1.ff.fc_out.weight", + # adanorm block + r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$": r"adaln_single.emb.mlp.fc_in.\1", + r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$": r"adaln_single.emb.mlp.fc_out.\1", + # caption projection + r"^caption_projection\.linear_1\.(weight|bias)$": r"caption_projection.fc_in.\1", + r"^caption_projection\.linear_2\.(weight|bias)$": r"caption_projection.fc_out.\1", + } + ) + + num_attention_heads: int = 48 + attention_head_dim: int = 128 + in_channels: int = 64 + out_channels: int | None = 64 + num_layers: int = 48 + dropout: float = 0.0 + patch_size: int = 1 + norm_type: str = "ada_norm_single" + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + caption_channels: int | list[int] | tuple[int, ...] | None = field( + default_factory=lambda: [6144, 1024] + ) + attention_type: str | None = "torch" + use_additional_conditions: bool | None = False + exclude_lora_layers: list[str] = field(default_factory=lambda: []) + + def __post_init__(self): + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.out_channels = ( + self.in_channels if self.out_channels is None else self.out_channels + ) + self.num_channels_latents = self.out_channels + + +@dataclass +class StepVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig) + + prefix: str = "StepVideo" diff --git a/worldcompass/fastvideo/configs/models/dits/wanvideo.py b/worldcompass/fastvideo/configs/models/dits/wanvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..778e5cf9ba98726f68ef0ea97c9080aff3dc95d8 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/dits/wanvideo.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class WanVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1", + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": r"condition_embedder.image_embedder.ff.fc_in.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": r"condition_embedder.image_embedder.ff.fc_out.\1", + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", + } + ) + + # Reverse mapping for saving checkpoints: custom -> hf + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Some LoRA adapters use the original official layer names instead of hf layer names, + # so apply this before the param_names_mapping + lora_param_names_mapping: dict = field( + default_factory=lambda: { + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + } + ) + + patch_size: tuple[int, int, int] = (1, 2, 2) + text_len = 512 + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + image_dim: int | None = None + added_kv_proj_dim: int | None = None + rope_max_seq_len: int = 1024 + pos_embed_seq_len: int | None = None + exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class WanVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig) + + prefix: str = "Wan" diff --git a/worldcompass/fastvideo/configs/models/encoders/__init__.py b/worldcompass/fastvideo/configs/models/encoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ecc7ef9ff9b582f3f6d08905569f95ad7eb331 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/encoders/__init__.py @@ -0,0 +1,23 @@ +from fastvideo.configs.models.encoders.base import ( + BaseEncoderOutput, + EncoderConfig, + ImageEncoderConfig, + TextEncoderConfig, +) +from fastvideo.configs.models.encoders.clip import ( + CLIPTextConfig, + CLIPVisionConfig, +) +from fastvideo.configs.models.encoders.llama import LlamaConfig +from fastvideo.configs.models.encoders.t5 import T5Config + +__all__ = [ + "EncoderConfig", + "TextEncoderConfig", + "ImageEncoderConfig", + "BaseEncoderOutput", + "CLIPTextConfig", + "CLIPVisionConfig", + "LlamaConfig", + "T5Config", +] diff --git a/worldcompass/fastvideo/configs/models/encoders/base.py b/worldcompass/fastvideo/configs/models/encoders/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1cf11954836709b94ac8e8d01e6eaf5edef68781 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/encoders/base.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Any + +import torch + +from fastvideo.configs.models.base import ArchConfig, ModelConfig +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.platforms import AttentionBackendEnum + + +@dataclass +class EncoderArchConfig(ArchConfig): + architectures: list[str] = field(default_factory=lambda: []) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + output_hidden_states: bool = False + use_return_dict: bool = True + + +@dataclass +class TextEncoderArchConfig(EncoderArchConfig): + vocab_size: int = 0 + hidden_size: int = 0 + num_hidden_layers: int = 0 + num_attention_heads: int = 0 + pad_token_id: int = 0 + eos_token_id: int = 0 + text_len: int = 0 + hidden_state_skip_layer: int = 0 + decoder_start_token_id: int = 0 + output_past: bool = True + scalable_attention: bool = True + tie_word_embeddings: bool = False + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=list + ) # mapping from huggingface weight names to custom names + tokenizer_kwargs: dict[str, Any] = field(default_factory=dict) + _fsdp_shard_conditions: list = field(default_factory=lambda: []) + + def __post_init__(self) -> None: + self.tokenizer_kwargs = { + "truncation": True, + "max_length": self.text_len, + "return_tensors": "pt", + } + + +@dataclass +class ImageEncoderArchConfig(EncoderArchConfig): + pass + + +@dataclass +class BaseEncoderOutput: + last_hidden_state: torch.FloatTensor | None = None + pooler_output: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + attention_mask: torch.Tensor | None = None + + +@dataclass +class EncoderConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=EncoderArchConfig) + + prefix: str = "" + quant_config: QuantizationConfig | None = None + lora_config: Any | None = None + + +@dataclass +class TextEncoderConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig) + + +@dataclass +class ImageEncoderConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig) diff --git a/worldcompass/fastvideo/configs/models/encoders/clip.py b/worldcompass/fastvideo/configs/models/encoders/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..617ae24b7d6c0976250f3c6084128714bf1e3d45 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/encoders/clip.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import ( + ImageEncoderArchConfig, + ImageEncoderConfig, + TextEncoderArchConfig, + TextEncoderConfig, +) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embeddings") + + +@dataclass +class CLIPTextArchConfig(TextEncoderArchConfig): + vocab_size: int = 49408 + hidden_size: int = 512 + intermediate_size: int = 2048 + projection_dim: int = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 8 + max_position_embeddings: int = 77 + hidden_act: str = "quick_gelu" + layer_norm_eps: float = 1e-5 + dropout: float = 0.0 + attention_dropout: float = 0.0 + initializer_range: float = 0.02 + initializer_factor: float = 1.0 + pad_token_id: int = 1 + bos_token_id: int = 49406 + eos_token_id: int = 49407 + text_len: int = 77 + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + ) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_transformer_layer, _is_embeddings] + ) + + +@dataclass +class CLIPVisionArchConfig(ImageEncoderArchConfig): + hidden_size: int = 768 + intermediate_size: int = 3072 + projection_dim: int = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int = 224 + patch_size: int = 32 + hidden_act: str = "quick_gelu" + layer_norm_eps: float = 1e-5 + dropout: float = 0.0 + attention_dropout: float = 0.0 + initializer_range: float = 0.02 + initializer_factor: float = 1.0 + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + ) + + +@dataclass +class CLIPTextConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field( + default_factory=CLIPTextArchConfig + ) + + num_hidden_layers_override: int | None = None + require_post_norm: bool | None = None + prefix: str = "clip" + + +@dataclass +class CLIPVisionConfig(ImageEncoderConfig): + arch_config: ImageEncoderArchConfig = field( + default_factory=CLIPVisionArchConfig + ) + + num_hidden_layers_override: int | None = None + require_post_norm: bool | None = None + prefix: str = "clip" diff --git a/worldcompass/fastvideo/configs/models/encoders/llama.py b/worldcompass/fastvideo/configs/models/encoders/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..314ab823e4d0f0f5596f888f444b63d6f048afe6 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/encoders/llama.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m) -> bool: + return n.endswith("norm") + + +@dataclass +class LlamaArchConfig(TextEncoderArchConfig): + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int = 0 + bos_token_id: int = 1 + eos_token_id: int = 2 + pretraining_tp: int = 1 + tie_word_embeddings: bool = False + rope_theta: float = 10000.0 + rope_scaling: float | None = None + attention_bias: bool = False + attention_dropout: float = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + hidden_state_skip_layer: int = 2 + text_len: int = 256 + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), # type: ignore + (".gate_up_proj", ".up_proj", 1), # type: ignore + ] + ) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [ + _is_transformer_layer, + _is_embeddings, + _is_final_norm, + ] + ) + + +@dataclass +class LlamaConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig) + + prefix: str = "llama" diff --git a/worldcompass/fastvideo/configs/models/encoders/t5.py b/worldcompass/fastvideo/configs/models/encoders/t5.py new file mode 100644 index 0000000000000000000000000000000000000000..0abee14619f136a60160608fc6ecd688f55c1a4c --- /dev/null +++ b/worldcompass/fastvideo/configs/models/encoders/t5.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +def _is_transformer_layer(n: str, m) -> bool: + return "block" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("shared") + + +def _is_final_layernorm(n: str, m) -> bool: + return n.endswith("final_layer_norm") + + +@dataclass +class T5ArchConfig(TextEncoderArchConfig): + vocab_size: int = 32128 + d_model: int = 512 + d_kv: int = 64 + d_ff: int = 2048 + num_layers: int = 6 + num_decoder_layers: int | None = None + num_heads: int = 8 + relative_attention_num_buckets: int = 32 + relative_attention_max_distance: int = 128 + dropout_rate: float = 0.1 + layer_norm_epsilon: float = 1e-6 + initializer_factor: float = 1.0 + feed_forward_proj: str = "relu" + dense_act_fn: str = "" + is_gated_act: bool = False + is_encoder_decoder: bool = True + use_cache: bool = True + pad_token_id: int = 0 + eos_token_id: int = 1 + classifier_dropout: float = 0.0 + text_len: int = 512 + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q", "q"), + (".qkv_proj", ".k", "k"), + (".qkv_proj", ".v", "v"), + ] + ) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [ + _is_transformer_layer, + _is_embeddings, + _is_final_layernorm, + ] + ) + + # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py + def __post_init__(self): + super().__post_init__() + act_info = self.feed_forward_proj.split("-") + self.dense_act_fn: str = act_info[-1] + self.is_gated_act: bool = act_info[0] == "gated" + if self.feed_forward_proj == "gated-gelu": + self.dense_act_fn = "gelu_new" + + self.tokenizer_kwargs = { + "padding": "max_length", + "truncation": True, + "max_length": self.text_len, + "add_special_tokens": True, + "return_attention_mask": True, + "return_tensors": "pt", + } + + +@dataclass +class T5Config(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig) + + prefix: str = "t5" diff --git a/worldcompass/fastvideo/configs/models/vaes/__init__.py b/worldcompass/fastvideo/configs/models/vaes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..700c8de1bcd7dbfd1f93e1b9a0acd7680a8f5b58 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/vaes/__init__.py @@ -0,0 +1,9 @@ +from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig +from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig +from fastvideo.configs.models.vaes.wanvae import WanVAEConfig + +__all__ = [ + "HunyuanVAEConfig", + "WanVAEConfig", + "StepVideoVAEConfig", +] diff --git a/worldcompass/fastvideo/configs/models/vaes/base.py b/worldcompass/fastvideo/configs/models/vaes/base.py new file mode 100644 index 0000000000000000000000000000000000000000..10c0bd5884c795a45353f3893498b4309c9ce04e --- /dev/null +++ b/worldcompass/fastvideo/configs/models/vaes/base.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +import argparse +import dataclasses +from dataclasses import dataclass, field +from typing import Any + +import torch + +from fastvideo.configs.models.base import ArchConfig, ModelConfig +from fastvideo.utils import StoreBoolean + + +@dataclass +class VAEArchConfig(ArchConfig): + scaling_factor: float | torch.Tensor = 0 + + temporal_compression_ratio: int = 4 + spatial_compression_ratio: int = 8 + + +@dataclass +class VAEConfig(ModelConfig): + arch_config: VAEArchConfig = field(default_factory=VAEArchConfig) + + # FastVideoVAE-specific parameters + load_encoder: bool = True + load_decoder: bool = True + + tile_sample_min_height: int = 256 + tile_sample_min_width: int = 256 + tile_sample_min_num_frames: int = 16 + tile_sample_stride_height: int = 192 + tile_sample_stride_width: int = 192 + tile_sample_stride_num_frames: int = 12 + blend_num_frames: int = 0 + + use_tiling: bool = True + use_temporal_tiling: bool = True + use_parallel_tiling: bool = True + use_temporal_scaling_frames: bool = True + + def __post_init__(self): + self.blend_num_frames = ( + self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + ) + + @staticmethod + def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any: + """Add CLI arguments for VAEConfig fields.""" + parser.add_argument( + f"--{prefix}.load-encoder", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.load_encoder", + default=VAEConfig.load_encoder, + help="Whether to load the VAE encoder", + ) + parser.add_argument( + f"--{prefix}.load-decoder", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.load_decoder", + default=VAEConfig.load_decoder, + help="Whether to load the VAE decoder", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-height", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_height", + default=VAEConfig.tile_sample_min_height, + help="Minimum height for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-width", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_width", + default=VAEConfig.tile_sample_min_width, + help="Minimum width for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-min-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames", + default=VAEConfig.tile_sample_min_num_frames, + help="Minimum number of frames for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-height", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height", + default=VAEConfig.tile_sample_stride_height, + help="Stride height for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-width", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width", + default=VAEConfig.tile_sample_stride_width, + help="Stride width for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.tile-sample-stride-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames", + default=VAEConfig.tile_sample_stride_num_frames, + help="Stride number of frames for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.blend-num-frames", + type=int, + dest=f"{prefix.replace('-', '_')}.blend_num_frames", + default=VAEConfig.blend_num_frames, + help="Number of frames to blend for VAE tile sampling", + ) + parser.add_argument( + f"--{prefix}.use-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_tiling", + default=VAEConfig.use_tiling, + help="Whether to use tiling for VAE", + ) + parser.add_argument( + f"--{prefix}.use-temporal-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_temporal_tiling", + default=VAEConfig.use_temporal_tiling, + help="Whether to use temporal tiling for VAE", + ) + parser.add_argument( + f"--{prefix}.use-parallel-tiling", + action=StoreBoolean, + dest=f"{prefix.replace('-', '_')}.use_parallel_tiling", + default=VAEConfig.use_parallel_tiling, + help="Whether to use parallel tiling for VAE", + ) + + return parser + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig": + kwargs = {} + for attr in dataclasses.fields(cls): + value = getattr(args, attr.name, None) + if value is not None: + kwargs[attr.name] = value + return cls(**kwargs) diff --git a/worldcompass/fastvideo/configs/models/vaes/hunyuanvae.py b/worldcompass/fastvideo/configs/models/vaes/hunyuanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..178b18ffea045cb50c68d6ff4ff30a841039fce2 --- /dev/null +++ b/worldcompass/fastvideo/configs/models/vaes/hunyuanvae.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class HunyuanVAEArchConfig(VAEArchConfig): + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + down_block_types: tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ) + up_block_types: tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ) + block_out_channels: tuple[int, ...] = (128, 256, 512, 512) + layers_per_block: int = 2 + act_fn: str = "silu" + norm_num_groups: int = 32 + scaling_factor: float = 0.476986 + spatial_compression_ratio: int = 8 + temporal_compression_ratio: int = 4 + mid_block_add_attention: bool = True + + def __post_init__(self): + self.spatial_compression_ratio: int = 2 ** ( + len(self.block_out_channels) - 1 + ) + + +@dataclass +class HunyuanVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig) diff --git a/worldcompass/fastvideo/configs/models/vaes/stepvideovae.py b/worldcompass/fastvideo/configs/models/vaes/stepvideovae.py new file mode 100644 index 0000000000000000000000000000000000000000..07da2512efa6dce206fbf632de393cbbff4adcbb --- /dev/null +++ b/worldcompass/fastvideo/configs/models/vaes/stepvideovae.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class StepVideoVAEArchConfig(VAEArchConfig): + in_channels: int = 3 + out_channels: int = 3 + z_channels: int = 64 + num_res_blocks: int = 2 + version: int = 2 + frame_len: int = 17 + world_size: int = 1 + + spatial_compression_ratio: int = 16 + temporal_compression_ratio: int = 8 + + scaling_factor: float = 1.0 + + +@dataclass +class StepVideoVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig) + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + use_temporal_scaling_frames: bool = False diff --git a/worldcompass/fastvideo/configs/models/vaes/wanvae.py b/worldcompass/fastvideo/configs/models/vaes/wanvae.py new file mode 100644 index 0000000000000000000000000000000000000000..ce56c86078e21d21599ce62156468e35a275743d --- /dev/null +++ b/worldcompass/fastvideo/configs/models/vaes/wanvae.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class WanVAEArchConfig(VAEArchConfig): + base_dim: int = 96 + decoder_base_dim: int | None = None + z_dim: int = 16 + dim_mult: tuple[int, ...] = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple[float, ...] = () + temperal_downsample: tuple[bool, ...] = (False, True, True) + dropout: float = 0.0 + latents_mean: tuple[float, ...] = ( + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ) + latents_std: tuple[float, ...] = ( + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ) + is_residual: bool = False + in_channels: int = 3 + out_channels: int = 3 + patch_size: int | None = None + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 8 + clip_output: bool = True + + def __post_init__(self): + self.scaling_factor: torch.tensor = 1.0 / torch.tensor( + self.latents_std + ).view(1, self.z_dim, 1, 1, 1) + self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view( + 1, self.z_dim, 1, 1, 1 + ) + self.temporal_compression_ratio = self.scale_factor_temporal + self.spatial_compression_ratio = self.scale_factor_spatial + + +@dataclass +class WanVAEConfig(VAEConfig): + arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig) + use_feature_cache: bool = True + + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def __post_init__(self): + self.blend_num_frames = ( + self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + ) * 2 diff --git a/worldcompass/fastvideo/configs/pipelines/__init__.py b/worldcompass/fastvideo/configs/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4869820bd5fe2d49531e57930790c1220165038d --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/__init__.py @@ -0,0 +1,28 @@ +from fastvideo.configs.pipelines.base import ( + PipelineConfig, + SlidingTileAttnConfig, +) +from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig +from fastvideo.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name, +) +from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig +from fastvideo.configs.pipelines.wan import ( + WanI2V480PConfig, + WanI2V720PConfig, + WanT2V480PConfig, + WanT2V720PConfig, +) + +__all__ = [ + "HunyuanConfig", + "FastHunyuanConfig", + "PipelineConfig", + "SlidingTileAttnConfig", + "WanT2V480PConfig", + "WanI2V480PConfig", + "WanT2V720PConfig", + "WanI2V720PConfig", + "StepVideoT2VConfig", + "get_pipeline_config_cls_from_name", +] diff --git a/worldcompass/fastvideo/configs/pipelines/base.py b/worldcompass/fastvideo/configs/pipelines/base.py new file mode 100644 index 0000000000000000000000000000000000000000..bcb3dff98d958c435be2c3fbfe49edadd3a7571a --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/base.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +from collections.abc import Callable +from dataclasses import asdict, dataclass, field, fields +from enum import Enum +from typing import Any, cast + +import torch + +from fastvideo.configs.models import ( + DiTConfig, + EncoderConfig, + ModelConfig, + VAEConfig, +) +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.utils import update_config_from_args +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict + +logger = init_logger(__name__) + + +class STA_Mode(str, Enum): + """STA (Sliding Tile Attention) modes.""" + + STA_INFERENCE = "STA_inference" + STA_SEARCHING = "STA_searching" + STA_TUNING = "STA_tuning" + STA_TUNING_CFG = "STA_tuning_cfg" + NONE = None + + +def preprocess_text(prompt: str) -> str: + return prompt + + +def postprocess_text(output: BaseEncoderOutput) -> torch.tensor: + raise NotImplementedError + + +# config for a single pipeline +@dataclass +class PipelineConfig: + """Base configuration for all pipeline architectures.""" + + model_path: str = "" + pipeline_config_path: str | None = None + + # Video generation parameters + embedded_cfg_scale: float = 6.0 + flow_shift: float | None = None + disable_autocast: bool = False + + # Model configuration + dit_config: DiTConfig = field(default_factory=DiTConfig) + dit_precision: str = "bf16" + + # VAE configuration + vae_config: VAEConfig = field(default_factory=VAEConfig) + vae_precision: str = "fp32" + vae_tiling: bool = True + vae_sp: bool = True + + # Image encoder configuration + image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig) + image_encoder_precision: str = "fp32" + + # Text encoder configuration + DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32",) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (EncoderConfig(),) + ) + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp32",) + ) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (preprocess_text,) + ) + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput], torch.tensor], ... + ] = field(default_factory=lambda: (postprocess_text,)) + + # StepVideo specific parameters + pos_magic: str | None = None + neg_magic: str | None = None + timesteps_scale: bool | None = None + + # STA (Sliding Tile Attention) parameters + mask_strategy_file_path: str | None = None + STA_mode: STA_Mode = STA_Mode.STA_INFERENCE + skip_time_steps: int = 15 + + # DMD parameters + dmd_denoising_steps: list[int] | None = field(default=None) + + # Compilation + # enable_torch_compile: bool = False + + @staticmethod + def add_cli_args( + parser: FlexibleArgumentParser, prefix: str = "" + ) -> FlexibleArgumentParser: + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + + # model_path will be conflicting with the model_path in FastVideoArgs, + # so we add it separately if prefix is not empty + if prefix_with_dot != "": + parser.add_argument( + f"--{prefix_with_dot}model-path", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}model_path", + default=PipelineConfig.model_path, + help="Path to the pretrained model", + ) + + parser.add_argument( + f"--{prefix_with_dot}pipeline-config-path", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}pipeline_config_path", + default=PipelineConfig.pipeline_config_path, + help="Path to the pipeline config", + ) + parser.add_argument( + f"--{prefix_with_dot}embedded-cfg-scale", + type=float, + dest=f"{prefix_with_dot.replace('-', '_')}embedded_cfg_scale", + default=PipelineConfig.embedded_cfg_scale, + help="Embedded CFG scale", + ) + parser.add_argument( + f"--{prefix_with_dot}flow-shift", + type=float, + dest=f"{prefix_with_dot.replace('-', '_')}flow_shift", + default=PipelineConfig.flow_shift, + help="Flow shift parameter", + ) + + # DiT configuration + parser.add_argument( + f"--{prefix_with_dot}dit-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}dit_precision", + default=PipelineConfig.dit_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for the DiT model", + ) + + # VAE configuration + parser.add_argument( + f"--{prefix_with_dot}vae-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}vae_precision", + default=PipelineConfig.vae_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for VAE", + ) + parser.add_argument( + f"--{prefix_with_dot}vae-tiling", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}vae_tiling", + default=PipelineConfig.vae_tiling, + help="Enable VAE tiling", + ) + parser.add_argument( + f"--{prefix_with_dot}vae-sp", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}vae_sp", + help="Enable VAE spatial parallelism", + ) + + # Text encoder configuration + parser.add_argument( + f"--{prefix_with_dot}text-encoder-precisions", + nargs="+", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}text_encoder_precisions", + default=PipelineConfig.DEFAULT_TEXT_ENCODER_PRECISIONS, + choices=["fp32", "fp16", "bf16"], + help="Precision for each text encoder", + ) + + # Image encoder configuration + parser.add_argument( + f"--{prefix_with_dot}image-encoder-precision", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}image_encoder_precision", + default=PipelineConfig.image_encoder_precision, + choices=["fp32", "fp16", "bf16"], + help="Precision for image encoder", + ) + parser.add_argument( + f"--{prefix_with_dot}pos_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}pos_magic", + default=PipelineConfig.pos_magic, + help="Positive magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}neg_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}neg_magic", + default=PipelineConfig.neg_magic, + help="Negative magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}timesteps_scale", + type=bool, + dest=f"{prefix_with_dot.replace('-', '_')}timesteps_scale", + default=PipelineConfig.timesteps_scale, + help="Bool for applying scheduler scale in set_timesteps, used in stepvideo", + ) + + # DMD parameters + parser.add_argument( + f"--{prefix_with_dot}dmd-denoising-steps", + type=parse_int_list, + default=PipelineConfig.dmd_denoising_steps, + help="Comma-separated list of denoising steps (e.g., '1000,757,522')", + ) + + # Add VAE configuration arguments + from fastvideo.configs.models.vaes.base import VAEConfig + + VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") + + # Add DiT configuration arguments + from fastvideo.configs.models.dits.base import DiTConfig + + DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config") + + return parser + + def update_config_from_dict( + self, args: dict[str, Any], prefix: str = "" + ) -> None: + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + update_config_from_args(self, args, prefix, pop_args=True) + update_config_from_args( + self.vae_config, args, f"{prefix_with_dot}vae_config", pop_args=True + ) + update_config_from_args( + self.dit_config, args, f"{prefix_with_dot}dit_config", pop_args=True + ) + + @classmethod + def from_pretrained(cls, model_path: str) -> "PipelineConfig": + """Use the pipeline class setting from model_path to match the pipeline config.""" + from fastvideo.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name, + ) + + pipeline_config_cls = get_pipeline_config_cls_from_name(model_path) + + return cast(PipelineConfig, pipeline_config_cls(model_path=model_path)) + + @classmethod + def from_kwargs( + cls, kwargs: dict[str, Any], config_cli_prefix: str = "" + ) -> "PipelineConfig": + """Load PipelineConfig from kwargs Dictionary. + + kwargs: dictionary of kwargs + config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance + """ + from fastvideo.configs.pipelines.registry import ( + get_pipeline_config_cls_from_name, + ) + + prefix_with_dot = ( + f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else "" + ) + model_path: str | None = kwargs.get( + prefix_with_dot + "model_path", None + ) or kwargs.get("model_path") + pipeline_config_or_path: ( + str | PipelineConfig | dict[str, Any] | None + ) = kwargs.get(prefix_with_dot + "pipeline_config", None) or kwargs.get( + "pipeline_config" + ) + if model_path is None: + raise ValueError("model_path is required in kwargs") + + # 1. Get the pipeline config class from the registry + pipeline_config_cls = get_pipeline_config_cls_from_name(model_path) + + # 2. Instantiate PipelineConfig + if pipeline_config_cls is None: + logger.warning( + "Couldn't find pipeline config for %s. Using the default pipeline config.", + model_path, + ) + pipeline_config = cls() + else: + pipeline_config = pipeline_config_cls() + + # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided + if isinstance(pipeline_config_or_path, str): + pipeline_config.load_from_json(pipeline_config_or_path) + kwargs[prefix_with_dot + "pipeline_config_path"] = ( + pipeline_config_or_path + ) + elif isinstance(pipeline_config_or_path, PipelineConfig): + pipeline_config = pipeline_config_or_path + elif isinstance(pipeline_config_or_path, dict): + pipeline_config.update_pipeline_config(pipeline_config_or_path) + + # 4. Update PipelineConfig from CLI arguments if provided + kwargs[prefix_with_dot + "model_path"] = model_path + pipeline_config.update_config_from_dict(kwargs, config_cli_prefix) + return pipeline_config + + def check_pipeline_config(self) -> None: + if self.vae_sp and not self.vae_tiling: + raise ValueError( + "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True." + ) + + if len(self.text_encoder_configs) != len(self.text_encoder_precisions): + raise ValueError( + f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})" + ) + + if len(self.text_encoder_configs) != len(self.preprocess_text_funcs): + raise ValueError( + f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" + ) + + if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs): + raise ValueError( + f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" + ) + + def dump_to_json(self, file_path: str): + output_dict = shallow_asdict(self) + del_keys = [] + for key, value in output_dict.items(): + if isinstance(value, ModelConfig): + model_dict = asdict(value) + # Model Arch Config should be hidden away from the users + model_dict.pop("arch_config") + output_dict[key] = model_dict + elif isinstance(value, tuple) and all( + isinstance(v, ModelConfig) for v in value + ): + model_dicts = [] + for v in value: + model_dict = asdict(v) + # Model Arch Config should be hidden away from the users + model_dict.pop("arch_config") + model_dicts.append(model_dict) + output_dict[key] = model_dicts + elif isinstance(value, tuple) and all(callable(f) for f in value): + # Skip dumping functions + del_keys.append(key) + + for key in del_keys: + output_dict.pop(key, None) + + with open(file_path, "w") as f: + json.dump(output_dict, f, indent=2) + + def load_from_json(self, file_path: str): + with open(file_path) as f: + input_pipeline_dict = json.load(f) + self.update_pipeline_config(input_pipeline_dict) + + def update_pipeline_config( + self, source_pipeline_dict: dict[str, Any] + ) -> None: + for f in fields(self): + key = f.name + if key in source_pipeline_dict: + current_value = getattr(self, key) + new_value = source_pipeline_dict[key] + + # If it's a nested ModelConfig, update it recursively + if isinstance(current_value, ModelConfig): + current_value.update_model_config(new_value) + elif isinstance(current_value, tuple) and all( + isinstance(v, ModelConfig) for v in current_value + ): + assert len(current_value) == len( + new_value + ), "Users shouldn't delete or add text encoder config objects in your json" + for target_config, source_config in zip( + current_value, new_value, strict=True + ): + target_config.update_model_config(source_config) + else: + setattr(self, key, new_value) + + if hasattr(self, "__post_init__"): + self.__post_init__() + + +@dataclass +class SlidingTileAttnConfig(PipelineConfig): + """Configuration for sliding tile attention.""" + + # Override any BaseConfig defaults as needed + # Add sliding tile specific parameters + window_size: int = 16 + stride: int = 8 + + # You can provide custom defaults for inherited fields + height: int = 576 + width: int = 1024 + + # Additional configuration specific to sliding tile attention + pad_to_square: bool = False + use_overlap_optimization: bool = True + + +def parse_int_list(value: str) -> list[int]: + """Parse a comma-separated string of integers into a list.""" + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] diff --git a/worldcompass/fastvideo/configs/pipelines/hunyuan.py b/worldcompass/fastvideo/configs/pipelines/hunyuan.py new file mode 100644 index 0000000000000000000000000000000000000000..05f34f2dddf97158ef14de7e151659a127af91ca --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/hunyuan.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TypedDict + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits import HunyuanVideoConfig +from fastvideo.configs.models.encoders import ( + BaseEncoderOutput, + CLIPTextConfig, + LlamaConfig, +) +from fastvideo.configs.models.vaes import HunyuanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" +) + + +class PromptTemplate(TypedDict): + template: str + crop_start: int + + +prompt_template_video: PromptTemplate = { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, +} + + +def llama_preprocess_text(prompt: str) -> str: + return prompt_template_video["template"].format(prompt) + + +def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + hidden_state_skip_layer = 2 + assert outputs.hidden_states is not None + hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states + last_hidden_state: torch.tensor = hidden_states[ + -(hidden_state_skip_layer + 1) + ] + crop_start = prompt_template_video.get("crop_start", -1) + last_hidden_state = last_hidden_state[:, crop_start:] + return last_hidden_state + + +def clip_preprocess_text(prompt: str) -> str: + return prompt + + +def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + pooler_output: torch.tensor = outputs.pooler_output + return pooler_output + + +@dataclass +class HunyuanConfig(PipelineConfig): + """Base configuration for HunYuan pipeline architecture.""" + + # HunyuanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig) + # Denoising stage + embedded_cfg_scale: int = 6 + flow_shift: int = 7 + + # Text encoding stage + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (LlamaConfig(), CLIPTextConfig()) + ) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (llama_preprocess_text, clip_preprocess_text) + ) + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput], torch.tensor], ... + ] = field( + default_factory=lambda: (llama_postprocess_text, clip_postprocess_text) + ) + + # Precision for each component + dit_precision: str = "bf16" + vae_precision: str = "fp16" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp16", "fp16") + ) + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class FastHunyuanConfig(HunyuanConfig): + """Configuration specifically optimized for FastHunyuan weights.""" + + # Override HunyuanConfig defaults + flow_shift: int = 17 + + # No need to re-specify guidance_scale or embedded_cfg_scale as they + # already have the desired values from HunyuanConfig diff --git a/worldcompass/fastvideo/configs/pipelines/registry.py b/worldcompass/fastvideo/configs/pipelines/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..d8526fec041f18d088d1841889adc7b84f03684c --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/registry.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry for pipeline weight-specific configurations.""" + +import os +from collections.abc import Callable + +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig +from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig +from fastvideo.configs.pipelines.wan import ( + FastWan2_1_T2V_480P_Config, + FastWan2_2_TI2V_5B_Config, + WanI2V480PConfig, + WanI2V720PConfig, + WanT2V480PConfig, + WanT2V720PConfig, +) +from fastvideo.logger import init_logger +from fastvideo.utils import ( + maybe_download_model_index, + verify_model_config_and_directory, +) + +logger = init_logger(__name__) + +# Registry maps specific model weights to their config classes +PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { + "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig, + "hunyuanvideo-community/HunyuanVideo": HunyuanConfig, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig, + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig, + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWan2_1_T2V_480P_Config, + "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers": FastWan2_1_T2V_480P_Config, + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig, + "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig, + # Add other specific weight variants +} + +# For determining pipeline type from model ID +PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { + "hunyuan": lambda id: "hunyuan" in id.lower(), + "wanpipeline": lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), + "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), + "stepvideo": lambda id: "stepvideo" in id.lower(), + # Add other pipeline architecture detectors +} + +# Fallback configs when exact match isn't found but architecture is detected +PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { + "hunyuan": HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant + "wanpipeline": WanT2V480PConfig, # Base Wan config as fallback for any Wan variant + "wanimagetovideo": WanI2V480PConfig, + "wandmdpipeline": FastWan2_1_T2V_480P_Config, + "stepvideo": StepVideoT2VConfig, + # Other fallbacks by architecture +} + + +def get_pipeline_config_cls_from_name( + pipeline_name_or_path: str, +) -> type[PipelineConfig]: + """Get the appropriate configuration class for a given pipeline name or path. + + This function implements a multi-step lookup process to find the most suitable + configuration class for a given pipeline. It follows this order: + 1. Exact match in the PIPE_NAME_TO_CONFIG + 2. Partial match in the PIPE_NAME_TO_CONFIG + 3. Fallback to class name in the model_index.json + 4. else raise an error + + Args: + pipeline_name_or_path (str): The name or path of the pipeline. This can be: + - A registered model ID (e.g., "FastVideo/FastHunyuan-diffusers") + - A local path to a model directory + - A model ID that will be downloaded + + Returns: + Type[PipelineConfig]: The configuration class that best matches the pipeline. + This will be one of: + - A specific weight configuration class if an exact match is found + - A fallback configuration class based on the pipeline architecture + - The base PipelineConfig class if no matches are found + + Note: + - For local paths, the function will verify the model configuration + - For remote models, it will attempt to download the model index + - Warning messages are logged when falling back to less specific configurations + """ + + pipeline_config_cls: type[PipelineConfig] | None = None + + # First try exact match for specific weights + if pipeline_name_or_path in PIPE_NAME_TO_CONFIG: + pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in PIPE_NAME_TO_CONFIG.items(): + if registered_id in pipeline_name_or_path: + pipeline_config_cls = config_class + break + + # If no match, try to use the fallback config + if pipeline_config_cls is None: + if os.path.exists(pipeline_name_or_path): + config = verify_model_config_and_directory(pipeline_name_or_path) + else: + config = maybe_download_model_index(pipeline_name_or_path) + logger.warning( + "Trying to use the config from the model_index.json. FastVideo may not correctly identify the optimal config for this model in this situation." + ) + + pipeline_name = config["_class_name"] + # Try to determine pipeline architecture for fallback + for pipeline_type, detector in PIPELINE_DETECTOR.items(): + if detector(pipeline_name.lower()): + pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get( + pipeline_type + ) + break + + if pipeline_config_cls is not None: + logger.warning( + "No match found for pipeline %s, using fallback config %s.", + pipeline_name_or_path, + pipeline_config_cls, + ) + + if pipeline_config_cls is None: + raise ValueError( + f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path." + ) + + return pipeline_config_cls diff --git a/worldcompass/fastvideo/configs/pipelines/stepvideo.py b/worldcompass/fastvideo/configs/pipelines/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..247c8b9c570909e9229e20242544bad2eb114b34 --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/stepvideo.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig, VAEConfig +from fastvideo.configs.models.dits import StepVideoConfig +from fastvideo.configs.models.vaes import StepVideoVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class StepVideoT2VConfig(PipelineConfig): + """Base configuration for StepVideo pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=StepVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: int = 13 + timesteps_scale: bool = False + pos_magic: str = ( + "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + ) + neg_magic: str = ( + "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" + ) + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "bf16" diff --git a/worldcompass/fastvideo/configs/pipelines/wan.py b/worldcompass/fastvideo/configs/pipelines/wan.py new file mode 100644 index 0000000000000000000000000000000000000000..913ccbc772a88173c3ebc838b34c572226f404ff --- /dev/null +++ b/worldcompass/fastvideo/configs/pipelines/wan.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.configs.models.encoders import ( + BaseEncoderOutput, + CLIPVisionConfig, + T5Config, +) +from fastvideo.configs.models.vaes import WanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor: + mask: torch.tensor = outputs.attention_mask + hidden_state: torch.tensor = outputs.last_hidden_state + seq_lens = mask.gt(0).sum(dim=1).long() + assert torch.isnan(hidden_state).sum() == 0 + prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] + prompt_embeds_tensor: torch.tensor = torch.stack( + [ + torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0, + ) + return prompt_embeds_tensor + + +@dataclass +class WanT2V480PConfig(PipelineConfig): + """Base configuration for Wan T2V 1.3B pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=WanVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: int = 3 + + # Text encoding stage + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (T5Config(),) + ) + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput], torch.tensor], ... + ] = field(default_factory=lambda: (t5_postprocess_text,)) + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp32",) + ) + + # WanConfig-specific added parameters + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class WanT2V720PConfig(WanT2V480PConfig): + """Base configuration for Wan T2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 5 + + +@dataclass +class WanI2V480PConfig(WanT2V480PConfig): + """Base configuration for Wan I2V 14B 480P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Precision for each component + image_encoder_config: EncoderConfig = field( + default_factory=CLIPVisionConfig + ) + image_encoder_precision: str = "fp32" + + def __post_init__(self): + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class WanI2V720PConfig(WanI2V480PConfig): + """Base configuration for Wan I2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 5 + + +@dataclass +class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): + """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: int = 8 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 757, 522] + ) + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): + flow_shift: int = 5 + ti2v_task: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): + flow_shift: int = 5 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 757, 522] + ) + + +@dataclass +class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): + pass + + +@dataclass +class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): + pass diff --git a/worldcompass/fastvideo/configs/sample/__init__.py b/worldcompass/fastvideo/configs/sample/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92765fb634f142577f31cb1443208a3f16103f3d --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/__init__.py @@ -0,0 +1,3 @@ +from fastvideo.configs.sample.base import SamplingParam + +__all__ = ["SamplingParam"] diff --git a/worldcompass/fastvideo/configs/sample/base.py b/worldcompass/fastvideo/configs/sample/base.py new file mode 100644 index 0000000000000000000000000000000000000000..70d48265037e4b12a373e9dac786bcd024e8b869 --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/base.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass +from typing import Any + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class SamplingParam: + """Sampling parameters for video generation.""" + + # All fields below are copied from ForwardBatch + data_type: str = "video" + + # Image inputs + image_path: str | None = None + + # Text inputs + prompt: str | list[str] | None = None + negative_prompt: str = "" + prompt_path: str | None = None + output_path: str = "outputs/" + output_video_name: str | None = None + + # Batch info + num_videos_per_prompt: int = 1 + seed: int = 1024 + + # Original dimensions (before VAE scaling) + num_frames: int = 125 + num_frames_round_down: bool = ( + False # Whether to round down num_frames if it's not divisible by num_gpus + ) + height: int = 720 + width: int = 1280 + fps: int = 24 + + # Denoising parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_rescale: float = 0.0 + + # TeaCache parameters + enable_teacache: bool = False + + # Misc + save_video: bool = True + return_frames: bool = False + + def __post_init__(self) -> None: + self.data_type = "video" if self.num_frames > 1 else "image" + + def check_sampling_param(self): + if self.prompt_path and not self.prompt_path.endswith(".txt"): + raise ValueError("prompt_path must be a txt file") + + def update(self, source_dict: dict[str, Any]) -> None: + for key, value in source_dict.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + logger.exception("%s has no attribute %s", type(self).__name__, key) + + self.__post_init__() + + @classmethod + def from_pretrained(cls, model_path: str) -> "SamplingParam": + from fastvideo.configs.sample.registry import get_sampling_param_cls_for_name + + sampling_cls = get_sampling_param_cls_for_name(model_path) + if sampling_cls is not None: + sampling_param: SamplingParam = sampling_cls() + else: + logger.warning( + "Couldn't find an optimal sampling param for %s. Using the default sampling param.", + model_path, + ) + sampling_param = cls() + + return sampling_param + + @staticmethod + def add_cli_args(parser: Any) -> Any: + """Add CLI arguments for SamplingParam fields.""" + parser.add_argument( + "--prompt", + type=str, + default=SamplingParam.prompt, + help="Text prompt for video generation", + ) + parser.add_argument( + "--negative-prompt", + type=str, + default=SamplingParam.negative_prompt, + help="Negative text prompt for video generation", + ) + parser.add_argument( + "--prompt-path", + type=str, + default=SamplingParam.prompt_path, + help="Path to a text file containing the prompt", + ) + parser.add_argument( + "--output-path", + type=str, + default=SamplingParam.output_path, + help="Path to save the generated video", + ) + parser.add_argument( + "--output-video-name", + type=str, + default=SamplingParam.output_video_name, + help="Name of the output video", + ) + parser.add_argument( + "--num-videos-per-prompt", + type=int, + default=SamplingParam.num_videos_per_prompt, + help="Number of videos to generate per prompt", + ) + parser.add_argument( + "--seed", + type=int, + default=SamplingParam.seed, + help="Random seed for generation", + ) + parser.add_argument( + "--num-frames", + type=int, + default=SamplingParam.num_frames, + help="Number of frames to generate", + ) + parser.add_argument( + "--height", + type=int, + default=SamplingParam.height, + help="Height of generated video", + ) + parser.add_argument( + "--width", + type=int, + default=SamplingParam.width, + help="Width of generated video", + ) + parser.add_argument( + "--fps", + type=int, + default=SamplingParam.fps, + help="Frames per second for saved video", + ) + parser.add_argument( + "--num-inference-steps", + type=int, + default=SamplingParam.num_inference_steps, + help="Number of denoising steps", + ) + parser.add_argument( + "--guidance-scale", + type=float, + default=SamplingParam.guidance_scale, + help="Classifier-free guidance scale", + ) + parser.add_argument( + "--guidance-rescale", + type=float, + default=SamplingParam.guidance_rescale, + help="Guidance rescale factor", + ) + parser.add_argument( + "--save-video", + action="store_true", + default=SamplingParam.save_video, + help="Whether to save the video to disk", + ) + parser.add_argument( + "--no-save-video", + action="store_false", + dest="save_video", + help="Don't save the video to disk", + ) + parser.add_argument( + "--return-frames", + action="store_true", + default=SamplingParam.return_frames, + help="Whether to return the raw frames", + ) + parser.add_argument( + "--image-path", + type=str, + default=SamplingParam.image_path, + help="Path to input image for image-to-video generation", + ) + return parser + + +@dataclass +class CacheParams: + cache_type: str = "none" diff --git a/worldcompass/fastvideo/configs/sample/hunyuan.py b/worldcompass/fastvideo/configs/sample/hunyuan.py new file mode 100644 index 0000000000000000000000000000000000000000..e28e474ad9df65f5143299ab592d70d398ec4a8c --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/hunyuan.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.sample.base import SamplingParam +from fastvideo.configs.sample.teacache import TeaCacheParams + + +@dataclass +class HunyuanSamplingParam(SamplingParam): + num_inference_steps: int = 50 + + num_frames: int = 125 + height: int = 720 + width: int = 1280 + fps: int = 24 + + guidance_scale: float = 1.0 + + teacache_params: TeaCacheParams = field( + default_factory=lambda: TeaCacheParams( + teacache_thresh=0.15, + coefficients=[ + 7.33226126e02, + -4.01131952e02, + 6.75869174e01, + -3.14987800e00, + 9.61237896e-02, + ], + ) + ) + + +@dataclass +class FastHunyuanSamplingParam(HunyuanSamplingParam): + num_inference_steps: int = 6 diff --git a/worldcompass/fastvideo/configs/sample/registry.py b/worldcompass/fastvideo/configs/sample/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..71fa453744ec2d21d4a427865d96f0f7aee73d49 --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/registry.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from collections.abc import Callable +from typing import Any + +from fastvideo.configs.sample.hunyuan import ( + FastHunyuanSamplingParam, + HunyuanSamplingParam, +) +from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam + +# isort: off +from fastvideo.configs.sample.wan import ( + FastWanT2V480PConfig, + Wan2_2_I2V_A14B_SamplingParam, + Wan2_2_T2V_A14B_SamplingParam, + Wan2_2_TI2V_5B_SamplingParam, + WanI2V_14B_480P_SamplingParam, + WanI2V_14B_720P_SamplingParam, + WanT2V_1_3B_SamplingParam, + WanT2V_14B_SamplingParam, +) + +# isort: on +from fastvideo.logger import init_logger +from fastvideo.utils import ( + maybe_download_model_index, + verify_model_config_and_directory, +) + +logger = init_logger(__name__) +# Registry maps specific model weights to their config classes +SAMPLING_PARAM_REGISTRY: dict[str, Any] = { + "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, + "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, + # Add other specific weight variants +} + +# For determining pipeline type from model ID +SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { + "hunyuan": lambda id: "hunyuan" in id.lower(), + "wanpipeline": lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), + "stepvideo": lambda id: "stepvideo" in id.lower(), + # Add other pipeline architecture detectors +} + +# Fallback configs when exact match isn't found but architecture is detected +SAMPLING_FALLBACK_PARAM: dict[str, Any] = { + "hunyuan": HunyuanSamplingParam, # Base Hunyuan config as fallback for any Hunyuan variant + "wanpipeline": WanT2V_1_3B_SamplingParam, # Base Wan config as fallback for any Wan variant + "wanimagetovideo": WanI2V_14B_480P_SamplingParam, + "stepvideo": StepVideoT2VSamplingParam, + # Other fallbacks by architecture +} + + +def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: + """Get the appropriate sampling param for specific pretrained weights.""" + + if os.path.exists(pipeline_name_or_path): + config = verify_model_config_and_directory(pipeline_name_or_path) + logger.warning( + "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed." + ) + else: + config = maybe_download_model_index(pipeline_name_or_path) + + pipeline_name = config["_class_name"] + + # First try exact match for specific weights + if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: + return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): + if registered_id in pipeline_name_or_path: + return config_class + + # If no match, try to use the fallback config + fallback_config = None + # Try to determine pipeline architecture for fallback + for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items(): + if detector(pipeline_name.lower()): + fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type) + break + + logger.warning( + "No match found for pipeline %s, using fallback sampling param %s.", + pipeline_name_or_path, + fallback_config, + ) + return fallback_config diff --git a/worldcompass/fastvideo/configs/sample/stepvideo.py b/worldcompass/fastvideo/configs/sample/stepvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..ef0fee49a152dec86d076905221023cea1182356 --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/stepvideo.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + +from fastvideo.configs.sample.base import SamplingParam + + +@dataclass +class StepVideoT2VSamplingParam(SamplingParam): + # Video parameters + height: int = 720 + width: int = 1280 + num_frames: int = 81 + + # Denoising stage + guidance_scale: float = 9.0 + num_inference_steps: int = 50 + + # neg magic and pos magic + # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" diff --git a/worldcompass/fastvideo/configs/sample/teacache.py b/worldcompass/fastvideo/configs/sample/teacache.py new file mode 100644 index 0000000000000000000000000000000000000000..3bde0648a129c3c45279925e581d6919319459c9 --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/teacache.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.sample.base import CacheParams + + +@dataclass +class TeaCacheParams(CacheParams): + cache_type: str = "teacache" + teacache_thresh: float = 0.0 + coefficients: list[float] = field(default_factory=list) + + +@dataclass +class WanTeaCacheParams(CacheParams): + # Unfortunately, TeaCache is very different for Wan than other models + cache_type: str = "teacache" + teacache_thresh: float = 0.0 + use_ret_steps: bool = True + ret_steps_coeffs: list[float] = field(default_factory=list) + non_ret_steps_coeffs: list[float] = field(default_factory=list) + + @property + def coefficients(self) -> list[float]: + if self.use_ret_steps: + return self.ret_steps_coeffs + else: + return self.non_ret_steps_coeffs + + @property + def ret_steps(self) -> int: + if self.use_ret_steps: + return 5 * 2 + else: + return 1 * 2 + + def get_cutoff_steps(self, num_inference_steps: int) -> int: + if self.use_ret_steps: + return num_inference_steps * 2 + else: + return num_inference_steps * 2 - 2 diff --git a/worldcompass/fastvideo/configs/sample/wan.py b/worldcompass/fastvideo/configs/sample/wan.py new file mode 100644 index 0000000000000000000000000000000000000000..af26e218f72f824f7859909c476ed1970899a1f6 --- /dev/null +++ b/worldcompass/fastvideo/configs/sample/wan.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.sample.base import SamplingParam +from fastvideo.configs.sample.teacache import WanTeaCacheParams + + +@dataclass +class WanT2V_1_3B_SamplingParam(SamplingParam): + # Video parameters + height: int = 480 + width: int = 832 + num_frames: int = 81 + fps: int = 16 + + # Denoising stage + guidance_scale: float = 3.0 + negative_prompt: str = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + ) + num_inference_steps: int = 50 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.08, + ret_steps_coeffs=[ + -5.21862437e04, + 9.23041404e03, + -5.28275948e02, + 1.36987616e01, + -4.99875664e-02, + ], + non_ret_steps_coeffs=[ + 2.39676752e03, + -1.31110545e03, + 2.01331979e02, + -8.29855975e00, + 1.37887774e-01, + ], + ) + ) + + +@dataclass +class WanT2V_14B_SamplingParam(SamplingParam): + # Video parameters + height: int = 720 + width: int = 1280 + num_frames: int = 81 + fps: int = 16 + + # Denoising stage + guidance_scale: float = 5.0 + negative_prompt: str = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + ) + num_inference_steps: int = 50 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.20, + use_ret_steps=False, + ret_steps_coeffs=[ + -3.03318725e05, + 4.90537029e04, + -2.65530556e03, + 5.87365115e01, + -3.15583525e-01, + ], + non_ret_steps_coeffs=[ + -5784.54975374, + 5449.50911966, + -1811.16591783, + 256.27178429, + -13.02252404, + ], + ) + ) + + +@dataclass +class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam): + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 40 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.26, + ret_steps_coeffs=[ + -3.03318725e05, + 4.90537029e04, + -2.65530556e03, + 5.87365115e01, + -3.15583525e-01, + ], + non_ret_steps_coeffs=[ + -5784.54975374, + 5449.50911966, + -1811.16591783, + 256.27178429, + -13.02252404, + ], + ) + ) + + +@dataclass +class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam): + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 40 + + teacache_params: WanTeaCacheParams = field( + default_factory=lambda: WanTeaCacheParams( + teacache_thresh=0.3, + ret_steps_coeffs=[ + -3.03318725e05, + 4.90537029e04, + -2.65530556e03, + 5.87365115e01, + -3.15583525e-01, + ], + non_ret_steps_coeffs=[ + -5784.54975374, + 5449.50911966, + -1811.16591783, + 256.27178429, + -13.02252404, + ], + ) + ) + + +@dataclass +class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam): + # DMD parameters + # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) + num_inference_steps: int = 3 + num_frames: int = 61 + height: int = 448 + width: int = 832 + fps: int = 16 + + +# ============================================= +# ============= Wan2.2 TI2V Models ============= +# ============================================= +@dataclass +class Wan2_2_Base_SamplingParam(SamplingParam): + """Sampling parameters for Wan2.2 TI2V 5B model.""" + + negative_prompt: str | None = ( + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + ) + + +@dataclass +class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam): + """Sampling parameters for Wan2.2 TI2V 5B model.""" + + height: int = 704 + width: int = 1280 + num_frames: int = 121 + fps: int = 24 + guidance_scale: float = 5.0 + num_inference_steps: int = 50 + + +@dataclass +class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): + guidance_scale: float = 4.0 + guidance_scale_2: float = 3.0 + num_inference_steps: int = 40 + fps: int = 16 + + +@dataclass +class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): + guidance_scale: float = 3.5 + guidance_scale_2: float = 3.5 + num_inference_steps: int = 40 + fps: int = 16 diff --git a/worldcompass/fastvideo/configs/utils.py b/worldcompass/fastvideo/configs/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a4baee71d3064180910960eabab6fe96027b093b --- /dev/null +++ b/worldcompass/fastvideo/configs/utils.py @@ -0,0 +1,59 @@ +import argparse +from typing import Any + + +def update_config_from_args( + config: Any, + args_dict: dict[str, Any], + prefix: str = "", + pop_args: bool = False, +) -> bool: + """Update configuration object from arguments dictionary. + + Args: + config: The configuration object to update + args_dict: Dictionary containing arguments + prefix: Prefix for the configuration parameters in the args_dict. + If None, assumes direct attribute mapping without prefix. + """ + # Handle top-level attributes (no prefix) + args_not_to_remove = [ + "model_path", + ] + args_to_remove = [] + if prefix.strip() == "": + for key, value in args_dict.items(): + if hasattr(config, key) and value is not None: + if key == "text_encoder_precisions" and isinstance(value, list): + setattr(config, key, tuple(value)) + else: + setattr(config, key, value) + if pop_args: + args_to_remove.append(key) + else: + # Handle nested attributes with prefix + prefix_with_dot = f"{prefix}." + for key, value in args_dict.items(): + if key.startswith(prefix_with_dot) and value is not None: + attr_name = key[len(prefix_with_dot) :] + if hasattr(config, attr_name): + setattr(config, attr_name, value) + if pop_args: + args_to_remove.append(key) + + if pop_args: + for key in args_to_remove: + if key not in args_not_to_remove: + args_dict.pop(key) + + return len(args_to_remove) > 0 + + +def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]: + """Clean the arguments by removing the ones that not explicitly provided by the user.""" + provided_args = {} + for k, v in vars(args).items(): + if v is not None and hasattr(args, "_provided") and k in args._provided: + provided_args[k] = v + + return provided_args diff --git a/worldcompass/fastvideo/configs/wan_1.3B_t2v_pipeline.json b/worldcompass/fastvideo/configs/wan_1.3B_t2v_pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..75648de613ed7747b73905971626c90bc9914608 --- /dev/null +++ b/worldcompass/fastvideo/configs/wan_1.3B_t2v_pipeline.json @@ -0,0 +1,39 @@ +{ + "embedded_cfg_scale": 6.0, + "flow_shift": 3, + "dit_cpu_offload": true, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": false, + "vae_sp": false, + "vae_config": { + "load_encoder": false, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 8, + "use_tiling": false, + "use_temporal_tiling": false, + "use_parallel_tiling": false, + "use_feature_cache": true + }, + "dit_config": { + "prefix": "Wan", + "quant_config": null + }, + "text_encoder_precisions": ["fp32"], + "text_encoder_configs": [ + { + "prefix": "t5", + "quant_config": null, + "lora_config": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false +} diff --git a/worldcompass/fastvideo/configs/wan_14B_i2v_480p_pipeline.json b/worldcompass/fastvideo/configs/wan_14B_i2v_480p_pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..3a41fb24cdea88993ad7f5c015c47943cfddd2d7 --- /dev/null +++ b/worldcompass/fastvideo/configs/wan_14B_i2v_480p_pipeline.json @@ -0,0 +1,47 @@ +{ + "embedded_cfg_scale": 6.0, + "flow_shift": 3, + "dit_cpu_offload": true, + "disable_autocast": false, + "precision": "bf16", + "vae_precision": "fp32", + "vae_tiling": false, + "vae_sp": false, + "vae_config": { + "load_encoder": true, + "load_decoder": true, + "tile_sample_min_height": 256, + "tile_sample_min_width": 256, + "tile_sample_min_num_frames": 16, + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 192, + "tile_sample_stride_num_frames": 12, + "blend_num_frames": 8, + "use_tiling": false, + "use_temporal_tiling": false, + "use_parallel_tiling": false, + "use_feature_cache": true + }, + "dit_config": { + "prefix": "Wan", + "quant_config": null + }, + "text_encoder_precisions": ["fp32"], + "text_encoder_configs": [ + { + "prefix": "t5", + "quant_config": null, + "lora_config": null + } + ], + "mask_strategy_file_path": null, + "enable_torch_compile": false, + "image_encoder_config": { + "prefix": "clip", + "quant_config": null, + "lora_config": null, + "num_hidden_layers_override": null, + "require_post_norm": null + }, + "image_encoder_precision": "fp32" +} diff --git a/worldcompass/fastvideo/dataset/__init__.py b/worldcompass/fastvideo/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..91f152c4c0aaa1376eec6188dba4e0f62cff95b1 --- /dev/null +++ b/worldcompass/fastvideo/dataset/__init__.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +from torchvision import transforms +from torchvision.transforms import Lambda + +from fastvideo.dataset.camera_dataset import build_hy_camera_dataloader +from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset +from fastvideo.dataset.transform import ( + CenterCropResizeVideo, + Normalize255, + TemporalRandomCrop, +) +from fastvideo.dataset.validation_dataset import ValidationDataset + + +def getdataset(args) -> VideoCaptionMergedDataset: + if args.do_temporal_sample: + temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x + else: + temporal_sample = None + norm_fun = Lambda(lambda x: 2.0 * x - 1.0) + resize_topcrop = [ + CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), + ] + resize = [ + CenterCropResizeVideo((args.max_height, args.max_width)), + ] + transform = transforms.Compose( + [ + # Normalize255(), + *resize, + ] + ) + transform_topcrop = transforms.Compose( + [ + Normalize255(), + *resize_topcrop, + norm_fun, + ] + ) + return VideoCaptionMergedDataset( + data_merge_path=args.data_merge_path, + args=args, + transform=transform, + temporal_sample=temporal_sample, + transform_topcrop=transform_topcrop, + seed=args.seed, + ) + + +__all__ = [ + "ValidationDataset", + "build_camera_dataloader", + "build_ar_dataloader", + "build_hy_camera_dataloader", + "VideoCaptionMergedDataset", +] diff --git a/worldcompass/fastvideo/dataset/camera_dataset.py b/worldcompass/fastvideo/dataset/camera_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b5f69fbfaa24bf0fd83b1517129d870ac07ce4de --- /dev/null +++ b/worldcompass/fastvideo/dataset/camera_dataset.py @@ -0,0 +1,548 @@ +import json +import os +import sys + +sys.path.append(os.path.abspath(".")) +import torch +import pandas as pd +import numpy as np +import random +from pathlib import Path + +from torch.utils.data import Dataset, Sampler +from torchdata.stateful_dataloader import StatefulDataLoader +from scipy.spatial.transform import Rotation as R + +from fastvideo.distributed import ( + get_sp_world_size, + get_world_group, + get_world_rank, + get_world_size, + get_gpu_world_size, + get_gpu_group, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def camera_center_normalization(w2c, nframe): + c2w_view0 = w2c[::nframe].inverse() # [B,4,4] + c2w_view0 = c2w_view0.repeat_interleave(nframe, dim=0) # [BF,4,4] + w2c = c2w_view0 @ w2c + return w2c + + +# action keys +ACTION_KEYS = [ + "D", + "DL", + "DR", +] + + +class DP_SP_BatchSampler(Sampler[list[int]]): + """A simple sequential batch sampler that yields batches of indices.""" + + def __init__( + self, + batch_size: int, + dataset_size: int, + num_sp_groups: int, + sp_world_size: int, + global_rank: int, + drop_last: bool = True, + drop_first_row: bool = False, + seed: int = 0, + ): + self.batch_size = batch_size + self.dataset_size = dataset_size + self.drop_last = drop_last + self.seed = seed + self.num_sp_groups = num_sp_groups + self.global_rank = global_rank + self.sp_world_size = sp_world_size + + # ── epoch-level RNG ──────────────────────────────────────────────── + rng = torch.Generator().manual_seed(self.seed) + # Create a random permutation of all indices + global_indices = torch.randperm(self.dataset_size, generator=rng) + + if drop_first_row: + # drop 0 in global_indices + global_indices = global_indices[global_indices != 0] + self.dataset_size = self.dataset_size - 1 + + if self.drop_last: + # For drop_last=True, we: + # 1. Ensure total samples is divisible by (batch_size * num_sp_groups) + # 2. This guarantees each SP group gets same number of complete batches + # 3. Prevents uneven batch sizes across SP groups at end of epoch + num_batches = self.dataset_size // self.batch_size + num_global_batches = num_batches // self.num_sp_groups + global_indices = global_indices[ + : num_global_batches * self.num_sp_groups * self.batch_size + ] + else: + if self.dataset_size % (self.num_sp_groups * self.batch_size) != 0: + # add more indices to make it divisible by (batch_size * num_sp_groups) + padding_size = self.num_sp_groups * self.batch_size - ( + self.dataset_size % (self.num_sp_groups * self.batch_size) + ) + logger.info( + "Padding the dataset from %d to %d", + self.dataset_size, + self.dataset_size + padding_size, + ) + global_indices = torch.cat( + [global_indices, global_indices[:padding_size]] + ) + + # shard the indices to each sp group + ith_sp_group = self.global_rank // self.sp_world_size + sp_group_local_indices = global_indices[ + ith_sp_group :: self.num_sp_groups + ] + self.sp_group_local_indices = sp_group_local_indices + logger.info( + "Dataset size for each sp group: %d", len(sp_group_local_indices) + ) + + def __iter__(self): + indices = self.sp_group_local_indices + for i in range(0, len(indices), self.batch_size): + batch_indices = indices[i : i + self.batch_size] + yield batch_indices.tolist() + + def __len__(self): + return len(self.sp_group_local_indices) // self.batch_size + + +def angle_bucket(theta): + """把角度映射到 [-1, 0, +1] 三类.""" + out = torch.zeros_like(theta, dtype=torch.int) + out[(theta >= 0) & (theta < 30)] = 1 # 正 + out[(theta >= 150) & (theta <= 180)] = -1 # 负 + # 30–60 或 120–150 → 斜，标记为 2 或 -2 + out[(theta >= 30) & (theta < 60)] = 2 # 正斜 + out[(theta >= 120) & (theta < 150)] = -2 # 负斜 + return out + + +def angle_to_onehot(x): + """ + x: (B,2)，分别是 (θx, θy)，范围 [0,180] + return: (B,4) one-hot 向量 [右, 左, 后, 前] + """ + B = x.size(0) + one_hot = torch.zeros(B, 4, dtype=torch.int) + + bucket_x = angle_bucket(x[:, 0]) + bucket_z = angle_bucket(x[:, 1]) + + # θx 控制前/后 + one_hot[bucket_x == 1, 3] = 1 # 前 + one_hot[bucket_x == -1, 2] = 1 # 后 + one_hot[bucket_x == 2, 3] = 1 # 前斜 → 前 + one_hot[bucket_x == -2, 2] = 1 # 后斜 → 后 + + # θy 控制左/右 + one_hot[bucket_z == 1, 0] = 1 # 右 + one_hot[bucket_z == -1, 1] = 1 # 左 + one_hot[bucket_z == 2, 0] = 1 # 右斜 → 右 + one_hot[bucket_z == -2, 1] = 1 # 左斜 → 左 + + return one_hot + + +class HunyuanImageJsonDataset(Dataset): + def __init__( + self, + json_path, + causal, + window_frames, + batch_size, + cfg_rate, + i2v_rate, + drop_last, + drop_first_row, + seed, + random_pose_path="", + neg_prompt_path="", + neg_byt5_prompt_path="", + ): + self.json_data = json.load(open(json_path, "r")) + self.all_length = len(self.json_data) + self.causal = causal + self.window_frames = window_frames + self.cfg_rate = cfg_rate + self.rng = random.Random(seed) + self.i2v_rate = i2v_rate + + # Load random pose from config path + if not random_pose_path: + raise ValueError("random_pose_path must be provided") + self.random_pose = json.load(open(random_pose_path, "r")) + + self.sampler = DP_SP_BatchSampler( + batch_size=batch_size, + dataset_size=self.all_length, + num_sp_groups=get_world_size() // get_gpu_world_size(), + sp_world_size=get_gpu_world_size(), + global_rank=get_world_rank(), + drop_last=drop_last, + drop_first_row=drop_first_row, + seed=seed, + ) + self.mapping = { + (0, 0, 0, 0): 0, + (1, 0, 0, 0): 1, + (0, 1, 0, 0): 2, + (0, 0, 1, 0): 3, + (0, 0, 0, 1): 4, + (1, 0, 1, 0): 5, + (1, 0, 0, 1): 6, + (0, 1, 1, 0): 7, + (0, 1, 0, 1): 8, + } + + # Load negative prompt from config path + # if not neg_prompt_path: + # raise ValueError("neg_prompt_path must be provided") + # self.neg_prompt_pt = torch.load( + # neg_prompt_path, + # map_location="cpu", + # weights_only=True, + # ) + + # # Load negative byt5 prompt from config path + # if not neg_byt5_prompt_path: + # raise ValueError("neg_byt5_prompt_path must be provided") + # self.neg_byt5_pt = torch.load( + # neg_byt5_prompt_path, + # map_location="cpu", + # weights_only=True, + # ) + + def __len__(self): + return self.all_length + + def camera_center_normalization(self, w2c): + c2w = np.linalg.inv(w2c) + C0_inv = np.linalg.inv(c2w[0]) + c2w_aligned = np.array([C0_inv @ C for C in c2w]) + return np.linalg.inv(c2w_aligned) + + def one_hot_to_one_dimension(self, one_hot): + y = torch.tensor([self.mapping[tuple(row.tolist())] for row in one_hot]) + return y + + def __getitem__(self, idx): + while True: + try: + json_data = self.json_data[idx] + + latent_pt_path = json_data["latent_path"] + + latent_pt = torch.load( + os.path.join(latent_pt_path), + map_location="cpu", + weights_only=True, + ) + + # Extract latent from the new structure + latent = latent_pt["latent"] + if latent.ndim == 5: # [1, C, T, H, W] + latent = latent[0] # Remove batch dimension -> [C, T, H, W] + latent_length = latent.shape[1] + + # 将latent的第三维zero padding到self.window_frames + if latent.shape[1] < self.window_frames: + pad_size = self.window_frames - latent.shape[1] + pad_shape = list(latent.shape) + pad_shape[1] = pad_size + latent = torch.cat( + [latent, torch.zeros(pad_shape, dtype=latent.dtype)], + dim=1, + ) + elif latent.shape[1] > self.window_frames: + latent = latent[:, : self.window_frames, ...] + + # Get caption from json_data directly + prompt = json_data.get("caption", "") + if not prompt: + raise ValueError("caption is None or empty") + + # Extract text embeddings from the new structure + prompt_embed = latent_pt["prompt_embeds"] + if prompt_embed.ndim == 3: # [1, L, D] + prompt_embed = prompt_embed[ + 0 + ] # Remove batch dimension -> [L, D] + + prompt_mask = latent_pt["prompt_mask"] + if prompt_mask.ndim == 2: # [1, L] + prompt_mask = prompt_mask[ + 0 + ] # Remove batch dimension -> [L] + + # Extract image condition from the new structure + image_cond = latent_pt["image_cond"] + if image_cond.ndim == 5: # [1, C, 1, H, W] + image_cond = image_cond[ + 0 + ] # Remove batch dimension -> [C, 1, H, W] + + # Extract vision states + vision_states = latent_pt["vision_states"] + if vision_states.ndim == 3: # [1, N, D] + vision_states = vision_states[ + 0 + ] # Remove batch dimension -> [N, D] + + # Extract byT5 embeddings + byt5_text_states = latent_pt["byt5_text_states"] + if byt5_text_states.ndim == 3: # [1, 256, 1472] + byt5_text_states = byt5_text_states[ + 0 + ] # Remove batch dimension -> [256, 1472] + + byt5_text_mask = latent_pt["byt5_text_mask"] + if byt5_text_mask.ndim == 2: # [1, 256] + byt5_text_mask = byt5_text_mask[ + 0 + ] # Remove batch dimension -> [256] + + pose_json = self.random_pose[idx % len(self.random_pose)] + pose_keys = list(pose_json.keys()) + + intrinsic_list = [] + w2c_list = [] + for i in range(self.window_frames): + t_key = pose_keys[i] + + c2w = np.array(pose_json[t_key]["extrinsic"]) + w2c = np.linalg.inv(c2w) + + w2c_list.append(w2c) + intrinsic = np.array(pose_json[t_key]["K"]) + if (intrinsic > 2000.0).any(): + raise ValueError("intrinsic > 2000") + intrinsic[0, 0] /= intrinsic[0, 2] * 2 + intrinsic[1, 1] /= intrinsic[1, 2] * 2 + intrinsic[0, 2] = 0.5 + intrinsic[1, 2] = 0.5 + intrinsic_list.append(intrinsic) + + w2c_list = np.array(w2c_list) + # w2c_list = self.camera_center_normalization(w2c_list) + intrinsic_list = torch.tensor(np.array(intrinsic_list)) + + # 从这里开始是计算离散action + c2ws = np.linalg.inv(w2c_list) + C_inv = np.linalg.inv(c2ws[:-1]) + relative_c2w = np.zeros_like(c2ws) + relative_c2w[0, ...] = c2ws[0, ...] + relative_c2w[1:, ...] = C_inv @ c2ws[1:, ...] + trans_one_hot = np.zeros( + (relative_c2w.shape[0], 4), dtype=np.int32 + ) + rotate_one_hot = np.zeros( + (relative_c2w.shape[0], 4), dtype=np.int32 + ) + + move_norm_valid = 0.0001 + for i in range(1, relative_c2w.shape[0]): + move_dirs = relative_c2w[i, :3, 3] # 方向向量 + move_norms = np.linalg.norm(move_dirs) + # print(move_norms) + if move_norms > move_norm_valid: # 认为有移动 + move_norm_dirs = move_dirs / move_norms + # print(np.linalg.norm(move_norm_dirs)) + angles_rad = np.arccos( + move_norm_dirs.clip(-1.0, 1.0) + ) # 防止数值误差超出[-1,1] + trans_angles_deg = angles_rad * ( + 180.0 / torch.pi + ) # 转成角度 0-180，这里是translation的角度 + else: + trans_angles_deg = torch.zeros(3) + + R_rel = relative_c2w[i, :3, :3] + r = R.from_matrix(R_rel) + rot_angles_deg = r.as_euler( + "xyz", degrees=True + ) # 输出欧拉角，单位度 + + # 现在需要解决的就是scale的问题，norms以及多少度算是转动 + if move_norms > move_norm_valid: # 认为有移动 + # 先判断前进后退，这个就根据z轴的来就行 + if trans_angles_deg[2] < 60: + trans_one_hot[i, 0] = 1 # 前 + elif trans_angles_deg[2] > 120: + trans_one_hot[i, 1] = 1 # 后 + + if trans_angles_deg[0] < 60: + trans_one_hot[i, 2] = 1 # 右 + elif trans_angles_deg[0] > 120: + trans_one_hot[i, 3] = 1 # 左 + + # 接下来是rotate, 对于matrix来说角度为9比较合适，对于其他数据来说需要找一个合适的角度 + if rot_angles_deg[1] > 5e-2: + rotate_one_hot[i, 0] = 1 # 右 + elif rot_angles_deg[1] < -5e-2: + rotate_one_hot[i, 1] = 1 # 左 + + if rot_angles_deg[0] > 5e-2: + rotate_one_hot[i, 2] = 1 # 上 + elif rot_angles_deg[0] < -5e-2: + rotate_one_hot[i, 3] = 1 # 下 + + trans_one_hot = torch.tensor(trans_one_hot) + rotate_one_hot = torch.tensor(rotate_one_hot) + action = torch.cat([trans_one_hot, rotate_one_hot], dim=1) + + trans_one_label = self.one_hot_to_one_dimension(trans_one_hot) + rotate_one_label = self.one_hot_to_one_dimension(rotate_one_hot) + action_for_pe = trans_one_label * 9 + rotate_one_label + + i2v_mask = torch.ones_like(latent) + + batch = { + "i2v_mask": i2v_mask, + "latent": latent, + "prompt_embed": prompt_embed, + "w2c": torch.tensor(w2c_list), + "intrinsic": intrinsic_list, + "action": action_for_pe, + "action_for_pe": action_for_pe, + "context_frames_list": None, # selected context frames for each chunk + "prompt": prompt, + "image_cond": image_cond, + "vision_states": vision_states, + "prompt_mask": prompt_mask, + "byt5_text_states": byt5_text_states, + "byt5_text_mask": byt5_text_mask, + # "all_channel_latents": all_channel_latents, + } + break + except Exception as e: + print("error:", e, latent_pt_path) + idx = self.rng.randint(0, self.all_length - 1) + + return batch + + +def hunyuan_latent_collate_function(batch): + latent = torch.stack([b["latent"] for b in batch], dim=0) + prompt_embed = torch.stack([b["prompt_embed"] for b in batch], dim=0) + w2c = torch.stack([b["w2c"] for b in batch], dim=0) + intrinsic = torch.stack([b["intrinsic"] for b in batch], dim=0) + action = torch.stack([b["action"] for b in batch], dim=0) + action_for_pe = torch.stack([b["action_for_pe"] for b in batch], dim=0) + i2v_mask = torch.stack([b["i2v_mask"] for b in batch], dim=0) + + image_cond = torch.stack([b["image_cond"] for b in batch], dim=0) + vision_states = torch.stack([b["vision_states"] for b in batch], dim=0) + prompt_mask = torch.stack([b["prompt_mask"] for b in batch], dim=0) + byt5_text_states = torch.stack( + [b["byt5_text_states"] for b in batch], dim=0 + ) + byt5_text_mask = torch.stack([b["byt5_text_mask"] for b in batch], dim=0) + + context_frames_list = [b["context_frames_list"] for b in batch] + prompt = [b["prompt"] for b in batch] + + return { + "i2v_mask": i2v_mask, + "latent": latent, + "prompt_embed": prompt_embed, + "w2c": w2c, + "intrinsic": intrinsic, + "action": action, + "context_frames_list": context_frames_list, + "action_for_pe": action_for_pe, + "image_cond": image_cond, + "vision_states": vision_states, + "prompt_mask": prompt_mask, + "byt5_text_states": byt5_text_states, + "byt5_text_mask": byt5_text_mask, + "prompt": prompt, + } + + +def cycle(dl): + while True: + for data in dl: + yield data + + +def latent_collate_function(batch): + # dir_norm = torch.stack([b["dir_norm"] for b in batch], dim=0) + # trans_discrete_labels = torch.stack([b["trans_discrete_labels"] for b in batch], dim=0) + # norm_dirs = torch.stack([b["norm_dirs"] for b in batch], dim=0) + latent = torch.stack([b["latent"] for b in batch], dim=0) + prompt_embed = torch.stack([b["prompt_embed"] for b in batch], dim=0) + w2c = torch.stack([b["w2c"] for b in batch], dim=0) + intrinsic = torch.stack([b["intrinsic"] for b in batch], dim=0) + action = torch.stack([b["action"] for b in batch], dim=0) + action_for_pe = torch.stack([b["action_for_pe"] for b in batch], dim=0) + # rotate_degree = torch.stack([b["rotate_degree"] for b in batch], dim=0) + i2v_mask = torch.stack([b["i2v_mask"] for b in batch], dim=0) + # video_path = [b["video_path"] for b in batch] + prompt = [b["prompt"] for b in batch] + return { + "i2v_mask": i2v_mask, + # "dir_norm": dir_norm, + # "rotate_degree": rotate_degree, + # "trans_discrete_labels": trans_discrete_labels, + # "norm_dirs": norm_dirs, + "latent": latent, + "prompt_embed": prompt_embed, + "prompt": prompt, + "w2c": w2c, + "intrinsic": intrinsic, + "action": action, + "action_for_pe": action_for_pe, + } + + +def build_hy_camera_dataloader( + json_path, + causal, + window_frames, + batch_size, + num_data_workers, + drop_last, + drop_first_row, + seed, + cfg_rate, + i2v_rate, + random_pose_path="", + neg_prompt_path="", + neg_byt5_prompt_path="", +) -> tuple[HunyuanImageJsonDataset, StatefulDataLoader]: + dataset = HunyuanImageJsonDataset( + json_path, + causal, + window_frames, + batch_size, + cfg_rate, + i2v_rate, + drop_last=drop_last, + drop_first_row=drop_first_row, + seed=seed, + random_pose_path=random_pose_path, + neg_prompt_path=neg_prompt_path, + neg_byt5_prompt_path=neg_byt5_prompt_path, + ) + + loader = StatefulDataLoader( + dataset, + batch_sampler=dataset.sampler, + collate_fn=hunyuan_latent_collate_function, + num_workers=num_data_workers, + pin_memory=True, + persistent_workers=num_data_workers > 0, + ) + return dataset, loader diff --git a/worldcompass/fastvideo/dataset/dataloader/schema.py b/worldcompass/fastvideo/dataset/dataloader/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..455c43ee16e1285d885eec250c5ab94bdf24c222 --- /dev/null +++ b/worldcompass/fastvideo/dataset/dataloader/schema.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# schema.py +"""Unified data schema and format for saving and loading image/video data after preprocessing. + +It uses apache arrow in-memory format that can be consumed by modern data frameworks that can handle +parquet or lance file. +""" + +import pyarrow as pa + +pyarrow_schema_t2v = pa.schema( + [ + pa.field("id", pa.string()), + # --- Image/Video VAE latents --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("vae_latent_bytes", pa.binary()), + # e.g., [C, T, H, W] or [C, H, W] + pa.field("vae_latent_shape", pa.list_(pa.int64())), + # e.g., 'float32' + pa.field("vae_latent_dtype", pa.string()), + # --- Text encoder output tensor --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("text_embedding_bytes", pa.binary()), + # e.g., [SeqLen, Dim] + pa.field("text_embedding_shape", pa.list_(pa.int64())), + # e.g., 'bfloat16' or 'float32' + pa.field("text_embedding_dtype", pa.string()), + # --- Metadata --- + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), # 'image' or 'video' + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + # -- Video-specific (can be null/default for images) --- + # Number of frames processed (e.g., 1 for image, N for video) + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), + ] +) diff --git a/worldcompass/fastvideo/dataset/latent_datasets.py b/worldcompass/fastvideo/dataset/latent_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..0ec0d0ef0141d706e821663f889bc88921bc14da --- /dev/null +++ b/worldcompass/fastvideo/dataset/latent_datasets.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +import random + +import torch +from torch.utils.data import Dataset + + +class LatentDataset(Dataset): + def __init__( + self, + json_path, + num_latent_t, + cfg_rate, + seed: int = 42, + ) -> None: + # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path + self.json_path = json_path + self.cfg_rate = cfg_rate + self.seed = seed + # Create a seeded random generator for deterministic CFG + self.rng = random.Random(seed) + self.datase_dir_path = os.path.dirname(json_path) + self.video_dir = os.path.join(self.datase_dir_path, "video") + self.latent_dir = os.path.join(self.datase_dir_path, "latent") + self.prompt_embed_dir = os.path.join( + self.datase_dir_path, "prompt_embed" + ) + self.prompt_attention_mask_dir = os.path.join( + self.datase_dir_path, "prompt_attention_mask" + ) + with open(self.json_path) as f: + self.data_anno = json.load(f) + # json.load(f) already keeps the order + # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path']) + self.num_latent_t = num_latent_t + + self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32) + + self.uncond_prompt_mask = torch.zeros(256).bool() + self.lengths = [ + data_item.get("length", 1) for data_item in self.data_anno + ] + + def __getitem__(self, idx): + latent_file = self.data_anno[idx]["latent_path"] + prompt_embed_file = self.data_anno[idx]["prompt_embed_path"] + prompt_attention_mask_file = self.data_anno[idx][ + "prompt_attention_mask" + ] + # load + latent = torch.load( + os.path.join(self.latent_dir, latent_file), + map_location="cpu", + weights_only=True, + ) + latent = latent.squeeze(0)[:, -self.num_latent_t :] + if self.rng.random() < self.cfg_rate: + prompt_embed = self.uncond_prompt_embed + prompt_attention_mask = self.uncond_prompt_mask + else: + prompt_embed = torch.load( + os.path.join(self.prompt_embed_dir, prompt_embed_file), + map_location="cpu", + weights_only=True, + ) + prompt_attention_mask = torch.load( + os.path.join( + self.prompt_attention_mask_dir, prompt_attention_mask_file + ), + map_location="cpu", + weights_only=True, + ) + return latent, prompt_embed, prompt_attention_mask + + def __len__(self): + return len(self.data_anno) + + +def latent_collate_function(batch): + # return latent, prompt, latent_attn_mask, text_attn_mask + # latent_attn_mask: # b t h w + # text_attn_mask: b 1 l + # needs to check if the latent/prompt' size and apply padding & attn mask + latents, prompt_embeds, prompt_attention_masks = zip(*batch, strict=True) + # calculate max shape + max_t = max([latent.shape[1] for latent in latents]) + max_h = max([latent.shape[2] for latent in latents]) + max_w = max([latent.shape[3] for latent in latents]) + + # padding + latent_list: list[torch.Tensor] = [ + torch.nn.functional.pad( + latent, + ( + 0, + max_t - latent.shape[1], + 0, + max_h - latent.shape[2], + 0, + max_w - latent.shape[3], + ), + ) + for latent in latents + ] + # attn mask + latent_attn_mask = torch.ones(len(latent_list), max_t, max_h, max_w) + # set to 0 if padding + for i, latent in enumerate(latent_list): + latent_attn_mask[i, latent.shape[1] :, :, :] = 0 + latent_attn_mask[i, :, latent.shape[2] :, :] = 0 + latent_attn_mask[i, :, :, latent.shape[3] :] = 0 + + prompt_embeds = torch.stack(prompt_embeds, dim=0) + prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0) + latents = torch.stack(latent_list, dim=0) + return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks diff --git a/worldcompass/fastvideo/dataset/parquet_dataset_iterable_style.py b/worldcompass/fastvideo/dataset/parquet_dataset_iterable_style.py new file mode 100644 index 0000000000000000000000000000000000000000..35bb4887d14c6b1bf173e77c8453d0cd6b4684af --- /dev/null +++ b/worldcompass/fastvideo/dataset/parquet_dataset_iterable_style.py @@ -0,0 +1,350 @@ +import os +import pickle +import random + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import tqdm +from torch.utils.data import IterableDataset, get_worker_info +from torchdata.stateful_dataloader import StatefulDataLoader + +from fastvideo.dataset.utils import collate_latents_embs_masks +from fastvideo.distributed import ( + get_sp_world_size, + get_world_group, + get_world_rank, + get_world_size, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class BatchIterator: + # TODO: Implement state_dict and load_state_dict to support resume. + def __init__( + self, + files, + batch_size, + text_padding_length, + keys, + worker_num_samples, + read_batch_size, + ): + self.files = files + self.batch_size = batch_size + self.text_padding_length = text_padding_length + self.keys = keys + self.worker_num_samples = worker_num_samples + self.processed_samples = 0 + self.buffer = [] + self.read_batch_size = read_batch_size + + def __iter__(self): + for file in self.files: + if self.processed_samples >= self.worker_num_samples: + return + + reader = pq.ParquetFile(file) + for batch in reader.iter_batches(batch_size=self.read_batch_size): + if self.processed_samples >= self.worker_num_samples: + return + + self.buffer.extend(batch.to_pylist()) + + while len(self.buffer) >= self.batch_size: + if self.processed_samples >= self.worker_num_samples: + return + + batch_to_process = self.buffer[: self.batch_size] + self.buffer = self.buffer[self.batch_size :] + + all_latents, all_embs, all_masks, caption_text = ( + collate_latents_embs_masks( + batch_to_process, + self.text_padding_length, + self.keys, + ) + ) + self.processed_samples += self.batch_size + yield all_latents, all_embs, all_masks, caption_text + + +class LatentsParquetIterStyleDataset(IterableDataset): + """Efficient loader for video-text data from a directory of Parquet files.""" + + # Modify this in the future if we want to add more keys, for example, in image to video. + keys = [("vae_latent", "latent"), ("text_embedding")] + + def __init__( + self, + path: str, + batch_size: int = 1024, + cfg_rate: float = 0.1, + num_workers: int = 1, + drop_last: bool = True, + text_padding_length: int = 512, + seed: int = 42, + read_batch_size: int = 32, + parquet_schema: pa.Schema = None, + ): + super().__init__() + self.path = str(path) + self.batch_size = batch_size + self.parquet_schema = parquet_schema + self.cfg_rate = cfg_rate + self.text_padding_length = text_padding_length + self.seed = seed + self.read_batch_size = read_batch_size + # Get distributed training info + self.global_rank = get_world_rank() + self.world_size = get_world_size() + self.sp_world_size = get_sp_world_size() + self.num_sp_groups = self.world_size // self.sp_world_size + num_workers = 1 if num_workers == 0 else num_workers + # Get sharding info + shard_parquet_files, shard_total_samples, shard_parquet_lengths = ( + shard_parquet_files_across_sp_groups_and_workers( + self.path, self.num_sp_groups, num_workers, seed + ) + ) + + if drop_last: + self.worker_num_samples = ( + min(shard_total_samples) // batch_size * batch_size + ) + # Assign files to current rank's SP group + ith_sp_group = self.global_rank // self.sp_world_size + self.sp_group_parquet_files = shard_parquet_files[ + ith_sp_group :: self.num_sp_groups + ] + self.sp_group_parquet_lengths = shard_parquet_lengths[ + ith_sp_group :: self.num_sp_groups + ] + self.sp_group_num_samples = shard_total_samples[ + ith_sp_group :: self.num_sp_groups + ] + logger.info( + "In total %d parquet files, %d samples, after sharding we retain %d samples due to drop_last", + sum([len(shard) for shard in shard_parquet_files]), + sum(shard_total_samples), + self.worker_num_samples * self.num_sp_groups * num_workers, + ) + else: + raise ValueError("drop_last must be True") + logger.info( + "Each dataloader worker will load %d samples", + self.worker_num_samples, + ) + + def __iter__(self): + worker_info = get_worker_info() + worker_id = worker_info.id if worker_info is not None else 1 + + worker_files = self.sp_group_parquet_files[worker_id] + + batch_iterator = BatchIterator( + files=worker_files, + batch_size=self.batch_size, + text_padding_length=self.text_padding_length, + keys=self.keys, + worker_num_samples=self.worker_num_samples, + read_batch_size=self.read_batch_size, + ) # type: ignore + + yield from batch_iterator + + if batch_iterator.processed_samples != self.worker_num_samples: + raise ValueError( + "Rank %d, Worker %d: Not enough samples to process, this should not happen", + self.global_rank, + worker_id, + ) + + +def shard_parquet_files_across_sp_groups_and_workers( + path: str, + num_sp_groups: int, + num_workers: int, + seed: int = 42, +) -> tuple[list[list[str]], list[int], list[dict[str, int]]]: + """Shard parquet files across SP groups and workers in a balanced way. + + Args: + path: Directory containing parquet files + num_sp_groups: Number of SP groups to shard across + num_workers: Number of workers per SP group + seed: Random seed for shuffling + + Returns: + Tuple containing: + - List of lists of parquet files for each shard + - List of total samples per shard + - List of dictionaries mapping file paths to their lengths + """ + # Check if sharding plan already exists + sharding_info_dir = os.path.join( + path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers" + ) + + # Only rank 0 handles cache checking and file scanning + if get_world_rank() == 0: + cache_loaded = False + shard_parquet_files = None + shard_total_samples = None + shard_parquet_lengths = None + + # First try to load existing sharding plan + if os.path.exists(sharding_info_dir): + logger.info("Loading sharding plan from %s", sharding_info_dir) + try: + with open( + os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), + "rb", + ) as f: + shard_parquet_files = pickle.load(f) + with open( + os.path.join(sharding_info_dir, "shard_total_samples.pkl"), + "rb", + ) as f: + shard_total_samples = pickle.load(f) + with open( + os.path.join( + sharding_info_dir, "shard_parquet_lengths.pkl" + ), + "rb", + ) as f: + shard_parquet_lengths = pickle.load(f) + cache_loaded = True + logger.info("Successfully loaded sharding plan") + except Exception as e: + logger.error("Error loading sharding plan: %s", str(e)) + logger.info("Falling back to creating new sharding plan") + cache_loaded = False + + # If cache not loaded (either doesn't exist or failed to load), create sharding plan + if not cache_loaded: + logger.info("Creating new sharding plan") + logger.info("Scanning for parquet files in %s", path) + + # Find all parquet files + parquet_files = [] + + for root, _, files in os.walk(path): + for file in files: + if file.endswith(".parquet"): + parquet_files.append(os.path.join(root, file)) + + if not parquet_files: + raise ValueError("No parquet files found in %s", path) + + # Calculate file lengths efficiently using a single pass + logger.info("Calculating file lengths...") + lengths = [] + for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"): + lengths.append(pq.ParquetFile(file).metadata.num_rows) + + total_samples = sum(lengths) + logger.info( + "Found %d files with %d total samples", + len(parquet_files), + total_samples, + ) + + # Sort files by length for better balancing + sorted_indices = np.argsort(lengths) + sorted_files = [parquet_files[i] for i in sorted_indices] + sorted_lengths = [lengths[i] for i in sorted_indices] + + # Create shards + num_shards = num_sp_groups * num_workers + shard_parquet_files = [[] for _ in range(num_shards)] + shard_total_samples = [0] * num_shards + shard_parquet_lengths = [{} for _ in range(num_shards)] + + # Distribute files to shards using a greedy approach + logger.info("Distributing files to shards...") + for file, length in zip( + reversed(sorted_files), reversed(sorted_lengths), strict=True + ): + # Find shard with minimum current length + target_shard = np.argmin(shard_total_samples) + shard_parquet_files[target_shard].append(file) + shard_total_samples[target_shard] += length + shard_parquet_lengths[target_shard][file] = length + # randomize each shard + for shard in shard_parquet_files: + rng = random.Random(seed) + rng.shuffle(shard) + + # Save the sharding plan + os.makedirs(sharding_info_dir, exist_ok=True) + with open( + os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "wb" + ) as f: + pickle.dump(shard_parquet_files, f) + with open( + os.path.join(sharding_info_dir, "shard_total_samples.pkl"), "wb" + ) as f: + pickle.dump(shard_total_samples, f) + with open( + os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), + "wb", + ) as f: + pickle.dump(shard_parquet_lengths, f) + logger.info("Saved sharding info to %s", sharding_info_dir) + + # Wait for rank 0 to finish creating/loading sharding plan + world_group = get_world_group() + world_group.barrier() + + # Now all ranks load the sharding plan (it should exist and be valid now) + logger.info( + "Loading sharding plan from %s after barrier", sharding_info_dir + ) + with open( + os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "rb" + ) as f: + shard_parquet_files = pickle.load(f) + with open( + os.path.join(sharding_info_dir, "shard_total_samples.pkl"), "rb" + ) as f: + shard_total_samples = pickle.load(f) + with open( + os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), "rb" + ) as f: + shard_parquet_lengths = pickle.load(f) + + return shard_parquet_files, shard_total_samples, shard_parquet_lengths + + +def build_parquet_iterable_style_dataloader( + path: str, + batch_size: int, + num_data_workers: int, + cfg_rate: float = 0.0, + drop_last: bool = True, + text_padding_length: int = 512, + seed: int = 42, + read_batch_size: int = 32, +) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]: + """Build a dataloader for the LatentsParquetIterStyleDataset.""" + dataset = LatentsParquetIterStyleDataset( + path=path, + batch_size=batch_size, + cfg_rate=cfg_rate, + num_workers=num_data_workers, + drop_last=drop_last, + text_padding_length=text_padding_length, + seed=seed, + read_batch_size=read_batch_size, + ) + + loader = StatefulDataLoader( + dataset, + batch_size=1, + num_workers=num_data_workers, + pin_memory=True, + ) + return dataset, loader diff --git a/worldcompass/fastvideo/dataset/preprocessing_datasets.py b/worldcompass/fastvideo/dataset/preprocessing_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..01498202bf6155005a8f870bd6f2b05fefb1a48d --- /dev/null +++ b/worldcompass/fastvideo/dataset/preprocessing_datasets.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import math +import os +import random +from abc import ABC, abstractmethod +from collections import Counter +from dataclasses import dataclass +from os.path import join as opj +from typing import Any + +import numpy as np +import torch +import torchvision +from einops import rearrange +from PIL import Image +from transformers import AutoTokenizer + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class PreprocessBatch: + """Batch information for dataset processing stages. + + This class holds all the information about a video-caption or image-caption pair as it moves + through the processing pipeline. Fields are populated by different stages. + """ + + # Raw metadata + path: str + cap: str | list[str] + resolution: dict | None = None + fps: float | None = None + duration: float | None = None + + # Processed metadata + num_frames: int | None = None + sample_frame_index: list[int] | None = None + sample_num_frames: int | None = None + + # Processed data + pixel_values: torch.Tensor | None = None + text: str | None = None + input_ids: torch.Tensor | None = None + cond_mask: torch.Tensor | None = None + + @property + def is_video(self) -> bool: + """Check if this is a video item.""" + return self.path.endswith(".mp4") + + @property + def is_image(self) -> bool: + """Check if this is an image item.""" + return self.path.endswith(".jpg") + + +class DatasetStage(ABC): + """Abstract base class for dataset processing stages. + + Similar to PipelineStage but designed for dataset preprocessing operations. + """ + + @abstractmethod + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process the dataset batch. + + Args: + batch: Dataset batch to process + **kwargs: Additional processing parameters + + Returns: + Processed batch + """ + raise NotImplementedError + + +class DatasetFilterStage(ABC): + """Abstract base class for dataset filtering stages. + + These stages can filter out items during metadata processing. + """ + + @abstractmethod + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """Check if batch should be kept. + + Args: + batch: Dataset batch to check + **kwargs: Additional parameters + + Returns: + True if batch should be kept, False otherwise + """ + raise NotImplementedError + + @abstractmethod + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process the dataset batch (for non-filtering operations). + + Args: + batch: Dataset batch to process + **kwargs: Additional processing parameters + + Returns: + Processed batch + """ + raise NotImplementedError + + +class DataValidationStage(DatasetFilterStage): + """Stage for validating data items.""" + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """Validate data item. + + Args: + batch: Dataset batch to validate + + Returns: + True if valid, False if invalid + """ + # Check for caption + if batch.cap is None: + return False + + if batch.is_video: + # Validate video-specific fields + if batch.duration is None or batch.fps is None: + return False + elif not batch.is_image: + return False + + return True + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process does nothing for validation - filtering is handled by should_keep.""" + return batch + + +class ResolutionFilterStage(DatasetFilterStage): + """Stage for filtering data items based on resolution constraints.""" + + def __init__( + self, + max_h_div_w_ratio: float = 17 / 16, + min_h_div_w_ratio: float = 8 / 16, + max_height: int = 1024, + max_width: int = 1024, + ): + self.max_h_div_w_ratio = max_h_div_w_ratio + self.min_h_div_w_ratio = min_h_div_w_ratio + self.max_height = max_height + self.max_width = max_width + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """Check if data item passes resolution filtering. + + Args: + batch: Dataset batch with resolution information + + Returns: + True if passes filter, False otherwise + """ + # Only apply to videos + if not batch.is_video: + return True + + if batch.resolution is None: + return False + + height = batch.resolution.get("height", None) + width = batch.resolution.get("width", None) + if height is None or width is None: + return False + + # Check aspect ratio + aspect = self.max_height / self.max_width + hw_aspect_thr = 1.5 + + return self.filter_resolution( + height, + width, + max_h_div_w_ratio=hw_aspect_thr * aspect, + min_h_div_w_ratio=1 / hw_aspect_thr * aspect, + ) + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process does nothing for resolution filtering - filtering is handled by should_keep.""" + return batch + + def filter_resolution( + self, h: int, w: int, max_h_div_w_ratio: float, min_h_div_w_ratio: float + ) -> bool: + """Filter based on height/width ratio.""" + return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio + + +class FrameSamplingStage(DatasetFilterStage): + """Stage for temporal frame sampling and indexing.""" + + def __init__( + self, + num_frames: int, + train_fps: int, + speed_factor: int = 1, + video_length_tolerance_range: float = 5.0, + drop_short_ratio: float = 0.0, + seed: int = 42, + ): + self.num_frames = num_frames + self.train_fps = train_fps + self.speed_factor = speed_factor + self.video_length_tolerance_range = video_length_tolerance_range + self.drop_short_ratio = drop_short_ratio + # Create a seeded random generator for deterministic sampling + self.rng = random.Random(seed) + + def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: + """Check if video should be kept based on length constraints. + + Args: + batch: Dataset batch + + Returns: + True if should be kept, False otherwise + """ + if batch.is_image: + return True + + if batch.duration is None or batch.fps is None: + return False + + num_frames = math.ceil(batch.fps * batch.duration) + + # Check if video is too long + if num_frames / batch.fps > self.video_length_tolerance_range * ( + self.num_frames / self.train_fps * self.speed_factor + ): + return False + + # Resample frame indices to check length + frame_interval = batch.fps / self.train_fps + start_frame_idx = 0 + frame_indices = np.arange( + start_frame_idx, num_frames, frame_interval + ).astype(int) + + # Filter short videos + return not ( + len(frame_indices) < self.num_frames + and self.rng.random() < self.drop_short_ratio + ) + + def process( + self, batch: PreprocessBatch, temporal_sample_fn=None, **kwargs + ) -> PreprocessBatch: + """Process frame sampling for video data items. + + Args: + batch: Dataset batch + temporal_sample_fn: Function for temporal sampling + + Returns: + Updated batch with frame sampling info + """ + if batch.is_image: + # For images, just add sample info + batch.sample_frame_index = [0] + batch.sample_num_frames = 1 + return batch + + assert batch.duration is not None and batch.fps is not None + batch.num_frames = math.ceil(batch.fps * batch.duration) + + # Resample frame indices + frame_interval = batch.fps / self.train_fps + start_frame_idx = 0 + frame_indices = np.arange( + start_frame_idx, batch.num_frames, frame_interval + ).astype(int) + + # Temporal crop if too long + if len(frame_indices) > self.num_frames: + if temporal_sample_fn is not None: + begin_index, end_index = temporal_sample_fn(len(frame_indices)) + frame_indices = frame_indices[begin_index:end_index] + else: + frame_indices = frame_indices[: self.num_frames] + + batch.sample_frame_index = frame_indices.tolist() + batch.sample_num_frames = len(frame_indices) + + return batch + + +class VideoTransformStage(DatasetStage): + """Stage for video data transformation.""" + + def __init__(self, transform) -> None: + self.transform = transform + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Transform video data. + + Args: + batch: Dataset batch with video information + + Returns: + Batch with transformed video tensor + """ + if not batch.is_video: + return batch + + assert os.path.exists(batch.path), f"file {batch.path} do not exist!" + assert ( + batch.sample_frame_index is not None + ), "Frame indices must be set before transformation" + + torchvision_video, _, metadata = torchvision.io.read_video( + batch.path, output_format="TCHW" + ) + video = torchvision_video[batch.sample_frame_index] + if self.transform is not None: + video = self.transform(video) + video = rearrange(video, "t c h w -> c t h w") + video = video.to(torch.uint8) + + h, w = video.shape[-2:] + assert ( + h / w <= 17 / 16 and h / w >= 8 / 16 + ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}" + + video = video.float() / 127.5 - 1.0 + batch.pixel_values = video + return batch + + +class ImageTransformStage(DatasetStage): + """Stage for image data transformation.""" + + def __init__(self, transform, transform_topcrop) -> None: + self.transform = transform + self.transform_topcrop = transform_topcrop + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Transform image data. + + Args: + batch: Dataset batch with image information + + Returns: + Batch with transformed image tensor + """ + if not batch.is_image: + return batch + + image = Image.open(batch.path).convert("RGB") + image = torch.from_numpy(np.array(image)) + image = rearrange(image, "h w c -> c h w").unsqueeze(0) + + if self.transform_topcrop is not None: + image = self.transform_topcrop(image) + elif self.transform is not None: + image = self.transform(image) + + image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W] + image = image.float() / 127.5 - 1.0 + batch.pixel_values = image + return batch + + +class TextEncodingStage(DatasetStage): + """Stage for text tokenization and encoding.""" + + def __init__( + self, + tokenizer, + text_max_length: int, + cfg_rate: float = 0.0, + seed: int = 42, + ): + self.tokenizer = tokenizer + self.text_max_length = text_max_length + self.cfg_rate = cfg_rate + # Create a seeded random generator for deterministic CFG + self.rng = random.Random(seed) + + def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: + """Process text data. + + Args: + batch: Dataset batch with caption information + + Returns: + Batch with encoded text information + """ + text = batch.cap + if not isinstance(text, list): + text = [text] + text = [self.rng.choice(text)] + + text = text[0] if self.rng.random() > self.cfg_rate else "" + text_tokens_and_mask = self.tokenizer( + text, + max_length=self.text_max_length, + padding="max_length", + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors="pt", + ) + + batch.text = text + batch.input_ids = text_tokens_and_mask["input_ids"] + batch.cond_mask = text_tokens_and_mask["attention_mask"] + return batch + + +class VideoCaptionMergedDataset( + torch.utils.data.IterableDataset, + torch.distributed.checkpoint.stateful.Stateful, +): + """ + Merged dataset for video and caption data with stage-based processing. + Assumes that data_merge_path is a txt file with the following format: + , + + The folder should contain videos. + + The json file should be a list of dictionaries with the following format: + [ + { + "path": "1gGQy4nxyUo-Scene-016.mp4", + "resolution": { + "width": 1920, + "height": 1080 + }, + "size": 2439112, + "fps": 25.0, + "duration": 6.88, + "num_frames": 172, + "cap": [ + "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open." + ] + }, + ... + ] + + This dataset processes video and image data through a series of stages: + - Data validation + - Resolution filtering + - Frame sampling + - Transformation + - Text encoding + """ + + def __init__( + self, + data_merge_path: str, + args, + transform, + temporal_sample, + transform_topcrop, + start_idx: int = 0, + seed: int = 42, + ): + self.data_merge_path = data_merge_path + self.start_idx = start_idx + self.args = args + self.temporal_sample = temporal_sample + self.seed = seed + + # Initialize tokenizer + tokenizer_path = os.path.join(args.model_path, "tokenizer") + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, cache_dir=args.cache_dir + ) + + # Initialize processing stages + self._init_stages(args, transform, transform_topcrop, tokenizer) + + # Process metadata + self.processed_batches = self._process_metadata() + + def _init_stages( + self, args, transform, transform_topcrop, tokenizer + ) -> None: + """Initialize all processing stages.""" + self.validation_stage = DataValidationStage() + self.resolution_filter_stage = ResolutionFilterStage( + max_height=args.max_height, max_width=args.max_width + ) + self.frame_sampling_stage = FrameSamplingStage( + num_frames=args.num_frames, + train_fps=args.train_fps, + speed_factor=args.speed_factor, + video_length_tolerance_range=args.video_length_tolerance_range, + drop_short_ratio=args.drop_short_ratio, + seed=self.seed, + ) + self.video_transform_stage = VideoTransformStage(transform) + self.image_transform_stage = ImageTransformStage( + transform, transform_topcrop + ) + self.text_encoding_stage = TextEncodingStage( + tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=args.training_cfg_rate, + seed=self.seed, + ) + + def _load_raw_data(self) -> list[dict]: + """Load raw data from JSON files.""" + # Read folder-annotation pairs + with open(self.data_merge_path) as f: + folder_anno_pairs = [ + line.strip().split(",") for line in f if line.strip() + ] + assert ( + len(folder_anno_pairs) == 1 + ), "Only support one folder-annotation pair" + assert ( + len(folder_anno_pairs[0]) == 2 + ), "Folder-annotation pair should have two elements" + folder, annotation_file = folder_anno_pairs[0] + + data_items: list[dict] = [] + with open(annotation_file) as f: + data_items = json.load(f) + + # Update paths with folder prefix + for item in data_items: + item["path"] = opj(folder, item["path"]) + + return data_items + + def _process_metadata(self) -> list[PreprocessBatch]: + """Process the raw metadata through all filtering stages.""" + raw_data = self._load_raw_data() + processed_batches = [] + + # Initialize counters + filter_counts = { + "validation_failed": 0, + "resolution_failed": 0, + "frame_sampling_failed": 0, + } + sample_num_frames: list[int] = [] + + for item in raw_data: + batch = PreprocessBatch( + path=item["path"], + cap=item["cap"], + resolution=item.get("resolution"), + fps=item.get("fps"), + duration=item.get("duration"), + ) + + # Apply filtering stages + if not self._apply_filter_stages(batch, filter_counts): + continue + + # Apply frame sampling processing + batch = self.frame_sampling_stage.process( + batch, temporal_sample_fn=self.temporal_sample + ) + + processed_batches.append(batch) + assert batch.sample_num_frames is not None + sample_num_frames.append(batch.sample_num_frames) + + self._log_filtering_stats( + filter_counts, + sample_num_frames, + len(raw_data), + len(processed_batches), + ) + return processed_batches + + def _apply_filter_stages( + self, batch: PreprocessBatch, filter_counts: dict[str, int] + ) -> bool: + """Apply all filter stages and update counters. + + Returns True if batch should be kept. + """ + if not self.validation_stage.should_keep(batch): + filter_counts["validation_failed"] += 1 + return False + + if not self.resolution_filter_stage.should_keep(batch): + filter_counts["resolution_failed"] += 1 + return False + + if not self.frame_sampling_stage.should_keep(batch): + filter_counts["frame_sampling_failed"] += 1 + return False + + return True + + def _log_filtering_stats( + self, + filter_counts: dict[str, int], + sample_num_frames: list[int], + before_count: int, + after_count: int, + ): + """Log filtering statistics.""" + logger.info( + "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, " + "Counter(sample_num_frames): %s, before filter: %d, after filter: %d", + filter_counts["validation_failed"], + filter_counts["resolution_failed"], + filter_counts["frame_sampling_failed"], + Counter(sample_num_frames), + before_count, + after_count, + ) + + def __iter__(self): + """Iterate through processed data items.""" + for idx in range(len(self.processed_batches)): + yield self._get_item(idx) + + def __len__(self): + return len(self.processed_batches) + + def _get_item(self, idx: int) -> dict: + """Get a single processed data item.""" + batch = self.processed_batches[idx] + + # Apply transformation stages + batch = self.video_transform_stage.process(batch) + batch = self.image_transform_stage.process(batch) + batch = self.text_encoding_stage.process(batch) + + # Build result dictionary + result = { + "pixel_values": batch.pixel_values, + "text": batch.text, + "input_ids": batch.input_ids, + "cond_mask": batch.cond_mask, + "path": batch.path, + } + + # Add video-specific fields + if batch.is_video: + result.update({"fps": batch.fps, "duration": batch.duration}) + + return result + + def state_dict(self) -> dict[str, Any]: + """Return state dict for checkpointing.""" + return {"processed_batches": self.processed_batches} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load state dict from checkpoint.""" + self.processed_batches = state_dict["processed_batches"] diff --git a/worldcompass/fastvideo/dataset/transform.py b/worldcompass/fastvideo/dataset/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..cc33c91114e57080e733f7c9f607a34d7f3d3cde --- /dev/null +++ b/worldcompass/fastvideo/dataset/transform.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +import random + +import torch + + +def _is_tensor_video_clip(clip) -> bool: + if not torch.is_tensor(clip): + raise TypeError(f"clip should be Tensor. Got {type(clip)}") + + if not clip.ndimension() == 4: + raise ValueError(f"clip should be 4D. Got {clip.dim()}D") + + return True + + +def crop(clip, i, j, h, w) -> torch.Tensor: + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) + """ + if len(clip.size()) != 4: + raise ValueError("clip should be a 4D tensor") + return clip[..., i : i + h, j : j + w] + + +def resize(clip, target_size, interpolation_mode) -> torch.Tensor: + if len(target_size) != 2: + raise ValueError( + f"target size should be tuple (height, width), instead got {target_size}" + ) + return torch.nn.functional.interpolate( + clip, + size=target_size, + mode=interpolation_mode, + align_corners=True, + antialias=True, + ) + + +def center_crop_th_tw(clip, th, tw, top_crop) -> torch.Tensor: + if not _is_tensor_video_clip(clip): + raise ValueError("clip should be a 4D torch.tensor") + + # import ipdb;ipdb.set_trace() + h, w = clip.size(-2), clip.size(-1) + tr = th / tw + if h / w > tr: + new_h = int(w * tr) + new_w = w + else: + new_h = h + new_w = int(h / tr) + + i = 0 if top_crop else int(round((h - new_h) / 2.0)) + j = int(round((w - new_w) / 2.0)) + return crop(clip, i, j, new_h, new_w) + + +def normalize_video(clip) -> torch.Tensor: + """ + Convert tensor data type from uint8 to float, divide value by 255.0 and + permute the dimensions of clip tensor + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) + Return: + clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) + """ + _is_tensor_video_clip(clip) + if not clip.dtype == torch.uint8: + raise TypeError( + f"clip tensor should have data type uint8. Got {clip.dtype}" + ) + # return clip.float().permute(3, 0, 1, 2) / 255.0 + return clip.float() / 255.0 + + +class CenterCropResizeVideo: + """First use the short side for cropping length, center crop video, then resize to the specified + size.""" + + def __init__( + self, + size, + top_crop=False, + interpolation_mode="bilinear", + ) -> None: + if len(size) != 2: + raise ValueError( + f"size should be tuple (height, width), instead got {size}" + ) + self.size = size + self.top_crop = top_crop + self.interpolation_mode = interpolation_mode + + def __call__(self, clip) -> torch.Tensor: + """ + Args: + clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) + Returns: + torch.tensor: scale resized / center cropped video clip. + size is (T, C, crop_size, crop_size) + """ + clip_center_crop = center_crop_th_tw( + clip, self.size[0], self.size[1], top_crop=self.top_crop + ) + clip_center_crop_resize = resize( + clip_center_crop, + target_size=self.size, + interpolation_mode=self.interpolation_mode, + ) + return clip_center_crop_resize + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" + + +class Normalize255: + """Convert tensor data type from uint8 to float, divide value by 255.0 and.""" + + def __init__(self) -> None: + pass + + def __call__(self, clip) -> torch.Tensor: + """ + Args: + clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) + Return: + clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) + """ + return normalize_video(clip) + + def __repr__(self) -> str: + return self.__class__.__name__ + + +class TemporalRandomCrop: + """Temporally crop the given frame indices at a random location. + + Args: + size (int): Desired length of frames will be seen in the model. + """ + + def __init__(self, size) -> None: + self.size = size + + def __call__(self, total_frames) -> tuple[int, int]: + rand_end = max(0, total_frames - self.size - 1) + begin_index = random.randint(0, rand_end) + end_index = min(begin_index + self.size, total_frames) + return begin_index, end_index diff --git a/worldcompass/fastvideo/dataset/upload_to_hf.py b/worldcompass/fastvideo/dataset/upload_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..c42f666d4fb1665149728a5c4657e49dc7c101ed --- /dev/null +++ b/worldcompass/fastvideo/dataset/upload_to_hf.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +from huggingface_hub import HfApi, upload_folder + +api = HfApi() +repo_id = "weizhou03/HD-Mixkit-Finetune-Wan" # customize this +api.create_repo(repo_id=repo_id, repo_type="dataset") + +upload_folder( + repo_id=repo_id, + folder_path="/workspace/data/HD-Mixkit-Finetune-Wan", + repo_type="dataset", + path_in_repo="", +) diff --git a/worldcompass/fastvideo/dataset/utils.py b/worldcompass/fastvideo/dataset/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec7fa180c8670caddc7d53a72fbda29e3a4dcec --- /dev/null +++ b/worldcompass/fastvideo/dataset/utils.py @@ -0,0 +1,234 @@ +import random +from typing import Any, cast + +import numpy as np +import torch + + +def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor: + """Pad or crop an embedding [L, D] to exactly padding_length tokens. + + Return: + - [L, D] tensor in pinned CPU memory + - [L] attention mask in pinned CPU memory + """ + L, D = t.shape + if padding_length > L: # pad + pad = torch.zeros(padding_length - L, D, dtype=t.dtype, device=t.device) + return torch.cat([t, pad], 0), torch.cat( + [torch.ones(L), torch.zeros(padding_length - L)], 0 + ) + else: # crop + return t[:padding_length], torch.ones(padding_length) + + +def get_torch_tensors_from_row_dict( + row_dict, keys, cfg_rate, rng=None +) -> dict[str, Any]: + """Get the latents and prompts from a row dictionary.""" + return_dict = {} + for key in keys: + shape, bytes = None, None + if isinstance(key, tuple): + for k in key: + try: + shape = row_dict[f"{k}_shape"] + bytes = row_dict[f"{k}_bytes"] + except KeyError: + continue + key = key[0] + if shape is None or bytes is None: + raise ValueError(f"Key {key} not found in row_dict") + else: + shape = row_dict[f"{key}_shape"] + bytes = row_dict[f"{key}_bytes"] + + # TODO (peiyuan): read precision + if ( + key == "text_embedding" + and (rng.random() if rng else random.random()) < cfg_rate + ): + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() + data = torch.from_numpy(data) + if len(data.shape) == 3: + B, L, D = data.shape + assert B == 1, "Batch size must be 1" + data = data.squeeze(0) + return_dict[key] = data + return return_dict + + +def collate_latents_embs_masks( + batch_to_process, text_padding_length, keys, cfg_rate=0.0, rng=None +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]: + # Initialize tensors to hold padded embeddings and masks + all_latents = [] + all_embs = [] + all_masks = [] + caption_text = [] + # Process each row individually + for i, row in enumerate(batch_to_process): + # Get tensors from row + data = get_torch_tensors_from_row_dict(row, keys, cfg_rate, rng) + latents, emb = data["vae_latent"], data["text_embedding"] + + padded_emb, mask = pad(emb, text_padding_length) + # Store in batch tensors + all_latents.append(latents) + all_embs.append(padded_emb) + all_masks.append(mask) + # TODO(py): remove this once we fix preprocess + try: + caption_text.append(row["prompt"]) + except KeyError: + caption_text.append(row["caption"]) + + # Pin memory for faster transfer to GPU + all_latents = torch.stack(all_latents) + all_embs = torch.stack(all_embs) + all_masks = torch.stack(all_masks) + + return all_latents, all_embs, all_masks, caption_text + + +def collate_rows_from_parquet_schema( + rows, parquet_schema, text_padding_length, cfg_rate=0.0, rng=None +) -> dict[str, Any]: + """Collate rows from parquet files based on the provided schema. Dynamically processes tensor + fields based on schema and returns batched data. + + Args: + rows: List of row dictionaries from parquet files + parquet_schema: PyArrow schema defining the structure of the data + + Returns: + Dict containing batched tensors and metadata + """ + if not rows: + return cast(dict[str, Any], {}) + + # Initialize containers for different data types + batch_data: dict[str, Any] = {} + + # Get tensor and metadata field names from schema (fields ending with '_bytes') + tensor_fields = [] + metadata_fields = [] + for field in parquet_schema.names: + if field.endswith("_bytes"): + shape_field = field.replace("_bytes", "_shape") + dtype_field = field.replace("_bytes", "_dtype") + tensor_name = field.replace("_bytes", "") + tensor_fields.append(tensor_name) + assert ( + shape_field in parquet_schema.names + ), f"Shape field {shape_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors." + assert ( + dtype_field in parquet_schema.names + ), f"Dtype field {dtype_field} not found in schema for field {field}. Currently we only support *_bytes fields for tensors." + elif not field.endswith("_shape") and not field.endswith("_dtype"): + # Only add actual metadata fields, not the shape/dtype helper fields + metadata_fields.append(field) + + # Process each tensor field + for tensor_name in tensor_fields: + tensor_list = [] + + for row in rows: + # Get tensor data from row using the existing helper function pattern + shape_key = f"{tensor_name}_shape" + bytes_key = f"{tensor_name}_bytes" + + if shape_key in row and bytes_key in row: + shape = row[shape_key] + bytes_data = row[bytes_key] + + if len(bytes_data) == 0: + tensor = torch.zeros(0, dtype=torch.bfloat16) + else: + # Convert bytes to tensor using float32 as default + if ( + tensor_name == "text_embedding" + and (rng.random() if rng else random.random()) + < cfg_rate + ): + data = np.zeros((512, 4096), dtype=np.float32) + else: + data = ( + np.frombuffer(bytes_data, dtype=np.float32) + .reshape(shape) + .copy() + ) + tensor = torch.from_numpy(data) + # if len(data.shape) == 3: + # B, L, D = tensor.shape + # assert B == 1, "Batch size must be 1" + # tensor = tensor.squeeze(0) + + tensor_list.append(tensor) + else: + # Handle missing tensor data + tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) + + # Stack tensors with special handling for text embeddings + if tensor_name == "text_embedding": + # Handle text embeddings with padding + padded_tensors = [] + attention_masks = [] + + for tensor in tensor_list: + if tensor.numel() > 0: + padded_tensor, mask = pad(tensor, text_padding_length) + padded_tensors.append(padded_tensor) + attention_masks.append(mask) + else: + # Handle empty embeddings - assume default embedding dimension + padded_tensors.append( + torch.zeros( + text_padding_length, 768, dtype=torch.bfloat16 + ) + ) + attention_masks.append(torch.zeros(text_padding_length)) + + batch_data[tensor_name] = torch.stack(padded_tensors) + batch_data["text_attention_mask"] = torch.stack(attention_masks) + else: + # Stack all tensors to preserve batch consistency + # Don't filter out None or empty tensors as this breaks batch sizing + try: + batch_data[tensor_name] = torch.stack(tensor_list) + except ValueError as e: + shapes = [ + ( + t.shape + if t is not None and hasattr(t, "shape") + else "None/Invalid" + ) + for t in tensor_list + ] + raise ValueError( + f"Failed to stack tensors for field '{tensor_name}'. " + f"Tensor shapes: {shapes}. " + f"All tensors in a batch must have compatible shapes. " + f"Original error: {e}" + ) from e + + # Process metadata fields into info_list + info_list = [] + for row in rows: + info = {} + for field in metadata_fields: + info[field] = row.get(field, "") + + # Add prompt field for backward compatibility + info["prompt"] = info.get("caption", "") + info_list.append(info) + + batch_data["info_list"] = info_list + + # Add caption_text for backward compatibility + if info_list and "caption" in info_list[0]: + batch_data["caption_text"] = [info["caption"] for info in info_list] + + return batch_data diff --git a/worldcompass/fastvideo/dataset/validation_dataset.py b/worldcompass/fastvideo/dataset/validation_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e53049f53675172ce25d66b520024131401be5 --- /dev/null +++ b/worldcompass/fastvideo/dataset/validation_dataset.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py +import os +import pathlib + +import datasets +from torch.utils.data import IterableDataset + +from fastvideo.distributed import ( + get_sp_world_size, + get_world_rank, + get_world_size, +) +from fastvideo.logger import init_logger +from fastvideo.models.vision_utils import load_image, load_video + +logger = init_logger(__name__) + + +class ValidationDataset(IterableDataset): + def __init__(self, filename: str): + super().__init__() + + self.filename = pathlib.Path(filename) + # get directory of filename + self.dir = os.path.abspath(self.filename.parent) + + if not self.filename.exists(): + raise FileNotFoundError( + f"File {self.filename.as_posix()} does not exist" + ) + + if self.filename.suffix == ".csv": + data = datasets.load_dataset( + "csv", data_files=self.filename.as_posix(), split="train" + ) + elif self.filename.suffix == ".json": + data = datasets.load_dataset( + "json", + data_files=self.filename.as_posix(), + split="train", + field="data", + ) + elif self.filename.suffix == ".parquet": + data = datasets.load_dataset( + "parquet", data_files=self.filename.as_posix(), split="train" + ) + elif self.filename.suffix == ".arrow": + data = datasets.load_dataset( + "arrow", data_files=self.filename.as_posix(), split="train" + ) + else: + _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"] + raise ValueError( + f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}" + ) + + # Get distributed training info + self.global_rank = get_world_rank() + self.world_size = get_world_size() + self.sp_world_size = get_sp_world_size() + self.num_sp_groups = self.world_size // self.sp_world_size + + # Convert to list to get total samples + self.all_samples = list(data) + self.original_total_samples = len(self.all_samples) + + # Extend samples to be a multiple of DP degree (num_sp_groups) + remainder = self.original_total_samples % self.num_sp_groups + if remainder != 0: + samples_to_add = self.num_sp_groups - remainder + + # Duplicate samples cyclically to reach the target + additional_samples = [] + for i in range(samples_to_add): + additional_samples.append( + self.all_samples[i % self.original_total_samples] + ) + + self.all_samples.extend(additional_samples) + + self.total_samples = len(self.all_samples) + + # Calculate which SP group this rank belongs to + self.sp_group_id = self.global_rank // self.sp_world_size + + # Now all SP groups will have equal number of samples + self.samples_per_sp_group = self.total_samples // self.num_sp_groups + + # Calculate start and end indices for this SP group + self.start_idx = self.sp_group_id * self.samples_per_sp_group + self.end_idx = self.start_idx + self.samples_per_sp_group + + # Get samples for this SP group + self.sp_group_samples = self.all_samples[self.start_idx : self.end_idx] + + logger.info( + "Rank %s (SP group %s): " + "Original samples: %s, " + "Extended samples: %s, " + "SP group samples: %s, " + "Range: [%s:%s]", + self.global_rank, + self.sp_group_id, + self.original_total_samples, + self.total_samples, + len(self.sp_group_samples), + self.start_idx, + self.end_idx, + local_main_process_only=False, + ) + + def __len__(self): + """Return the number of samples for this SP group.""" + return len(self.sp_group_samples) + + def __iter__(self): + for sample in self.sp_group_samples: + # For consistency reasons, we mandate that "caption" is always present in the validation dataset. + # However, since the model specifications use "prompt", we create an alias here. + sample["prompt"] = sample["caption"] + + # Load image or video if the path is provided + # TODO(aryan): need to handle custom columns here for control conditions + sample["image"] = None + sample["video"] = None + + if sample.get("image_path", None) is not None: + image_path = sample["image_path"] + image_path = os.path.join(self.dir, image_path) + sample["image_path"] = image_path + if not pathlib.Path( + image_path + ).is_file() and not image_path.startswith("http"): + logger.warning("Image file %s does not exist.", image_path) + else: + sample["image"] = load_image(image_path) + + if sample.get("video_path", None) is not None: + video_path = sample["video_path"] + video_path = os.path.join(self.dir, video_path) + sample["video_path"] = video_path + if not pathlib.Path( + video_path + ).is_file() and not video_path.startswith("http"): + logger.warning("Video file %s does not exist.", video_path) + else: + sample["video"] = load_video(video_path) + + if sample.get("control_image_path", None) is not None: + control_image_path = sample["control_image_path"] + control_image_path = os.path.join(self.dir, control_image_path) + sample["control_image_path"] = control_image_path + if not pathlib.Path( + control_image_path + ).is_file() and not control_image_path.startswith("http"): + logger.warning( + "Control Image file %s does not exist.", + control_image_path, + ) + else: + sample["control_image"] = load_image(control_image_path) + + if sample.get("control_video_path", None) is not None: + control_video_path = sample["control_video_path"] + control_video_path = os.path.join(self.dir, control_video_path) + sample["control_video_path"] = control_video_path + if not pathlib.Path( + control_video_path + ).is_file() and not control_video_path.startswith("http"): + logger.warning( + "Control Video file %s does not exist.", + control_video_path, + ) + else: + sample["control_video"] = load_video(control_video_path) + + sample = {k: v for k, v in sample.items() if v is not None} + yield sample diff --git a/worldcompass/fastvideo/distributed/__init__.py b/worldcompass/fastvideo/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3c99da1a15d168a936abc5e9a197f81ea354d7f6 --- /dev/null +++ b/worldcompass/fastvideo/distributed/__init__.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.distributed.communication_op import * +from fastvideo.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + get_dp_group, + get_dp_rank, + get_dp_world_size, + get_local_torch_device, + get_sp_group, + get_sp_parallel_rank, + get_sp_world_size, + get_tp_group, + get_tp_rank, + get_tp_world_size, + get_world_group, + get_world_rank, + get_world_size, + init_distributed_environment, + initialize_model_parallel, + maybe_init_distributed_environment_and_model_parallel, + initialize_gpu_parallel_group, + model_parallel_is_initialized, + get_gpu_group, + get_gpu_rank, + get_gpu_world_size, +) +from fastvideo.distributed.utils import * + +__all__ = [ + # Initialization + "init_distributed_environment", + "initialize_model_parallel", + "cleanup_dist_env_and_memory", + "model_parallel_is_initialized", + "maybe_init_distributed_environment_and_model_parallel", + "initialize_gpu_parallel_group", + # World group + "get_world_group", + "get_world_rank", + "get_world_size", + # GPU parallel group + "get_gpu_group", + "get_gpu_rank", + "get_gpu_world_size", + # Data parallel group + "get_dp_group", + "get_dp_rank", + "get_dp_world_size", + # Sequence parallel group + "get_sp_group", + "get_sp_parallel_rank", + "get_sp_world_size", + # Tensor parallel group + "get_tp_group", + "get_tp_rank", + "get_tp_world_size", + # Get torch device + "get_local_torch_device", +] diff --git a/worldcompass/fastvideo/distributed/communication_op.py b/worldcompass/fastvideo/distributed/communication_op.py new file mode 100644 index 0000000000000000000000000000000000000000..fe3436b1404c526832c228aaef214e4df9339fca --- /dev/null +++ b/worldcompass/fastvideo/distributed/communication_op.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py + +import torch +import torch.distributed + +from fastvideo.distributed.parallel_state import get_sp_group, get_tp_group + + +def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + return get_tp_group().all_reduce(input_) + + +def tensor_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_tp_group().all_gather(input_, dim) + + +# TODO: remove model, make it sequence_parallel +def sequence_model_parallel_all_to_all_4D( + input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 +) -> torch.Tensor: + """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group.""" + return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim) + + +def sequence_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_sp_group().all_gather(input_, dim) diff --git a/worldcompass/fastvideo/distributed/device_communicators/__init__.py b/worldcompass/fastvideo/distributed/device_communicators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/distributed/device_communicators/base_device_communicator.py b/worldcompass/fastvideo/distributed/device_communicators/base_device_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..513295e265a22f055798ccb5e81b404142f5cc9d --- /dev/null +++ b/worldcompass/fastvideo/distributed/device_communicators/base_device_communicator.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py + +from typing import Any + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.distributed import ProcessGroup, ReduceOp + + +class DistributedAutograd: + """Collection of autograd functions for distributed operations. + + This class provides custom autograd functions for distributed operations like all_reduce, + all_gather, and all_to_all. Each operation is implemented as a static inner class with proper + forward and backward implementations. + """ + + class AllReduce(torch.autograd.Function): + """Differentiable all_reduce operation. + + The gradient of all_reduce is another all_reduce operation since the operation combines + values from all ranks equally. + """ + + @staticmethod + def forward( + ctx: Any, + group: ProcessGroup, + input_: Tensor, + op: dist.ReduceOp | None = None, + ) -> Tensor: + ctx.group = group + ctx.op = op + output = input_.clone() + dist.all_reduce(output, group=group, op=op) + return output + + @staticmethod + def backward( + ctx: Any, grad_output: Tensor + ) -> tuple[None, Tensor, None]: + grad_output = grad_output.clone() + dist.all_reduce(grad_output, group=ctx.group, op=ctx.op) + return None, grad_output, None + + class AllGather(torch.autograd.Function): + """Differentiable all_gather operation. + + The operation gathers tensors from all ranks and concatenates them along a specified + dimension. The backward pass uses reduce_scatter to efficiently distribute gradients back to + source ranks. + """ + + @staticmethod + def forward( + ctx: Any, + group: ProcessGroup, + input_: Tensor, + world_size: int, + dim: int, + ) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.dim = dim + ctx.input_shape = input_.shape + + input_size = input_.size() + output_size = (input_size[0] * world_size,) + input_size[1:] + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + + dist.all_gather_into_tensor(output_tensor, input_, group=group) + + output_tensor = output_tensor.reshape((world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + return output_tensor + + @staticmethod + def backward( + ctx: Any, grad_output: Tensor + ) -> tuple[None, Tensor, None, None]: + # Split the gradient tensor along the gathered dimension + dim_size = grad_output.size(ctx.dim) // ctx.world_size + grad_chunks = grad_output.reshape( + grad_output.shape[: ctx.dim] + + (ctx.world_size, dim_size) + + grad_output.shape[ctx.dim + 1 :] + ) + grad_chunks = grad_chunks.movedim(ctx.dim, 0) + + # Each rank only needs its corresponding gradient + grad_input = torch.empty( + ctx.input_shape, + dtype=grad_output.dtype, + device=grad_output.device, + ) + dist.reduce_scatter_tensor( + grad_input, grad_chunks.contiguous(), group=ctx.group + ) + + return None, grad_input, None, None + + class AllToAll4D(torch.autograd.Function): + """Differentiable all_to_all operation specialized for 4D tensors. + + This operation is particularly useful for attention operations where we need to + redistribute data across ranks for efficient parallel processing. + + The operation supports two modes: + 1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads + 2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions + """ + + @staticmethod + def forward( + ctx: Any, + group: ProcessGroup, + input_: Tensor, + world_size: int, + scatter_dim: int, + gather_dim: int, + ) -> Tensor: + ctx.group = group + ctx.world_size = world_size + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + + if world_size == 1: + return input_ + + assert ( + input_.dim() == 4 + ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + + if scatter_dim == 2 and gather_dim == 1: + bs, shard_seqlen, hn, hd = input_.shape + seqlen = shard_seqlen * world_size + shard_hn = hn // world_size + + input_ = input_.transpose( + 0, 2 + ).contiguous() # hn, shard_seqlen, bs, hd + output = torch.empty_like(input_) + + dist.all_to_all_single( + output, input_, group=group + ) # hn, shard_seqlen, bs, hd + + output = torch.cat( + output.split(shard_hn), dim=1 + ) # sharded hn, seqlen, bs, hd + + output = output.transpose( + 0, 2 + ).contiguous() # bs, seqlen, sharded_hn, hd + + return output + elif scatter_dim == 1 and gather_dim == 2: + bs, seqlen, shard_hn, hd = input_.shape + hn = shard_hn * world_size + shard_seqlen = seqlen // world_size + + input_ = input_.transpose( + 0, 2 + ).contiguous() # shard_hn, seqlen, bs, hd + + input_ = ( + input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd) + .transpose(0, 1) + .reshape(shard_hn * world_size, shard_seqlen, bs, hd) + .contiguous() + ) + + output = torch.empty_like(input_) + + dist.all_to_all_single(output, input_, group=group) + + output = output.transpose( + 0, 2 + ).contiguous() # bs, seqlen, sharded_hn, hd + + return output + else: + raise RuntimeError( + f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " + f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported." + ) + + @staticmethod + def backward( + ctx: Any, grad_output: Tensor + ) -> tuple[None, Tensor, None, None, None]: + if ctx.world_size == 1: + return None, grad_output, None, None, None + + # For backward pass, we swap scatter_dim and gather_dim + output = DistributedAutograd.AllToAll4D.apply( + ctx.group, + grad_output, + ctx.world_size, + ctx.gather_dim, + ctx.scatter_dim, + ) + return None, output, None, None, None + + +class DeviceCommunicatorBase: + """Base class for device-specific communicator with autograd support. + + It can use the `cpu_group` to initialize the communicator. + If the device has PyTorch integration (PyTorch can recognize its + communication backend), the `device_group` will also be given. + """ + + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + self.device = device or torch.device("cpu") + self.cpu_group = cpu_group + self.device_group = device_group + self.unique_name = unique_name + self.rank = dist.get_rank(cpu_group) + self.world_size = dist.get_world_size(cpu_group) + self.ranks = dist.get_process_group_ranks(cpu_group) + self.global_rank = dist.get_rank() + self.global_world_size = dist.get_world_size() + self.rank_in_group = dist.get_group_rank( + self.cpu_group, self.global_rank + ) + + def all_reduce( + self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM + ) -> torch.Tensor: + """Performs an all_reduce operation with gradient support.""" + return DistributedAutograd.AllReduce.apply( + self.device_group, input_, op + ) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Performs an all_gather operation with gradient support.""" + if dim < 0: + dim += input_.dim() + return DistributedAutograd.AllGather.apply( + self.device_group, input_, self.world_size, dim + ) + + def all_to_all_4D( + self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 + ) -> torch.Tensor: + """Performs a 4D all-to-all operation with gradient support.""" + return DistributedAutograd.AllToAll4D.apply( + self.device_group, input_, self.world_size, scatter_dim, gather_dim + ) + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + # Gather. + torch.distributed.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way.""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + pass diff --git a/worldcompass/fastvideo/distributed/device_communicators/cpu_communicator.py b/worldcompass/fastvideo/distributed/device_communicators/cpu_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..34173f43a814a5d219c7df0c36df3183d43cd53c --- /dev/null +++ b/worldcompass/fastvideo/distributed/device_communicators/cpu_communicator.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py + +import os + +import torch +from torch.distributed import ProcessGroup + +from fastvideo.platforms import current_platform +from fastvideo.platforms.interface import CpuArchEnum + +from .base_device_communicator import DeviceCommunicatorBase + + +class CpuCommunicator(DeviceCommunicatorBase): + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + super().__init__(cpu_group, device, device_group, unique_name) + self.dist_module = torch.distributed + + if ( + (current_platform.get_cpu_architecture() == CpuArchEnum.X86) + and hasattr(torch.ops._C, "init_shm_manager") + and unique_name.startswith("tp") + ): + self.dist_module = _CPUSHMDistributed(self) + + def all_reduce( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM, + ) -> torch.Tensor: + self.dist_module.all_reduce(input_, group=self.device_group, op=op) + return input_ + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + + # Gather. + self.dist_module.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + input_size = input_.size() + # NOTE: we have to use concat-style all-gather here, + # stack-style all-gather has compatibility issues with + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 + output_size = (input_size[0] * self.world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + # All-gather. + self.dist_module.all_gather_into_tensor( + output_tensor, input_, group=self.device_group + ) + + # Reshape + output_tensor = output_tensor.reshape((self.world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (self.world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + return output_tensor + + +class _CPUSHMDistributed: + def __init__(self, communicator: CpuCommunicator): + instance_identifier = os.environ["VLLM_DIST_IDENT"] + unique_name = communicator.unique_name + instance_identifier = f"{instance_identifier}-{unique_name}" + self.communicator = communicator + + group_ranks = [str(rank) for rank in self.communicator.ranks] + shm_group_identifier = f"[{'-'.join(group_ranks)}]" + self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm" + + self.handle = self._init_cpu_shm() + + def _init_cpu_shm(self) -> int: + handle = torch.ops._C.init_shm_manager( + self.group_name, + self.communicator.world_size, + self.communicator.rank, + ) + torch.distributed.barrier(self.communicator.device_group) + torch.ops._C.join_shm_manager( + handle, + self.group_name, + ) + torch.distributed.barrier(self.communicator.device_group) + + return int(handle) + + def all_reduce( + self, input: torch.Tensor, group: ProcessGroup | None = None + ) -> None: + torch.ops._C.shm_allreduce(self.handle, input) + + def gather( + self, + input: torch.Tensor, + gather_list: list[torch.Tensor] | None, + dst: int = -1, + group: ProcessGroup | None = None, + ) -> None: + # Note: different from the torch gather, here we use local dst rank. + torch.ops._C.shm_gather( + self.handle, + input, + gather_list, + torch.distributed.get_group_rank(group, dst), + ) + + def all_gather_into_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + group: ProcessGroup | None = None, + ) -> None: + torch.ops._C.shm_all_gather(self.handle, input, output) diff --git a/worldcompass/fastvideo/distributed/device_communicators/cuda_communicator.py b/worldcompass/fastvideo/distributed/device_communicators/cuda_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..47e8f2503013b9713b916ba554218a050c1ff085 --- /dev/null +++ b/worldcompass/fastvideo/distributed/device_communicators/cuda_communicator.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py + +import torch +from torch.distributed import ProcessGroup + +from fastvideo.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, +) + + +class CudaCommunicator(DeviceCommunicatorBase): + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + ): + super().__init__(cpu_group, device, device_group, unique_name) + + from fastvideo.distributed.device_communicators.pynccl import ( + PyNcclCommunicator, + ) + + self.pynccl_comm: PyNcclCommunicator | None = None + if self.world_size > 1: + self.pynccl_comm = PyNcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None): + pynccl_comm = self.pynccl_comm + assert pynccl_comm is not None + out = pynccl_comm.all_reduce(input_, op=op) + if out is None: + # fall back to the default all-reduce using PyTorch. + # this usually happens during testing. + # when we run the model, allreduce only happens for the TP + # group, where we always have either custom allreduce or pynccl. + out = input_.clone() + torch.distributed.all_reduce(out, group=self.device_group, op=op) + return out + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way.""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + if self.pynccl_comm is not None: + self.pynccl_comm = None diff --git a/worldcompass/fastvideo/distributed/device_communicators/pynccl.py b/worldcompass/fastvideo/distributed/device_communicators/pynccl.py new file mode 100644 index 0000000000000000000000000000000000000000..5c6ac1d68c9440919298cfe2c25ad71744c9c436 --- /dev/null +++ b/worldcompass/fastvideo/distributed/device_communicators/pynccl.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from fastvideo.distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, + buffer_type, + cudaStream_t, + ncclComm_t, + ncclDataTypeEnum, + ncclRedOpTypeEnum, + ncclUniqueId, +) +from fastvideo.distributed.utils import StatelessProcessGroup +from fastvideo.logger import init_logger +from fastvideo.utils import current_stream + +logger = init_logger(__name__) + + +class PyNcclCommunicator: + def __init__( + self, + group: ProcessGroup | StatelessProcessGroup, + device: int | str | torch.device, + library_path: str | None = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyNcclCommunicator to. If None, + it will be bind to f"cuda:{local_rank}". + library_path: the path to the NCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "PyNcclCommunicator should be attached to a non-NCCL group." + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + return + try: + self.nccl = NCCLLibrary(library_path) + except Exception: + # disable because of missing NCCL library + # e.g. in a non-GPU environment + self.available = False + self.disabled = True + return + + self.available = True + self.disabled = False + + logger.info("FastVideo is using nccl==%s", self.nccl.ncclGetVersion()) + + if self.rank == 0: + # get the unique id from NCCL + self.unique_id = self.nccl.ncclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = ncclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + # nccl communicator and stream will use this device + # `torch.cuda.device` is a context manager that changes the + # current cuda device to the specified one + with torch.cuda.device(device): + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank + ) + + stream = current_stream() + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + if stream is not None: + stream.synchronize() + del data + + def all_reduce( + self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None + ) -> torch.Tensor: + if self.disabled: + return None + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert in_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {in_tensor.device}" + ) + + out_tensor = torch.empty_like(in_tensor) + + if stream is None: + stream = current_stream() + self.nccl.ncclAllReduce( + buffer_type(in_tensor.data_ptr()), + buffer_type(out_tensor.data_ptr()), + in_tensor.numel(), + ncclDataTypeEnum.from_torch(in_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + return out_tensor + + def all_gather( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + stream=None, + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclAllGather( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + input_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def reduce_scatter( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None, + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclReduceScatter( + buffer_type(input_tensor.data_ptr()), + buffer_type(output_tensor.data_ptr()), + output_tensor.numel(), + ncclDataTypeEnum.from_torch(input_tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def send(self, tensor: torch.Tensor, dst: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclSend( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + dst, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def recv(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + self.nccl.ncclRecv( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def broadcast(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = current_stream() + if src == self.rank: + sendbuff = buffer_type(tensor.data_ptr()) + # NCCL requires the sender also to have a receive buffer + recvbuff = buffer_type(tensor.data_ptr()) + else: + sendbuff = buffer_type() + recvbuff = buffer_type(tensor.data_ptr()) + self.nccl.ncclBroadcast( + sendbuff, + recvbuff, + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) diff --git a/worldcompass/fastvideo/distributed/device_communicators/pynccl_wrapper.py b/worldcompass/fastvideo/distributed/device_communicators/pynccl_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..6738648b101f2472d6ef89e86ab4aff82bbece7f --- /dev/null +++ b/worldcompass/fastvideo/distributed/device_communicators/pynccl_wrapper.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py + +# This file is a pure Python wrapper for the NCCL library. +# The main purpose is to use NCCL combined with CUDA graph. +# Before writing this script, we tried the following approach: +# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself +# often gets stuck when initializing the NCCL communicator. +# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` +# contains many other potential cuda APIs, that are not allowed during +# capturing the CUDA graph. For further details, please check +# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . +# +# Another rejected idea is to write a C/C++ binding for NCCL. It is usually +# doable, but we often encounter issues related with nccl versions, and need +# to switch between different versions of NCCL. See +# https://github.com/NVIDIA/nccl/issues/1234 for more details. +# A C/C++ binding is not flexible enough to handle this. It requires +# recompilation of the code every time we want to switch between different +# versions. This current implementation, with a **pure** Python wrapper, is +# more flexible. We can easily switch between different versions of NCCL by +# changing the environment variable `FASTVIDEO_NCCL_SO_PATH`, or the `so_file` +# variable in the code. + +# TODO(will): support FASTVIDEO_NCCL_SO_PATH + +import ctypes +import platform +from dataclasses import dataclass +from typing import Any + +import torch +from torch.distributed import ReduceOp + +from fastvideo.logger import init_logger +from fastvideo.utils import find_nccl_library + +logger = init_logger(__name__) + +# === export types and functions from nccl to Python === +# for the original nccl definition, please check +# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in + +ncclResult_t = ctypes.c_int +ncclComm_t = ctypes.c_void_p + + +class ncclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +cudaStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +ncclDataType_t = ctypes.c_int + + +class ncclDataTypeEnum: + ncclInt8 = 0 + ncclChar = 0 + ncclUint8 = 1 + ncclInt32 = 2 + ncclInt = 2 + ncclUint32 = 3 + ncclInt64 = 4 + ncclUint64 = 5 + ncclFloat16 = 6 + ncclHalf = 6 + ncclFloat32 = 7 + ncclFloat = 7 + ncclFloat64 = 8 + ncclDouble = 8 + ncclBfloat16 = 9 + ncclNumTypes = 10 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.ncclInt8 + if dtype == torch.uint8: + return cls.ncclUint8 + if dtype == torch.int32: + return cls.ncclInt32 + if dtype == torch.int64: + return cls.ncclInt64 + if dtype == torch.float16: + return cls.ncclFloat16 + if dtype == torch.float32: + return cls.ncclFloat32 + if dtype == torch.float64: + return cls.ncclFloat64 + if dtype == torch.bfloat16: + return cls.ncclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +ncclRedOp_t = ctypes.c_int + + +class ncclRedOpTypeEnum: + ncclSum = 0 + ncclProd = 1 + ncclMax = 2 + ncclMin = 3 + ncclAvg = 4 + ncclNumOps = 5 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.ncclSum + if op == ReduceOp.PRODUCT: + return cls.ncclProd + if op == ReduceOp.MAX: + return cls.ncclMax + if op == ReduceOp.MIN: + return cls.ncclMin + if op == ReduceOp.AVG: + return cls.ncclAvg + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: list[Any] + + +class NCCLLibrary: + exported_functions = [ + # const char* ncclGetErrorString(ncclResult_t result) + Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), + # ncclResult_t ncclGetVersion(int *version); + Function( + "ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)] + ), + # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); + Function( + "ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)] + ), + # ncclResult_t ncclCommInitRank( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); + # note that ncclComm_t is a pointer type, so the first argument + # is a pointer to a pointer + Function( + "ncclCommInitRank", + ncclResult_t, + [ + ctypes.POINTER(ncclComm_t), + ctypes.c_int, + ncclUniqueId, + ctypes.c_int, + ], + ), + # ncclResult_t ncclAllReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllReduce", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclAllGather( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclAllGather", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclReduceScatter( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function( + "ncclReduceScatter", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ncclRedOp_t, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclSend( + # const void* sendbuff, size_t count, ncclDataType_t datatype, + # int dest, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclSend", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclRecv( + # void* recvbuff, size_t count, ncclDataType_t datatype, + # int src, ncclComm_t comm, cudaStream_t stream); + Function( + "ncclRecv", + ncclResult_t, + [ + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # ncclResult_t ncclBroadcast( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, int root, ncclComm_t comm, + # cudaStream_t stream); + Function( + "ncclBroadcast", + ncclResult_t, + [ + buffer_type, + buffer_type, + ctypes.c_size_t, + ncclDataType_t, + ctypes.c_int, + ncclComm_t, + cudaStream_t, + ], + ), + # be cautious! this is a collective call, it will block until all + # processes in the communicator have called this function. + # because Python object destruction can happen in random order, + # it is better not to call it at all. + # ncclResult_t ncclCommDestroy(ncclComm_t comm); + Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: dict[str, dict[str, Any]] = {} + + def __init__(self, so_file: str | None = None): + + so_file = so_file or find_nccl_library() + + try: + if so_file not in NCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + NCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = NCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + logger.error( + "Failed to load NCCL library from %s ." + "It is expected if you are not running on NVIDIA/AMD GPUs." + "Otherwise, the nccl library might not exist, be corrupted " + "or it does not support the current platform %s." + "If you already have the library, please set the " + "environment variable FASTVIDEO_NCCL_SO_PATH" + " to point to the correct nccl library path.", + so_file, + platform.platform(), + ) + raise e + + if so_file not in NCCLLibrary.path_to_dict_mapping: + _funcs: dict[str, Any] = {} + for func in NCCLLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + NCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] + + def ncclGetErrorString(self, result: ncclResult_t) -> str: + return str(self._funcs["ncclGetErrorString"](result).decode("utf-8")) + + def NCCL_CHECK(self, result: ncclResult_t) -> None: + if result != 0: + error_str = self.ncclGetErrorString(result) + raise RuntimeError(f"NCCL error: {error_str}") + + def ncclGetVersion(self) -> str: + version = ctypes.c_int() + self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) + version_str = str(version.value) + # something like 21903 --> "2.19.3" + major = version_str[0].lstrip("0") + minor = version_str[1:3].lstrip("0") + patch = version_str[3:].lstrip("0") + return f"{major}.{minor}.{patch}" + + def ncclGetUniqueId(self) -> ncclUniqueId: + unique_id = ncclUniqueId() + self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) + return unique_id + + def ncclCommInitRank( + self, world_size: int, unique_id: ncclUniqueId, rank: int + ) -> ncclComm_t: + comm = ncclComm_t() + self.NCCL_CHECK( + self._funcs["ncclCommInitRank"]( + ctypes.byref(comm), world_size, unique_id, rank + ) + ) + return comm + + def ncclAllReduce( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllReduce"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclReduceScatter( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + op: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclReduceScatter"]( + sendbuff, recvbuff, count, datatype, op, comm, stream + ) + ) + + def ncclAllGather( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + # `datatype` actually should be `ncclDataType_t` + # which is an aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK( + self._funcs["ncclAllGather"]( + sendbuff, recvbuff, count, datatype, comm, stream + ) + ) + + def ncclSend( + self, + sendbuff: buffer_type, + count: int, + datatype: int, + dest: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclSend"]( + sendbuff, count, datatype, dest, comm, stream + ) + ) + + def ncclRecv( + self, + recvbuff: buffer_type, + count: int, + datatype: int, + src: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclRecv"]( + recvbuff, count, datatype, src, comm, stream + ) + ) + + def ncclBroadcast( + self, + sendbuff: buffer_type, + recvbuff: buffer_type, + count: int, + datatype: int, + root: int, + comm: ncclComm_t, + stream: cudaStream_t, + ) -> None: + self.NCCL_CHECK( + self._funcs["ncclBroadcast"]( + sendbuff, recvbuff, count, datatype, root, comm, stream + ) + ) + + def ncclCommDestroy(self, comm: ncclComm_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) + + +__all__ = [ + "NCCLLibrary", + "ncclDataTypeEnum", + "ncclRedOpTypeEnum", + "ncclUniqueId", + "ncclComm_t", + "cudaStream_t", + "buffer_type", +] diff --git a/worldcompass/fastvideo/distributed/parallel_state.py b/worldcompass/fastvideo/distributed/parallel_state.py new file mode 100644 index 0000000000000000000000000000000000000000..fdb0af376220f8b0ba180e3aaed18347b1656ed4 --- /dev/null +++ b/worldcompass/fastvideo/distributed/parallel_state.py @@ -0,0 +1,1408 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Adapted from +"""FastVideo distributed state. +It takes over the control of the distributed environment from PyTorch. +The typical workflow is: + +- call `init_distributed_environment` to initialize the distributed environment. +- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to + initialize the model parallel groups. + +- any code dealing with the distributed stuff + +- call `destroy_model_parallel` to destroy the model parallel groups. +- call `destroy_distributed_environment` to destroy the distributed environment. + +If you only need to use the distributed environment without model parallelism, + you can skip the model parallel initialization and destruction steps. +""" + +import contextlib +import os +import pickle +import weakref +from collections import namedtuple +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from multiprocessing import shared_memory +from typing import Any, Optional +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup, ReduceOp + +import fastvideo.envs as envs +from fastvideo.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, +) +from fastvideo.distributed.device_communicators.cpu_communicator import ( + CpuCommunicator, +) +from fastvideo.distributed.utils import StatelessProcessGroup +from fastvideo.logger import init_logger +from fastvideo.platforms import current_platform + +logger = init_logger(__name__) + + +@dataclass +class GraphCaptureContext: + stream: torch.cuda.Stream | None + + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + + +def _split_tensor_dict( + tensor_dict: dict[str, torch.Tensor | Any], +) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]: + """Split the tensor dictionary into two parts: + 1. A list of (key, value) pairs. If the value is a tensor, it is replaced + by its metadata. + 2. A list of tensors. + """ + metadata_list: list[tuple[str, Any]] = [] + tensor_list: list[torch.Tensor] = [] + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor): + # Note: we cannot use `value.device` here, + # because it contains not only the device type but also the device + # index (e.g. "cuda:0"). We only need the device type. + # receiving side will set the device index. + device = value.device.type + metadata_list.append( + (key, TensorMetadata(device, value.dtype, value.size())) + ) + tensor_list.append(value) + else: + metadata_list.append((key, value)) + return metadata_list, tensor_list + + +_group_name_counter: dict[str, int] = {} + + +def _get_unique_name(name: str) -> str: + """Get a unique name for the group. + + Example: + _get_unique_name("tp") -> "tp:0" + _get_unique_name("tp") -> "tp:1" + """ + if name not in _group_name_counter: + _group_name_counter[name] = 0 + newname = f"{name}:{_group_name_counter[name]}" + _group_name_counter[name] += 1 + return newname + + +_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._all_reduce_out_place(tensor) + + +def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + return torch.empty_like(tensor) + + +class GroupCoordinator: + """PyTorch ProcessGroup wrapper for a group of processes. + + PyTorch ProcessGroup is bound to one specific communication backend, e.g. NCCL, Gloo, MPI, + etc. GroupCoordinator takes charge of all the communication operations among the processes + in the group. It manages both CPU and device communication. + """ + + # available attributes: + rank: int # global rank + ranks: list[int] # global ranks in the group + world_size: int # size of the group + # difference between `local_rank` and `rank_in_group`: + # if we have a group of size 4 across two nodes: + # Process | Node | Rank | Local Rank | Rank in Group + # 0 | 0 | 0 | 0 | 0 + # 1 | 0 | 1 | 1 | 1 + # 2 | 1 | 2 | 0 | 2 + # 3 | 1 | 3 | 1 | 3 + local_rank: int # local rank used to assign devices + rank_in_group: int # rank inside the group + cpu_group: ProcessGroup # group for CPU communication + device_group: ProcessGroup # group for device communication + use_device_communicator: bool # whether to use device communicator + device_communicator: DeviceCommunicatorBase # device communicator + mq_broadcaster: Any | None # shared memory broadcaster + + def __init__( + self, + group_ranks: list[list[int]], + local_rank: int, + torch_distributed_backend: str | Backend, + use_device_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, + ): + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + + for ranks in group_ranks: + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend + ) + # a group with `gloo` backend, to allow direct coordination between + # processes through the CPU. + cpu_group = torch.distributed.new_group(ranks, backend="gloo") + if self.rank in ranks: + self.ranks = ranks + self.world_size = len(ranks) + self.rank_in_group = ranks.index(self.rank) + self.device_group = device_group + self.cpu_group = cpu_group + try: + assert self.cpu_group is not None + assert self.device_group is not None + except Exception as e: + print(f"rank: {self.rank} group not found") + raise e + + from fastvideo.platforms import current_platform + + # TODO: fix it for other platforms + self.device = get_local_torch_device() + + self.use_device_communicator = use_device_communicator + + self.device_communicator: DeviceCommunicatorBase = None # type: ignore + if use_device_communicator and self.world_size > 1: + # Platform-aware device communicator selection + if current_platform.is_cuda_alike(): + from fastvideo.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + + self.device_communicator = CudaCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + else: + # For MPS and CPU, use the CPU communicator + self.device_communicator = CpuCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) + + self.mq_broadcaster = None + + from fastvideo.platforms import current_platform + + # TODO(will): check if this is needed + # self.use_custom_op_call = current_platform.is_cuda_alike() + self.use_custom_op_call = False + + @property + def first_rank(self): + """Return the global rank of the first process in the group.""" + return self.ranks[0] + + @property + def last_rank(self): + """Return the global rank of the last process in the group.""" + return self.ranks[-1] + + @property + def is_first_rank(self): + """Return whether the caller is the first process in the group.""" + return self.rank == self.first_rank + + @property + def is_last_rank(self): + """Return whether the caller is the last process in the group.""" + return self.rank == self.last_rank + + @property + def next_rank(self): + """Return the global rank of the process that follows the caller.""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group + 1) % world_size] + + @property + def prev_rank(self): + """Return the global rank of the process that precedes the caller.""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group - 1) % world_size] + + @contextmanager + def graph_capture( + self, graph_capture_context: GraphCaptureContext | None = None + ): + # Platform-aware graph capture + from fastvideo.platforms import current_platform + + if current_platform.is_cuda_alike(): + if graph_capture_context is None: + stream = torch.cuda.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = torch.cuda.current_stream() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with torch.cuda.stream(stream): + yield graph_capture_context + else: + # For non-CUDA platforms (MPS, CPU), just yield the context without stream management + if graph_capture_context is None: + # Create a dummy context for non-CUDA platforms + graph_capture_context = GraphCaptureContext(None) + yield graph_capture_context + + def all_reduce( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = ReduceOp.SUM, + ) -> torch.Tensor: + """User-facing all-reduce function before we actually call the all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we always make the all-reduce operation + out-of-place. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + if self.use_custom_op_call: + return torch.ops.vllm.all_reduce( + input_, group_name=self.unique_name + ) + else: + return self._all_reduce_out_place(input_, op=op) + + def _all_reduce_out_place( + self, + input_: torch.Tensor, + op: torch.distributed.ReduceOp | None = ReduceOp.SUM, + ) -> torch.Tensor: + return self.device_communicator.all_reduce(input_, op=op) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + return self.device_communicator.all_gather(input_, dim) + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + return self.device_communicator.gather(input_, dst, dim) + + def all_to_all_4D( + self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 + ) -> torch.Tensor: + if self.world_size == 1: + return input_ + return self.device_communicator.all_to_all_4D( + input_, scatter_dim, gather_dim + ) + + def broadcast(self, input_: torch.Tensor, src: int = 0): + """Broadcast the input tensor. + + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + # Broadcast. + torch.distributed.broadcast( + input_, src=self.ranks[src], group=self.device_group + ) + return input_ + + def broadcast_object(self, obj: Any | None = None, src: int = 0): + """Broadcast the input object. + + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj + if self.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_broadcaster.broadcast_object(obj) + if self.rank_in_group == src: + torch.distributed.broadcast_object_list( + [obj], src=self.ranks[src], group=self.cpu_group + ) + return obj + else: + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=self.ranks[src], group=self.cpu_group + ) + return recv[0] + + def broadcast_object_list( + self, + obj_list: list[Any], + src: int = 0, + group: ProcessGroup | None = None, + ): + """Broadcast the input object list. + + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj_list + # Broadcast. + torch.distributed.broadcast_object_list( + obj_list, src=self.ranks[src], group=self.device_group + ) + return obj_list + + def send_object(self, obj: Any, dst: int) -> None: + """Send the input object list to the destination rank.""" + """NOTE: `dst` is the local rank of the destination rank.""" + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank." + ) + + # Serialize object to tensor and get the size as well + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) + + size_tensor = torch.tensor( + [object_tensor.numel()], dtype=torch.long, device="cpu" + ) + + # Send object size + + torch.distributed.send( + size_tensor, dst=self.ranks[dst], group=self.cpu_group + ) + + # Send object + torch.distributed.send( + object_tensor, dst=self.ranks[dst], group=self.cpu_group + ) + + return None + + def recv_object(self, src: int) -> Any: + """Receive the input object list from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + + assert src < self.world_size, f"Invalid src rank ({src})" + + assert ( + src != self.rank_in_group + ), "Invalid source rank. Source rank is the same as the current rank." + + size_tensor = torch.empty(1, dtype=torch.long, device="cpu") + + # Receive object size + rank_size = torch.distributed.recv( + size_tensor, src=self.ranks[src], group=self.cpu_group + ) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device="cpu", + ) + + rank_object = torch.distributed.recv( + object_tensor, src=self.ranks[src], group=self.cpu_group + ) + + assert ( + rank_object == rank_size + ), "Received object sender rank does not match the size sender rank." + + obj = pickle.loads(object_tensor.numpy().tobytes()) + + return obj + + def broadcast_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any] | None = None, + src: int = 0, + group: ProcessGroup | None = None, + metadata_group: ProcessGroup | None = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Broadcast the input tensor dictionary. + + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + group = self.device_group + metadata_group = self.cpu_group + assert src < self.world_size, f"Invalid src rank ({src})" + + rank_in_group = self.rank_in_group + if rank_in_group == src: + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `broadcast_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.broadcast_object(metadata_list, src=src) + async_handles = [] + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + for async_handle in async_handles: + async_handle.wait() + + else: + metadata_list = self.broadcast_object(None, src=src) + tensor_dict = {} + async_handles = [] + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty( + value.size, dtype=value.dtype, device=value.device + ) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=group, + async_op=True, + ) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + for async_handle in async_handles: + async_handle.wait() + return tensor_dict + + def send_tensor_dict( + self, + tensor_dict: dict[str, torch.Tensor | Any], + dst: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Send the input tensor dictionary. + + NOTE: `dst` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + all_gather_size = ( + 1 if all_gather_group is None else all_gather_group.world_size + ) + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + metadata_list: list[tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `send_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.send_object(metadata_list, dst=dst) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ): + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.send( + tensor, dst=self.ranks[dst], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.send(tensor, dst=self.ranks[dst], group=group) + return None + + def recv_tensor_dict( + self, + src: int | None = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> dict[str, torch.Tensor | Any] | None: + """Recv the input tensor dictionary. + + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return None + + all_gather_size = ( + 1 if all_gather_group is None else all_gather_group.world_size + ) + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: dict[str, Any] = {} + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty( + value.size, dtype=value.dtype, device=value.device + ) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, -1)[ + all_gather_rank + ] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.recv( + tensor, src=self.ranks[src], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.recv( + tensor, src=self.ranks[src], group=group + ) + if use_all_gather: + # do the allgather + tensor = all_gather_group.all_gather( # type: ignore + tensor, dim=0 + ) + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + return tensor_dict + + def barrier(self) -> None: + """Barrier synchronization among the group. + + NOTE: don't use `device_group` here! `barrier` in NCCL is + terrible because it is internally a broadcast operation with + secretly created GPU tensors. It is easy to mess up the current + device. Use the CPU group instead. + """ + torch.distributed.barrier(group=self.cpu_group) + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way.""" + """NOTE: `dst` is the local rank of the destination rank.""" + self.device_communicator.send(tensor, dst) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: int | None = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + return self.device_communicator.recv(size, dtype, src) + + def destroy(self) -> None: + if self.device_group is not None: + torch.distributed.destroy_process_group(self.device_group) + self.device_group = None + if self.cpu_group is not None: + torch.distributed.destroy_process_group(self.cpu_group) + self.cpu_group = None + if self.device_communicator is not None: + self.device_communicator.destroy() + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: GroupCoordinator | None = None +_NODE: GroupCoordinator | None = None + + +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, "world group is not initialized" + return _WORLD + + +def init_world_group( + ranks: list[int], local_rank: int, backend: str +) -> GroupCoordinator: + return GroupCoordinator( + group_ranks=[ranks], + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: list[list[int]], + local_rank: int, + backend: str, + use_message_queue_broadcaster: bool = False, + group_name: str | None = None, +) -> GroupCoordinator: + + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_device_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: GroupCoordinator | None = None + + +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +_GPU: GroupCoordinator | None = None + + +def get_gpu_group() -> GroupCoordinator: + assert _GPU is not None, "gpu parallel group is not initialized" + return _GPU + + +_ENABLE_CUSTOM_ALL_REDUCE = True + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def init_distributed_environment( + world_size: int = 1, + rank: int = 0, + distributed_init_method: str = "env://", + local_rank: int = 0, + backend: str = "nccl", + device_id: torch.device | None = None, +): + # Determine the appropriate backend based on the platform + from fastvideo.platforms import current_platform + + if backend == "nccl" and not current_platform.is_cuda_alike(): + # Use gloo backend for non-CUDA platforms (MPS, CPU) + backend = "gloo" + logger.info( + "Using gloo backend for %s platform", current_platform.device_name + ) + + logger.debug( + "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s", + world_size, + rank, + local_rank, + distributed_init_method, + backend, + ) + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment" + ) + + # For MPS, don't pass device_id as it doesn't support device indices + if current_platform.is_mps(): + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + ) + else: + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + device_id=device_id, + ) + # set the local rank + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = envs.LOCAL_RANK + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert ( + _WORLD.world_size == torch.distributed.get_world_size() + ), "world group already initialized with a different world size" + + +_SP: GroupCoordinator | None = None + + +def get_sp_group() -> GroupCoordinator: + assert _SP is not None, "sequence model parallel group is not initialized" + return _SP + + +_DP: GroupCoordinator | None = None + + +def get_dp_group() -> GroupCoordinator: + assert _DP is not None, "data parallel group is not initialized" + return _DP + + +def initialize_gpu_parallel_group( + gpu_parallel_size: int = 1, + backend: str | None = None, +) -> None: + global _GPU + assert _GPU is None, "gpu parallel group is not initialized" + world_size: int = get_world_size() + num_tensor_model_parallel_groups: int = world_size // gpu_parallel_size + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list(range(i * gpu_parallel_size, (i + 1) * gpu_parallel_size)) + group_ranks.append(ranks) + + _GPU = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="gp", + ) + + return _GPU + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + sequence_model_parallel_size: int = 1, + data_parallel_size: int = 1, + backend: str | None = None, +) -> None: + """Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism (used for language encoder). + sequence_model_parallel_size: number of GPUs used for sequence model + parallelism (used for DiT). + """ + # Get world size and rank. Ensure some consistencies. + assert ( + _WORLD is not None + ), "world group is not initialized, please call init_distributed_environment first" + world_size: int = get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group + ) + assert ( + world_size >= tensor_model_parallel_size + ), f"world_size({world_size}) must be greater than or equal to tensor_model_parallel_size({tensor_model_parallel_size})" + num_tensor_model_parallel_groups: int = ( + world_size // tensor_model_parallel_size + ) + global _TP + assert _TP is None, "tensor model parallel group is already initialized" + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range( + i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size, + ) + ) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="tp", + ) + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = ( + world_size // sequence_model_parallel_size + ) + global _SP + assert _SP is None, "sequence model parallel group is already initialized" + group_ranks = [] + + # Since SP is incompatible with TP and PP, we can use a simpler group creation logic + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range( + i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size, + ) + ) + group_ranks.append(ranks) + + _SP = init_model_parallel_group( + group_ranks, get_world_group().local_rank, backend, group_name="sp" + ) + + # Build the data parallel groups. + num_data_parallel_groups: int = sequence_model_parallel_size + global _DP + assert _DP is None, "data parallel group is already initialized" + group_ranks = [] + + for i in range(num_data_parallel_groups): + ranks = list(range(i, world_size, num_data_parallel_groups)) + group_ranks.append(ranks) + + _DP = init_model_parallel_group( + group_ranks, get_world_group().local_rank, backend, group_name="dp" + ) + + +def get_sp_world_size() -> int: + """Return world size for the sequence model parallel group.""" + return get_sp_group().world_size + + +def get_sp_parallel_rank() -> int: + """Return my rank for the sequence model parallel group.""" + return get_sp_group().rank_in_group + + +def get_gpu_world_size() -> int: + """Return world size for the gpu parallel group.""" + return get_gpu_group().world_size + + +def get_gpu_rank() -> int: + """Return my rank for the gpu parallel group.""" + return get_gpu_group().rank_in_group + + +def get_world_size() -> int: + """Return world size for the world group.""" + return get_world_group().world_size + + +def get_world_rank() -> int: + """Return my rank for the world group.""" + return get_world_group().rank + + +def get_dp_world_size() -> int: + """Return world size for the data parallel group.""" + return get_dp_group().world_size + + +def get_dp_rank() -> int: + """Return my rank for the data parallel group.""" + return get_dp_group().rank_in_group + + +def get_local_torch_device() -> torch.device: + """Return the torch device for the current rank.""" + return ( + torch.device(f"cuda:{envs.LOCAL_RANK}") + if current_platform.is_cuda_alike() + else torch.device("mps") + ) + + +def maybe_init_distributed_environment_and_model_parallel( + tp_size: int, sp_size: int, distributed_init_method: str = "env://" +): + if _WORLD is not None and model_parallel_is_initialized(): + # make sure the tp and sp sizes are correct + assert ( + get_tp_world_size() == tp_size + ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" + assert ( + get_sp_world_size() == sp_size + ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + return + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + rank = int(os.environ.get("RANK", 0)) + device = get_local_torch_device() + logger.info( + "Initializing distributed environment with world_size=%d, device=%s", + world_size, + device, + ) + + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method=distributed_init_method, + device_id=device, + ) + print(tp_size, sp_size) + initialize_model_parallel( + tensor_model_parallel_size=tp_size, sequence_model_parallel_size=sp_size + ) + # Only set CUDA device if we're on a CUDA platform + if current_platform.is_cuda_alike(): + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + + +def model_parallel_is_initialized() -> bool: + """Check if tensor, sequence parallel groups are initialized.""" + return _TP is not None and _SP is not None and _DP is not None + + +_TP_STATE_PATCHED = False + + +@contextmanager +def patch_tensor_parallel_group(tp_group: GroupCoordinator): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _TP_STATE_PATCHED + assert not _TP_STATE_PATCHED, "Should not call when it's already patched" + + _TP_STATE_PATCHED = True + old_tp_group = get_tp_group() + global _TP + _TP = tp_group + try: + yield + finally: + # restore the original state + _TP_STATE_PATCHED = False + _TP = old_tp_group + + +def get_tp_world_size() -> int: + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tp_rank() -> int: + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def destroy_model_parallel() -> None: + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _SP + if _SP: + _SP.destroy() + _SP = None + + global _DP + if _DP: + _DP.destroy() + _DP = None + + +def destroy_distributed_environment() -> None: + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def cleanup_dist_env_and_memory(shutdown_ray: bool = False): + destroy_model_parallel() + destroy_distributed_environment() + with contextlib.suppress(AssertionError): + torch.distributed.destroy_process_group() + if shutdown_ray: + import ray # Lazy import Ray + + ray.shutdown() + + +def is_the_same_node_as( + pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0 +) -> list[int]: + """This is a collective operation that returns if each rank is in the same node as the source + rank. + + It tests if processes are attached to the same memory system (shared access to shared memory). + """ + if isinstance(pg, ProcessGroup): + assert ( + torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL + ), "in_the_same_node_as should be tested with a non-NCCL group." + # local rank inside the group + rank = torch.distributed.get_rank(group=pg) + world_size = torch.distributed.get_world_size(group=pg) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + else: + rank = pg.rank + world_size = pg.world_size + ranks = list(range(world_size)) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + magic_message = b"magic_message" + shm = None + + try: + with contextlib.suppress(OSError): + if rank == source_rank: + # create a shared memory segment + shm = shared_memory.SharedMemory(create=True, size=128) + shm.buf[: len(magic_message)] = magic_message + if isinstance(pg, ProcessGroup): + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg + ) + else: + pg.broadcast_obj(shm.name, src=source_rank) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + if isinstance(pg, ProcessGroup): + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg + ) + name = recv[0] + else: + name = pg.broadcast_obj(None, src=source_rank) + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[: len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as e: + logger.error("Error ignored in is_in_the_same_node: %s", e) + finally: + if shm: + shm.close() + + if isinstance(pg, ProcessGroup): + torch.distributed.barrier(group=pg) + else: + pg.barrier() + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + + if isinstance(pg, ProcessGroup): + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + aggregated_data = is_in_the_same_node + else: + aggregated_data = torch.zeros_like(is_in_the_same_node) + for i in range(world_size): + rank_data = pg.broadcast_obj(is_in_the_same_node, src=i) + aggregated_data += rank_data + + return [x == 1 for x in aggregated_data.tolist()] + + +def initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group + ) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % tensor_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = ( + world_size // tensor_model_parallel_size + ) + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range( + i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size, + ) + ) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_model_parallel_group( + tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name, + ) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group + ) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % sequence_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = ( + world_size // sequence_model_parallel_size + ) + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range( + i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size, + ) + ) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_model_parallel_group( + sp_group_ranks, + get_world_group().local_rank, + backend, + group_name=group_name, + ) + + return sp_group diff --git a/worldcompass/fastvideo/distributed/utils.py b/worldcompass/fastvideo/distributed/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..083b696e5f1f5083ded4098a5ab8e54d54f781a3 --- /dev/null +++ b/worldcompass/fastvideo/distributed/utils.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py + +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +import dataclasses +import pickle +import time +from collections import deque +from collections.abc import Sequence +from typing import Any + +import torch +from torch.distributed import TCPStore + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def ensure_divisibility(numerator, denominator) -> None: + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, "{} is not divisible by {}".format( + numerator, denominator + ) + + +def divide(numerator: int, denominator: int) -> int: + """Ensure that numerator is divisible by the denominator and return the division value.""" + ensure_divisibility(numerator, denominator) + return numerator // denominator + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> Sequence[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = divide(tensor.size()[last_dim], num_partitions) + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # NOTE: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tuple(tensor_list) + + +@dataclasses.dataclass +class StatelessProcessGroup: + """A dataclass to hold a metadata store, and the rank, world_size of the group. + + Only use it to communicate metadata between processes. For data-plane communication, create + NCCL-related objects. + """ + + rank: int + world_size: int + store: torch._C._distributed_c10d.Store + data_expiration_seconds: int = 3600 # 1 hour + + # dst rank -> counter + send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict) + # src rank -> counter + recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) + broadcast_send_counter: int = 0 + broadcast_recv_src_counter: dict[int, int] = dataclasses.field( + default_factory=dict + ) + + # A deque to store the data entries, with key and timestamp. + entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque) + + def __post_init__(self): + assert self.rank < self.world_size + self.send_dst_counter = {i: 0 for i in range(self.world_size)} + self.recv_src_counter = {i: 0 for i in range(self.world_size)} + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def send_obj(self, obj: Any, dst: int): + """Send an object to a destination rank.""" + self.expire_data() + key = f"send_to/{dst}/{self.send_dst_counter[dst]}" + self.store.set(key, pickle.dumps(obj)) + self.send_dst_counter[dst] += 1 + self.entries.append((key, time.perf_counter())) + + def expire_data(self) -> None: + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.perf_counter() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def recv_obj(self, src: int) -> Any: + """Receive an object from a source rank.""" + obj = pickle.loads( + self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}") + ) + self.recv_src_counter[src] += 1 + return obj + + def broadcast_obj(self, obj: Any | None, src: int) -> Any: + """Broadcast an object from a source rank to all other ranks. + + It does not clean up after all ranks have received the object. Use it for limited times, + e.g., for initialization. + """ + if self.rank == src: + self.expire_data() + key = f"broadcast_from/{src}/{self.broadcast_send_counter}" + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.perf_counter())) + return obj + else: + key = f"broadcast_from/{src}/{self.broadcast_recv_src_counter[src]}" + recv_obj = pickle.loads(self.store.get(key)) + self.broadcast_recv_src_counter[src] += 1 + return recv_obj + + def all_gather_obj(self, obj: Any) -> list[Any]: + """All gather an object from all ranks.""" + gathered_objs = [] + for i in range(self.world_size): + if i == self.rank: + gathered_objs.append(obj) + self.broadcast_obj(obj, src=self.rank) + else: + recv_obj = self.broadcast_obj(None, src=i) + gathered_objs.append(recv_obj) + return gathered_objs + + def barrier(self): + """A barrier to synchronize all ranks.""" + for i in range(self.world_size): + if i == self.rank: + self.broadcast_obj(None, src=self.rank) + else: + self.broadcast_obj(None, src=i) + + @staticmethod + def create( + host: str, + port: int, + rank: int, + world_size: int, + data_expiration_seconds: int = 3600, + ) -> "StatelessProcessGroup": + """A replacement for `torch.distributed.init_process_group` that does not pollute the global + state. + + If we have process A and process B called `torch.distributed.init_process_group` + to form a group, and then we want to form another group with process A, B, C, + D, it is not possible in PyTorch, because process A and process B have already + formed a group, and process C and process D cannot join that group. This + function is a workaround for this issue. + + `torch.distributed.init_process_group` is a global call, while this function + is a stateless call. It will return a `StatelessProcessGroup` object that can be + used for exchanging metadata. With this function, process A and process B + can call `StatelessProcessGroup.create` to form a group, and then process A, B, + C, and D can call `StatelessProcessGroup.create` to form another group. + """ # noqa + store = TCPStore( + host_name=host, + port=port, + world_size=world_size, + is_master=(rank == 0), + ) + + return StatelessProcessGroup( + rank=rank, + world_size=world_size, + store=store, + data_expiration_seconds=data_expiration_seconds, + ) diff --git a/worldcompass/fastvideo/entrypoints/__init__.py b/worldcompass/fastvideo/entrypoints/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/entrypoints/cli/__init__.py b/worldcompass/fastvideo/entrypoints/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/entrypoints/cli/cli_types.py b/worldcompass/fastvideo/entrypoints/cli/cli_types.py new file mode 100644 index 0000000000000000000000000000000000000000..c28919bd74c3b8deddef87e5d2ef974154b5f316 --- /dev/null +++ b/worldcompass/fastvideo/entrypoints/cli/cli_types.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py + +import argparse + +from fastvideo.utils import FlexibleArgumentParser + + +class CLISubcommand: + """Base class for CLI subcommands.""" + + name: str + + def cmd(self, args: argparse.Namespace) -> None: + """Execute the command with the given arguments.""" + raise NotImplementedError + + def validate(self, args: argparse.Namespace) -> None: + """Validate the arguments for this command.""" + pass + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + """Initialize the subparser for this command.""" + raise NotImplementedError diff --git a/worldcompass/fastvideo/entrypoints/cli/generate.py b/worldcompass/fastvideo/entrypoints/cli/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..53b546dd29212005baf915179854b3c35278e9c1 --- /dev/null +++ b/worldcompass/fastvideo/entrypoints/cli/generate.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py + +import argparse +import dataclasses +import os +from typing import cast + +from fastvideo import VideoGenerator +from fastvideo.configs.sample.base import SamplingParam +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.cli.utils import RaiseNotImplementedAction +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser + +logger = init_logger(__name__) + + +class GenerateSubcommand(CLISubcommand): + """The `generate` subcommand for the FastVideo CLI.""" + + def __init__(self) -> None: + self.name = "generate" + super().__init__() + self.init_arg_names = self._get_init_arg_names() + self.generation_arg_names = self._get_generation_arg_names() + + def _get_init_arg_names(self) -> list[str]: + """Get names of arguments for VideoGenerator initialization.""" + return ["num_gpus", "tp_size", "sp_size", "model_path"] + + def _get_generation_arg_names(self) -> list[str]: + """Get names of arguments for generate_video method.""" + return [field.name for field in dataclasses.fields(SamplingParam)] + + def cmd(self, args: argparse.Namespace) -> None: + excluded_args = ["subparser", "config", "dispatch_function"] + + provided_args = {} + for k, v in vars(args).items(): + if ( + k not in excluded_args + and v is not None + and hasattr(args, "_provided") + and k in args._provided + ): + provided_args[k] = v + + if "model_path" in vars(args) and args.model_path is not None: + provided_args["model_path"] = args.model_path + + if "prompt" in vars(args) and args.prompt is not None: + provided_args["prompt"] = args.prompt + + merged_args = {**provided_args} + + logger.info("CLI Args: %s", merged_args) + + if "model_path" not in merged_args or not merged_args["model_path"]: + raise ValueError( + "model_path must be provided either in config file or via --model-path" + ) + + # Check if either prompt or prompt_txt is provided + has_prompt = "prompt" in merged_args and merged_args["prompt"] + has_prompt_txt = ( + "prompt_txt" in merged_args and merged_args["prompt_txt"] + ) + + if not (has_prompt or has_prompt_txt): + raise ValueError("Either prompt or prompt_txt must be provided") + + if has_prompt and has_prompt_txt: + raise ValueError( + "Cannot provide both 'prompt' and 'prompt_txt'. Use only one of them." + ) + + init_args = { + k: v + for k, v in merged_args.items() + if k not in self.generation_arg_names + } + generation_args = { + k: v + for k, v in merged_args.items() + if k in self.generation_arg_names + } + + model_path = init_args.pop("model_path") + prompt = generation_args.pop("prompt", None) + + generator = VideoGenerator.from_pretrained( + model_path=model_path, **init_args + ) + + # Call generate_video - it handles both single and batch modes + generator.generate_video(prompt=prompt, **generation_args) + + def validate(self, args: argparse.Namespace) -> None: + """Validate the arguments for this command.""" + if args.num_gpus is not None and args.num_gpus <= 0: + raise ValueError("Number of gpus must be positive") + + if args.config and not os.path.exists(args.config): + raise ValueError(f"Config file not found: {args.config}") + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + generate_parser = subparsers.add_parser( + "generate", + help="Run inference on a model", + usage="fastvideo generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]", + ) + + generate_parser.add_argument( + "--config", + type=str, + default="", + required=False, + help="Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional.", + ) + + generate_parser = FastVideoArgs.add_cli_args(generate_parser) + generate_parser = SamplingParam.add_cli_args(generate_parser) + + generate_parser.add_argument( + "--text-encoder-configs", + action=RaiseNotImplementedAction, + help="JSON array of text encoder configurations (NOT YET IMPLEMENTED)", + ) + + return cast(FlexibleArgumentParser, generate_parser) + + +def cmd_init() -> list[CLISubcommand]: + return [GenerateSubcommand()] diff --git a/worldcompass/fastvideo/entrypoints/cli/main.py b/worldcompass/fastvideo/entrypoints/cli/main.py new file mode 100644 index 0000000000000000000000000000000000000000..173c376ce33dfe7844066485bddd0704d413de80 --- /dev/null +++ b/worldcompass/fastvideo/entrypoints/cli/main.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py + +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.cli.generate import cmd_init as generate_cmd_init +from fastvideo.utils import FlexibleArgumentParser + + +def cmd_init() -> list[CLISubcommand]: + """Initialize all commands from separate modules.""" + commands = [] + commands.extend(generate_cmd_init()) + return commands + + +def main() -> None: + parser = FlexibleArgumentParser(description="FastVideo CLI") + parser.add_argument("-v", "--version", action="version", version="0.1.0") + + subparsers = parser.add_subparsers(required=False, dest="subparser") + + cmds = {} + for cmd in cmd_init(): + cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) + cmds[cmd.name] = cmd + args = parser.parse_args() + if args.subparser in cmds: + cmds[args.subparser].validate(args) + + if hasattr(args, "dispatch_function"): + args.dispatch_function(args) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/worldcompass/fastvideo/entrypoints/cli/utils.py b/worldcompass/fastvideo/entrypoints/cli/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..af53f8527ab86e1d8078e29f58a49eef4dc7d6b8 --- /dev/null +++ b/worldcompass/fastvideo/entrypoints/cli/utils.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import subprocess +import sys + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class RaiseNotImplementedAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + raise NotImplementedError( + f"The {option_string} option is not yet implemented" + ) + + +def launch_distributed( + num_gpus: int, args: list[str], master_port: int | None = None +) -> int: + """Launch a distributed job with the given arguments. + + Args: + num_gpus: Number of GPUs to use + args: Arguments to pass to v1_fastvideo_inference.py (defaults to sys.argv[1:]) + master_port: Port for the master process (default: random) + """ + + current_env = os.environ.copy() + python_executable = sys.executable + project_root = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../..") + ) + main_script = os.path.join( + project_root, "fastvideo/sample/v1_fastvideo_inference.py" + ) + + cmd = [ + python_executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={num_gpus}", + ] + + if master_port is not None: + cmd.append(f"--master_port={master_port}") + + cmd.append(main_script) + cmd.extend(args) + + logger.info("Running inference with %d GPU(s)", num_gpus) + logger.info("Launching command: %s", " ".join(cmd)) + + current_env["PYTHONIOENCODING"] = "utf-8" + process = subprocess.Popen( + cmd, + env=current_env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1, + encoding="utf-8", + errors="replace", + ) + + if process.stdout: + for line in iter(process.stdout.readline, ""): + print(line.strip()) + + return process.wait() diff --git a/worldcompass/fastvideo/entrypoints/video_generator.py b/worldcompass/fastvideo/entrypoints/video_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..e38c72f4663d0c517027e8ae66b03fd2c63eaac7 --- /dev/null +++ b/worldcompass/fastvideo/entrypoints/video_generator.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VideoGenerator module for FastVideo. + +This module provides a consolidated interface for generating videos using diffusion models. +""" + +import math +import os +import time +from copy import deepcopy +from typing import Any + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +from fastvideo.configs.sample import SamplingParam +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ForwardBatch +from fastvideo.utils import align_to, shallow_asdict +from fastvideo.worker.executor import Executor + +logger = init_logger(__name__) + + +class VideoGenerator: + """A unified class for generating videos using diffusion models. + + This class provides a simple interface for video generation with rich customization options, + similar to popular frameworks like HF Diffusers. + """ + + def __init__( + self, + fastvideo_args: FastVideoArgs, + executor_class: type[Executor], + log_stats: bool, + ): + """Initialize the video generator. + + Args: + fastvideo_args: The inference arguments + executor_class: The executor class to use for inference + """ + self.fastvideo_args = fastvideo_args + self.executor = executor_class(fastvideo_args) + + @classmethod + def from_pretrained( + cls, + model_path: str, + device: str | None = None, + torch_dtype: torch.dtype | None = None, + **kwargs, + ) -> "VideoGenerator": + """Create a video generator from a pretrained model. + + Args: + model_path: Path or identifier for the pretrained model + device: Device to load the model on (e.g., "cuda", "cuda:0", "cpu") + torch_dtype: Data type for model weights (e.g., torch.float16) + pipeline_config: Pipeline config to use for inference + **kwargs: Additional arguments to customize model loading, set any FastVideoArgs or PipelineConfig attributes here. + + Returns: + The created video generator + + Priority level: Default pipeline config < User's pipeline config < User's kwargs + """ + # If users also provide some kwargs, it will override the FastVideoArgs and PipelineConfig. + kwargs["model_path"] = model_path + fastvideo_args = FastVideoArgs.from_kwargs(**kwargs) + + return cls.from_fastvideo_args(fastvideo_args) + + @classmethod + def from_fastvideo_args( + cls, fastvideo_args: FastVideoArgs + ) -> "VideoGenerator": + """Create a video generator with the specified arguments. + + Args: + fastvideo_args: The inference arguments + + Returns: + The created video generator + """ + # Initialize distributed environment if needed + # initialize_distributed_and_parallelism(fastvideo_args) + + executor_class = Executor.get_class(fastvideo_args) + return cls( + fastvideo_args=fastvideo_args, + executor_class=executor_class, + log_stats=False, # TODO: implement + ) + + def generate_video( + self, + prompt: str | None = None, + sampling_param: SamplingParam | None = None, + **kwargs, + ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]: + """Generate a video based on the given prompt. + + Args: + prompt: The prompt to use for generation (optional if prompt_txt is provided) + negative_prompt: The negative prompt to use (overrides the one in fastvideo_args) + output_path: Path to save the video (overrides the one in fastvideo_args) + output_video_name: Name of the video file to save. Default is the first 100 characters of the prompt. + save_video: Whether to save the video to disk + return_frames: Whether to return the raw frames + num_inference_steps: Number of denoising steps (overrides fastvideo_args) + guidance_scale: Classifier-free guidance scale (overrides fastvideo_args) + num_frames: Number of frames to generate (overrides fastvideo_args) + height: Height of generated video (overrides fastvideo_args) + width: Width of generated video (overrides fastvideo_args) + fps: Frames per second for saved video (overrides fastvideo_args) + seed: Random seed for generation (overrides fastvideo_args) + callback: Callback function called after each step + callback_steps: Number of steps between each callback + + Returns: + Either the output dictionary, list of frames, or list of results for batch processing + """ + # Handle batch processing from text file + if self.fastvideo_args.prompt_txt is not None: + prompt_txt_path = self.fastvideo_args.prompt_txt + if not os.path.exists(prompt_txt_path): + raise FileNotFoundError( + f"Prompt text file not found: {prompt_txt_path}" + ) + + # Read prompts from file + with open(prompt_txt_path, encoding="utf-8") as f: + prompts = [line.strip() for line in f if line.strip()] + + if not prompts: + raise ValueError(f"No prompts found in file: {prompt_txt_path}") + + logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path) + + if sampling_param is not None: + original_output_video_name = sampling_param.output_video_name + else: + original_output_video_name = None + + results = [] + for i, batch_prompt in enumerate(prompts): + logger.info( + "Processing prompt %d/%d: %s...", + i + 1, + len(prompts), + batch_prompt[:100], + ) + + try: + # Generate video for this prompt using the same logic below + if ( + sampling_param is not None + and original_output_video_name is not None + ): + sampling_param.output_video_name = ( + original_output_video_name + f"_{i}" + ) + result = self._generate_single_video( + batch_prompt, sampling_param, **kwargs + ) + + # Add prompt info to result + if isinstance(result, dict): + result["prompt_index"] = i + result["prompt"] = batch_prompt + + results.append(result) + logger.info( + "Successfully generated video for prompt %d", i + 1 + ) + + except Exception as e: + logger.error( + "Failed to generate video for prompt %d: %s", i + 1, e + ) + continue + + logger.info( + "Completed batch processing. Generated %d videos successfully.", + len(results), + ) + return results + + # Single prompt generation (original behavior) + if prompt is None: + raise ValueError("Either prompt or prompt_txt must be provided") + + return self._generate_single_video(prompt, sampling_param, **kwargs) + + def _generate_single_video( + self, + prompt: str, + sampling_param: SamplingParam | None = None, + **kwargs, + ) -> dict[str, Any] | list[np.ndarray]: + """Internal method for single video generation.""" + # Create a copy of inference args to avoid modifying the original + fastvideo_args = self.fastvideo_args + pipeline_config = fastvideo_args.pipeline_config + + # Validate inputs + if not isinstance(prompt, str): + raise TypeError( + f"`prompt` must be a string, but got {type(prompt)}" + ) + prompt = prompt.strip() + if sampling_param is None: + sampling_param = SamplingParam.from_pretrained( + fastvideo_args.model_path + ) + else: + sampling_param = deepcopy(sampling_param) + + kwargs["prompt"] = prompt + sampling_param.update(kwargs) + + # Process negative prompt + if sampling_param.negative_prompt is not None: + sampling_param.negative_prompt = ( + sampling_param.negative_prompt.strip() + ) + + # Validate dimensions + if ( + sampling_param.height <= 0 + or sampling_param.width <= 0 + or sampling_param.num_frames <= 0 + ): + raise ValueError( + f"Height, width, and num_frames must be positive integers, got " + f"height={sampling_param.height}, width={sampling_param.width}, " + f"num_frames={sampling_param.num_frames}" + ) + + temporal_scale_factor = ( + pipeline_config.vae_config.arch_config.temporal_compression_ratio + ) + num_frames = sampling_param.num_frames + num_gpus = fastvideo_args.num_gpus + use_temporal_scaling_frames = ( + pipeline_config.vae_config.use_temporal_scaling_frames + ) + + # Adjust number of frames based on number of GPUs + if use_temporal_scaling_frames: + orig_latent_num_frames = ( + num_frames - 1 + ) // temporal_scale_factor + 1 + else: # stepvideo only + orig_latent_num_frames = sampling_param.num_frames // 17 * 3 + + if orig_latent_num_frames % fastvideo_args.num_gpus != 0: + # Adjust latent frames to be divisible by number of GPUs + if sampling_param.num_frames_round_down: + # Ensure we have at least 1 batch per GPU + new_latent_num_frames = ( + max(1, (orig_latent_num_frames // num_gpus)) * num_gpus + ) + else: + new_latent_num_frames = ( + math.ceil(orig_latent_num_frames / num_gpus) * num_gpus + ) + + if use_temporal_scaling_frames: + # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor + new_num_frames = ( + new_latent_num_frames - 1 + ) * temporal_scale_factor + 1 + else: # stepvideo only + # Find the least common multiple of 3 and num_gpus + divisor = math.lcm(3, num_gpus) + # Round up to the nearest multiple of this LCM + new_latent_num_frames = ( + (new_latent_num_frames + divisor - 1) // divisor + ) * divisor + # Convert back to actual frames using the StepVideo formula + new_num_frames = new_latent_num_frames // 3 * 17 + + logger.info( + "Adjusting number of frames from %s to %s based on number of GPUs (%s)", + sampling_param.num_frames, + new_num_frames, + fastvideo_args.num_gpus, + ) + sampling_param.num_frames = new_num_frames + + # Calculate sizes + target_height = align_to(sampling_param.height, 16) + target_width = align_to(sampling_param.width, 16) + + # Calculate latent sizes + latents_size = [ + (sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, + sampling_param.width // 8, + ] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + # Log parameters + debug_str = f""" + height: {target_height} + width: {target_width} + video_length: {sampling_param.num_frames} + prompt: {prompt} + image_path: {sampling_param.image_path} + neg_prompt: {sampling_param.negative_prompt} + seed: {sampling_param.seed} + infer_steps: {sampling_param.num_inference_steps} + num_videos_per_prompt: {sampling_param.num_videos_per_prompt} + guidance_scale: {sampling_param.guidance_scale} + n_tokens: {n_tokens} + flow_shift: {fastvideo_args.pipeline_config.flow_shift} + embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale} + save_video: {sampling_param.save_video} + output_path: {sampling_param.output_path} + """ # type: ignore[attr-defined] + logger.info(debug_str) + + # Prepare batch + batch = ForwardBatch( + **shallow_asdict(sampling_param), + eta=0.0, + n_tokens=n_tokens, + VSA_sparsity=fastvideo_args.VSA_sparsity, + extra={}, + ) + + # Use prompt[:100] for video name + if batch.output_video_name is None: + batch.output_video_name = prompt[:100] + + # Run inference + start_time = time.perf_counter() + output_batch = self.executor.execute_forward(batch, fastvideo_args) + samples = output_batch.output + logging_info = output_batch.logging_info + + gen_time = time.perf_counter() - start_time + logger.info("Generated successfully in %.2f seconds", gen_time) + + # Process outputs + videos = rearrange(samples, "b c t h w -> t b c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + # Save video if requested + if batch.save_video: + output_path = batch.output_path + if output_path: + os.makedirs(output_path, exist_ok=True) + video_path = os.path.join( + output_path, f"{batch.output_video_name}.mp4" + ) + imageio.mimsave(video_path, frames, fps=batch.fps, format="mp4") + logger.info("Saved video to %s", video_path) + else: + logger.warning("No output path provided, video not saved") + + if batch.return_frames: + return frames + else: + return { + "samples": samples, + "frames": frames, + "prompts": prompt, + "size": (target_height, target_width, batch.num_frames), + "generation_time": gen_time, + "logging_info": logging_info, + } + + def set_lora_adapter( + self, lora_nickname: str, lora_path: str | None = None + ) -> None: + self.executor.set_lora_adapter(lora_nickname, lora_path) + + def shutdown(self): + """Shutdown the video generator.""" + self.executor.shutdown() + del self.executor diff --git a/worldcompass/fastvideo/envs.py b/worldcompass/fastvideo/envs.py new file mode 100644 index 0000000000000000000000000000000000000000..3f10538f4f7a7154951df5b9150a025be3644947 --- /dev/null +++ b/worldcompass/fastvideo/envs.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60 + FASTVIDEO_NCCL_SO_PATH: str | None = None + LD_LIBRARY_PATH: str | None = None + LOCAL_RANK: int = 0 + CUDA_VISIBLE_DEVICES: str | None = None + FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo") + FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo") + FASTVIDEO_CONFIGURE_LOGGING: int = 1 + FASTVIDEO_LOGGING_LEVEL: str = "INFO" + FASTVIDEO_LOGGING_PREFIX: str = "" + FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None + FASTVIDEO_TRACE_FUNCTION: int = 0 + FASTVIDEO_ATTENTION_BACKEND: str | None = None + FASTVIDEO_ATTENTION_CONFIG: str | None = None + FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork" + FASTVIDEO_TARGET_DEVICE: str = "cuda" + MAX_JOBS: str | None = None + NVCC_THREADS: str | None = None + CMAKE_BUILD_TYPE: str | None = None + VERBOSE: bool = False + FASTVIDEO_SERVER_DEV_MODE: bool = False + FASTVIDEO_STAGE_LOGGING: bool = False + + +def get_default_cache_root() -> str: + return os.getenv( + "XDG_CACHE_HOME", + os.path.join(os.path.expanduser("~"), ".cache"), + ) + + +def get_default_config_root() -> str: + return os.getenv( + "XDG_CONFIG_HOME", + os.path.join(os.path.expanduser("~"), ".config"), + ) + + +def maybe_convert_int(value: str | None) -> int | None: + if value is None: + return None + return int(value) + + +# The begin-* and end* here are used by the documentation generator +# to extract the used env vars. + +# begin-env-vars-definition + +environment_variables: dict[str, Callable[[], Any]] = { + # ================== Installation Time Env Vars ================== + # Target device of FastVideo, supporting [cuda (by default), + # rocm, neuron, cpu, openvino] + "FASTVIDEO_TARGET_DEVICE": lambda: os.getenv( + "FASTVIDEO_TARGET_DEVICE", "cuda" + ), + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + # Number of threads to use for nvcc + # By default this is 1. + # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. + "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), + # If set, fastvideo will use precompiled binaries (*.so) + "FASTVIDEO_USE_PRECOMPILED": lambda: ( + bool(os.environ.get("FASTVIDEO_USE_PRECOMPILED")) + or bool(os.environ.get("FASTVIDEO_PRECOMPILED_WHEEL_LOCATION")) + ), + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"), + # If set, fastvideo will print verbose logs during installation + "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), + # Root directory for FASTVIDEO configuration files + # Defaults to `~/.config/fastvideo` unless `XDG_CONFIG_HOME` is set + # Note that this not only affects how fastvideo finds its configuration files + # during runtime, but also affects how fastvideo installs its configuration + # files during **installation**. + "FASTVIDEO_CONFIG_ROOT": lambda: os.path.expanduser( + os.getenv( + "FASTVIDEO_CONFIG_ROOT", + os.path.join(get_default_config_root(), "fastvideo"), + ) + ), + # ================== Runtime Env Vars ================== + # Root directory for FASTVIDEO cache files + # Defaults to `~/.cache/fastvideo` unless `XDG_CACHE_HOME` is set + "FASTVIDEO_CACHE_ROOT": lambda: os.path.expanduser( + os.getenv( + "FASTVIDEO_CACHE_ROOT", + os.path.join(get_default_cache_root(), "fastvideo"), + ) + ), + # Interval in seconds to log a warning message when the ring buffer is full + "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL": lambda: int( + os.environ.get("FASTVIDEO_RINGBUFFER_WARNING_INTERVAL", "60") + ), + # Path to the NCCL library file. It is needed because nccl>=2.19 brought + # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 + "FASTVIDEO_NCCL_SO_PATH": lambda: os.environ.get( + "FASTVIDEO_NCCL_SO_PATH", None + ), + # when `FASTVIDEO_NCCL_SO_PATH` is not set, fastvideo will try to find the nccl + # library file in the locations specified by `LD_LIBRARY_PATH` + "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None), + # Internal flag to enable Dynamo fullgraph capture + "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool( + os.environ.get("FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0" + ), + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": lambda: os.environ.get( + "CUDA_VISIBLE_DEVICES", None + ), + # timeout for each iteration in the engine + "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S": lambda: int( + os.environ.get("FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S", "60") + ), + # Logging configuration + # If set to 0, fastvideo will not configure logging + # If set to 1, fastvideo will configure logging using the default configuration + # or the configuration file specified by FASTVIDEO_LOGGING_CONFIG_PATH + "FASTVIDEO_CONFIGURE_LOGGING": lambda: int( + os.getenv("FASTVIDEO_CONFIGURE_LOGGING", "1") + ), + "FASTVIDEO_LOGGING_CONFIG_PATH": lambda: os.getenv( + "FASTVIDEO_LOGGING_CONFIG_PATH" + ), + # this is used for configuring the default logging level + "FASTVIDEO_LOGGING_LEVEL": lambda: os.getenv( + "FASTVIDEO_LOGGING_LEVEL", "INFO" + ), + # if set, FASTVIDEO_LOGGING_PREFIX will be prepended to all log messages + "FASTVIDEO_LOGGING_PREFIX": lambda: os.getenv( + "FASTVIDEO_LOGGING_PREFIX", "" + ), + # Trace function calls + # If set to 1, fastvideo will trace function calls + # Useful for debugging + "FASTVIDEO_TRACE_FUNCTION": lambda: int( + os.getenv("FASTVIDEO_TRACE_FUNCTION", "0") + ), + # Backend for attention computation + # Available options: + # - "TORCH_SDPA": use torch.nn.MultiheadAttention + # - "FLASH_ATTN": use FlashAttention + # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention + # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention + # - "SAGE_ATTN": use Sage Attention + "FASTVIDEO_ATTENTION_BACKEND": lambda: os.getenv( + "FASTVIDEO_ATTENTION_BACKEND", None + ), + # Path to the attention configuration file. Only used for sliding tile + # attention for now. + "FASTVIDEO_ATTENTION_CONFIG": lambda: ( + None + if os.getenv("FASTVIDEO_ATTENTION_CONFIG", None) is None + else os.path.expanduser(os.getenv("FASTVIDEO_ATTENTION_CONFIG", ".")) + ), + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "FASTVIDEO_WORKER_MULTIPROC_METHOD": lambda: os.getenv( + "FASTVIDEO_WORKER_MULTIPROC_METHOD", "fork" + ), + # Enables torch profiler if set. Path to the directory where torch profiler + # traces are saved. Note that it must be an absolute path. + "FASTVIDEO_TORCH_PROFILER_DIR": lambda: ( + None + if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None + else os.path.expanduser(os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", ".")) + ), + # If set, fastvideo will run in development mode, which will enable + # some additional endpoints for developing and debugging, + # e.g. `/reset_prefix_cache` + "FASTVIDEO_SERVER_DEV_MODE": lambda: bool( + int(os.getenv("FASTVIDEO_SERVER_DEV_MODE", "0")) + ), + # If set, fastvideo will enable stage logging, which will print the time + # taken for each stage + "FASTVIDEO_STAGE_LOGGING": lambda: bool( + int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0")) + ), +} + +# end-env-vars-definition + + +def __getattr__(name: str): + # lazy evaluation of environment variables + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return list(environment_variables.keys()) diff --git a/worldcompass/fastvideo/fastvideo_args.py b/worldcompass/fastvideo/fastvideo_args.py new file mode 100644 index 0000000000000000000000000000000000000000..9090e6344359834feb01688fc11cf9ffd4b82ac3 --- /dev/null +++ b/worldcompass/fastvideo/fastvideo_args.py @@ -0,0 +1,1378 @@ +# SPDX-License-Identifier: Apache-2.0 +# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py +"""The arguments of FastVideo Inference.""" + +import argparse +import dataclasses +from contextlib import contextmanager +from dataclasses import field +from enum import Enum +from typing import Any + +from fastvideo.configs.configs import PreprocessConfig +from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode +from fastvideo.configs.utils import clean_cli_args +from fastvideo.logger import init_logger +from fastvideo.platforms import current_platform +from fastvideo.utils import FlexibleArgumentParser, StoreBoolean + +logger = init_logger(__name__) + + +class ExecutionMode(str, Enum): + """Enumeration for different pipeline modes. + + Inherits from str to allow string comparison for backward compatibility. + """ + + INFERENCE = "inference" + PREPROCESS = "preprocess" + FINETUNING = "finetuning" + DISTILLATION = "distillation" + + @classmethod + def from_string(cls, value: str) -> "ExecutionMode": + """Convert string to ExecutionMode enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid mode: {value}. Must be one of: {', '.join([m.value for m in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings for argparse.""" + return [mode.value for mode in cls] + + +class WorkloadType(str, Enum): + """Enumeration for different workload types. + + Inherits from str to allow string comparison for backward compatibility. + """ + + I2V = "i2v" # Image to Video + T2V = "t2v" # Text to Video + T2I = "t2i" # Text to Image + I2I = "i2i" # Image to Image + + @classmethod + def from_string(cls, value: str) -> "WorkloadType": + """Convert string to WorkloadType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid workload type: {value}. Must be one of: {', '.join([m.value for m in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings for argparse.""" + return [workload.value for workload in cls] + + +# args for fastvideo framework +@dataclasses.dataclass +class FastVideoArgs: + # Model and path configuration (for convenience) + model_path: str + + # Running mode + mode: ExecutionMode = ExecutionMode.INFERENCE + + # Workload type + workload_type: WorkloadType = WorkloadType.T2V + + # Cache strategy + cache_strategy: str = "none" + + # Distributed executor backend + distributed_executor_backend: str = "mp" + + # HuggingFace specific parameters + trust_remote_code: bool = False + revision: str | None = None + + # Parallelism + num_gpus: int = 1 + tp_size: int = -1 + sp_size: int = -1 + gpu_para: int = 1 + hsdp_replicate_dim: int = 1 + hsdp_shard_dim: int = -1 + dist_timeout: int | None = None # timeout for torch.distributed + + pipeline_config: PipelineConfig = field(default_factory=PipelineConfig) + preprocess_config: PreprocessConfig | None = None + + # LoRA parameters + # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. + lora_path: str | None = None + lora_nickname: str = "default" # for swapping adapters in the pipeline + # can restrict layers to adapt, e.g. ["q_proj"] + # Will adapt only q, k, v, o by default. + lora_target_modules: list[str] | None = None + + output_type: str = "pil" + + # CPU offload parameters + dit_cpu_offload: bool = True + use_fsdp_inference: bool = True + text_encoder_cpu_offload: bool = True + image_encoder_cpu_offload: bool = True + vae_cpu_offload: bool = True + pin_cpu_memory: bool = True + + # STA (Sliding Tile Attention) parameters + mask_strategy_file_path: str | None = None + STA_mode: STA_Mode = STA_Mode.STA_INFERENCE + skip_time_steps: int = 15 + + # Compilation + enable_torch_compile: bool = False + + disable_autocast: bool = False + + # VSA parameters + VSA_sparsity: float = 0.0 # inference/validation sparsity + + # Master port for distributed training/inference + master_port: int | None = None + + # Stage verification + enable_stage_verification: bool = True + + # Prompt text file for batch processing + prompt_txt: str | None = None + + # model paths for correct deallocation + model_paths: dict[str, str] = field(default_factory=dict) + model_loaded: dict[str, bool] = field( + default_factory=lambda: { + "transformer": True, + "vae": True, + } + ) + + eval_only: bool = False + cls_name: str | None = None + load_from_dir: str | None = None + module_name: str | None = None + ar_action_load_from_dir: str | None = None + + # # DMD parameters + # dmd_denoising_steps: List[int] | None = field(default=None) + + # MoE parameters used by Wan2.2 + boundary_ratio: float | None = None + + @property + def training_mode(self) -> bool: + return not self.inference_mode + + def __post_init__(self): + self.check_fastvideo_args() + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + # Model and path configuration + parser.add_argument( + "--model-path", + type=str, + help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", + ) + parser.add_argument( + "--model-dir", + type=str, + help="Directory containing StepVideo model", + ) + + # Running mode + parser.add_argument( + "--mode", + type=str, + choices=ExecutionMode.choices(), + default=FastVideoArgs.mode.value, + help="The mode to run FastVideo", + ) + + # Workload type + parser.add_argument( + "--workload-type", + type=str, + choices=WorkloadType.choices(), + default=FastVideoArgs.workload_type.value, + help="The workload type", + ) + + # distributed_executor_backend + parser.add_argument( + "--distributed-executor-backend", + type=str, + choices=["mp"], + default=FastVideoArgs.distributed_executor_backend, + help="The distributed executor backend to use", + ) + + # HuggingFace specific parameters + parser.add_argument( + "--trust-remote-code", + action=StoreBoolean, + default=FastVideoArgs.trust_remote_code, + help="Trust remote code when loading HuggingFace models", + ) + parser.add_argument( + "--revision", + type=str, + default=FastVideoArgs.revision, + help="The specific model version to use (can be a branch name, tag name, or commit id)", + ) + + # Parallelism + parser.add_argument( + "--num-gpus", + type=int, + default=FastVideoArgs.num_gpus, + help="The number of GPUs to use.", + ) + parser.add_argument( + "--tp-size", + type=int, + default=FastVideoArgs.tp_size, + help="The tensor parallelism size.", + ) + parser.add_argument( + "--sp-size", + type=int, + default=FastVideoArgs.sp_size, + help="The sequence parallelism size.", + ) + parser.add_argument( + "--gpu-para", + type=int, + default=FastVideoArgs.gpu_para, + help="The GPU parallelism size.", + ) + parser.add_argument( + "--hsdp-replicate-dim", + type=int, + default=FastVideoArgs.hsdp_replicate_dim, + help="The data parallelism size.", + ) + parser.add_argument( + "--hsdp-shard-dim", + type=int, + default=FastVideoArgs.hsdp_shard_dim, + help="The data parallelism shards.", + ) + parser.add_argument( + "--dist-timeout", + type=int, + default=FastVideoArgs.dist_timeout, + help="Set timeout for torch.distributed initialization.", + ) + + # Output type + parser.add_argument( + "--output-type", + type=str, + default=FastVideoArgs.output_type, + choices=["pil"], + help="Output type for the generated video", + ) + + # Prompt text file for batch processing + parser.add_argument( + "--prompt-txt", + type=str, + default=FastVideoArgs.prompt_txt, + help="Path to a text file containing prompts (one per line) for batch processing", + ) + + # STA (Sliding Tile Attention) parameters + parser.add_argument( + "--STA-mode", + type=str, + default=FastVideoArgs.STA_mode.value, + choices=[mode.value for mode in STA_Mode], + help="STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None", + ) + parser.add_argument( + "--skip-time-steps", + type=int, + default=FastVideoArgs.skip_time_steps, + help="Number of time steps to warmup (full attention) for STA", + ) + parser.add_argument( + "--mask-strategy-file-path", + type=str, + help="Path to mask strategy JSON file for STA", + ) + parser.add_argument( + "--enable-torch-compile", + action=StoreBoolean, + default=FastVideoArgs.enable_torch_compile, + help="Use torch.compile to speed up DiT inference." + + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)", + ) + + parser.add_argument( + "--dit-cpu-offload", + action=StoreBoolean, + help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.", + ) + parser.add_argument( + "--use-fsdp-inference", + action=StoreBoolean, + help="Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.", + ) + parser.add_argument( + "--text-encoder-cpu-offload", + action=StoreBoolean, + help="Use CPU offload for text encoder. Enable if run out of memory.", + ) + parser.add_argument( + "--image-encoder-cpu-offload", + action=StoreBoolean, + help="Use CPU offload for image encoder. Enable if run out of memory.", + ) + parser.add_argument( + "--vae-cpu-offload", + action=StoreBoolean, + help="Use CPU offload for VAE. Enable if run out of memory.", + ) + parser.add_argument( + "--pin-cpu-memory", + action=StoreBoolean, + help='Pin memory for CPU offload. Only added as a temp workaround if it throws "CUDA error: invalid argument". ' + "Should be enabled in almost all cases", + ) + parser.add_argument( + "--disable-autocast", + action=StoreBoolean, + help="Disable autocast for denoising loop and vae decoding in pipeline sampling", + ) + + # VSA parameters + parser.add_argument( + "--VSA-sparsity", + type=float, + default=FastVideoArgs.VSA_sparsity, + help="Validation sparsity for VSA", + ) + + # Master port for distributed training/inference + parser.add_argument( + "--master-port", + type=int, + default=FastVideoArgs.master_port, + help="Master port for distributed training/inference", + ) + + # Stage verification + parser.add_argument( + "--enable-stage-verification", + action=StoreBoolean, + default=FastVideoArgs.enable_stage_verification, + help="Enable input/output verification for pipeline stages", + ) + + parser.add_argument( + "--cls-name", + type=str, + help="model package name", + ) + parser.add_argument( + "--load-from-dir", + type=str, + help="ar model checkpoint directory", + ) + parser.add_argument( + "--ar-action-load-from-dir", + type=str, + help="ar action model checkpoint directory", + ) + parser.add_argument( + "--eval-only", + action=StoreBoolean, + help="Only evaluate the model", + ) + # Add pipeline configuration arguments + PipelineConfig.add_cli_args(parser) + + # Add preprocessing configuration arguments + PreprocessConfig.add_cli_args(parser) + + return parser + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs": + provided_args = clean_cli_args(args) + # Get all fields from the dataclass + attrs = [attr.name for attr in dataclasses.fields(cls)] + + # Create a dictionary of attribute values, with defaults for missing attributes + kwargs: dict[str, Any] = {} + for attr in attrs: + if attr == "pipeline_config": + if "HunyuanTransformer" in provided_args["cls_name"]: + continue + pipeline_config = PipelineConfig.from_kwargs(provided_args) + kwargs["pipeline_config"] = pipeline_config + elif attr == "preprocess_config": + preprocess_config = PreprocessConfig.from_kwargs(provided_args) + kwargs["preprocess_config"] = preprocess_config + elif attr == "mode": + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, FastVideoArgs.mode.value) + kwargs["mode"] = ( + ExecutionMode.from_string(mode_value) + if isinstance(mode_value, str) + else mode_value + ) + elif attr == "workload_type": + # Convert string to WorkloadType enum + workload_type_value = getattr( + args, "workload_type", FastVideoArgs.workload_type.value + ) + kwargs["workload_type"] = ( + WorkloadType.from_string(workload_type_value) + if isinstance(workload_type_value, str) + else workload_type_value + ) + # Use getattr with default value from the dataclass for potentially missing attributes + else: + # Get the field to check if it has a default_factory + field = dataclasses.fields(cls)[ + next( + i + for i, f in enumerate(dataclasses.fields(cls)) + if f.name == attr + ) + ] + if field.default_factory is not dataclasses.MISSING: + # Use the default_factory to create the default value + default_value = field.default_factory() + else: + default_value = getattr(cls, attr, None) + value = getattr(args, attr, default_value) + kwargs[attr] = value # type: ignore + + return cls(**kwargs) # type: ignore + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs": + # Convert mode string to enum if necessary + if "mode" in kwargs and isinstance(kwargs["mode"], str): + kwargs["mode"] = ExecutionMode.from_string(kwargs["mode"]) + + # Convert workload_type string to enum if necessary + if "workload_type" in kwargs and isinstance( + kwargs["workload_type"], str + ): + kwargs["workload_type"] = WorkloadType.from_string( + kwargs["workload_type"] + ) + + kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs) + kwargs["preprocess_config"] = PreprocessConfig.from_kwargs(kwargs) + return cls(**kwargs) + + def check_fastvideo_args(self) -> None: + """Validate inference arguments for consistency.""" + if current_platform.is_mps(): + self.use_fsdp_inference = False + + # Validate mode and inference_mode consistency + assert isinstance( + self.mode, ExecutionMode + ), f"Mode must be an ExecutionMode enum, got {type(self.mode)}" + assert ( + self.mode in ExecutionMode.choices() + ), f"Invalid execution mode: {self.mode}" + + # Validate workload type + assert isinstance( + self.workload_type, WorkloadType + ), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}" + assert ( + self.workload_type in WorkloadType.choices() + ), f"Invalid workload type: {self.workload_type}" + + self.inference_mode = False + + if not self.inference_mode: + assert ( + self.hsdp_replicate_dim != -1 + ), "hsdp_replicate_dim must be set for training" + assert ( + self.hsdp_shard_dim != -1 + ), "hsdp_shard_dim must be set for training" + assert self.sp_size != -1, "sp_size must be set for training" + + if self.tp_size == -1: + self.tp_size = 1 + if self.sp_size == -1: + self.sp_size = self.num_gpus + if self.hsdp_shard_dim == -1: + self.hsdp_shard_dim = self.num_gpus + + assert ( + self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0 + ), "num_gpus must >= and be divisible by sp_size" + assert ( + self.hsdp_replicate_dim <= self.num_gpus + and self.num_gpus % self.hsdp_replicate_dim == 0 + ), "num_gpus must >= and be divisible by hsdp_replicate_dim" + assert ( + self.hsdp_shard_dim <= self.num_gpus + and self.num_gpus % self.hsdp_shard_dim == 0 + ), "num_gpus must >= and be divisible by hsdp_shard_dim" + + if self.num_gpus < max(self.tp_size, self.sp_size): + self.num_gpus = max(self.tp_size, self.sp_size) + + if self.pipeline_config is None: + raise ValueError("pipeline_config is not set in FastVideoArgs") + + self.pipeline_config.check_pipeline_config() + + # Add preprocessing config validation if needed + if self.mode == ExecutionMode.PREPROCESS: + if self.preprocess_config is None: + raise ValueError( + "preprocess_config is not set in FastVideoArgs when mode is PREPROCESS" + ) + if self.preprocess_config.model_path == "": + self.preprocess_config.model_path = self.model_path + if not self.pipeline_config.vae_config.load_encoder: + self.pipeline_config.vae_config.load_encoder = True + self.preprocess_config.check_preprocess_config() + + +_current_fastvideo_args = None + + +def prepare_fastvideo_args(argv: list[str]) -> FastVideoArgs: + """Prepare the inference arguments from the command line arguments. + + Args: + argv: The command line arguments. Typically, it should be `sys.argv[1:]` + to ensure compatibility with `parse_args` when no arguments are passed. + + Returns: + The inference arguments. + """ + parser = FlexibleArgumentParser() + FastVideoArgs.add_cli_args(parser) + raw_args = parser.parse_args(argv) + fastvideo_args = FastVideoArgs.from_cli_args(raw_args) + global _current_fastvideo_args + _current_fastvideo_args = fastvideo_args + return fastvideo_args + + +@contextmanager +def set_current_fastvideo_args(fastvideo_args: FastVideoArgs): + """Temporarily set the current fastvideo config. + + Used during model initialization. We save the current fastvideo config in a global variable, so + that all modules can access it, e.g. custom ops can access the fastvideo config to determine how + to dispatch. + """ + global _current_fastvideo_args + old_fastvideo_args = _current_fastvideo_args + try: + _current_fastvideo_args = fastvideo_args + yield + finally: + _current_fastvideo_args = old_fastvideo_args + + +def get_current_fastvideo_args() -> FastVideoArgs: + if _current_fastvideo_args is None: + # in ci, usually when we test custom ops/modules directly, + # we don't set the fastvideo config. In that case, we set a default + # config. + # TODO(will): may need to handle this for CI. + raise ValueError("Current fastvideo args is not set.") + return _current_fastvideo_args + + +@dataclasses.dataclass +class TrainingArgs(FastVideoArgs): + """Training arguments. + + Inherits from FastVideoArgs and adds training-specific arguments. If there are any conflicts, + the training arguments will take precedence. + """ + + data_path: str = "" + dataloader_num_workers: int = 0 + train_batch_size: int = 0 + group_frame: bool = False + group_resolution: bool = False + + # text encoder & vae & diffusion model + pretrained_model_name_or_path: str = "" + dit_model_name_or_path: str = "" + + # diffusion setting + ema_decay: float = 0.0 + ema_start_step: int = 0 + training_cfg_rate: float = 0.0 + precondition_outputs: bool = False + + # GRPO setting + std_type: str = "sample" + adv_clip_max: float = 5.0 + + sampling_steps: int = 10 + sampling_batch_size: int = 2 + grpo_generation_num: int = 2 + bestofn: int = 2 + train_timestep_fraction: float = 0.5 + single_chunk_size: int = 1 + action_reward_type: bool = False + hpsv3_reward_weight: float = 0.2 + hpsv3_quality_reward_weight: float = 0.2 + action_reward_weight: float = 1.0 + hpsv3_quality_drift_reward_weight: float = 0.2 + + # Chunk selection parameters + chunk_selection_strategy: str = ( + "min2max" # Options: "min2max", "max2min", "random" + ) + min_chunk_id: int = 1 + max_chunk_id: int = 16 + + # EMA parameters + ema_min_decay: float = 0.4 + ema_max_decay: float = 0.9 + ema_step_decay: float = 0.002 + ema_ckpt_decay: float = 0.9 + + # validation & logs + validation_dataset_file: str = "" + validation_preprocessed_path: str = "" + validation_sampling_steps: str = "" + validation_guidance_scale: str = "" + validation_steps: float = 0.0 + log_validation: bool = False + tracker_project_name: str = "" + wandb_run_name: str = "" + seed: int | None = None + + # output + output_dir: str = "" + checkpoints_total_limit: int = 0 + checkpointing_steps: int = 0 + resume_from_checkpoint: str = ( + "" # specify the checkpoint folder to resume from + ) + + # optimizer & scheduler + num_train_epochs: int = 0 + max_train_steps: int = 0 + gradient_accumulation_steps: int = 0 + learning_rate: float = 0.0 + scale_lr: bool = False + lr_scheduler: str = "constant" + lr_warmup_steps: int = 0 + max_grad_norm: float = 0.0 + enable_gradient_checkpointing_type: str | None = None + selective_checkpointing: float = 0.0 + mixed_precision: str = "" + train_sp_batch_size: int = 0 + fsdp_sharding_startegy: str = "" + + weighting_scheme: str = "" + logit_mean: float = 0.0 + logit_std: float = 1.0 + mode_scale: float = 0.0 + + num_euler_timesteps: int = 0 + lr_num_cycles: int = 0 + lr_power: float = 0.0 + min_lr_ratio: float = ( + 0.5 # minimum learning rate ratio for cosine_with_min_lr scheduler + ) + not_apply_cfg_solver: bool = False + distill_cfg: float = 0.0 + scheduler_type: str = "" + linear_quadratic_threshold: float = 0.0 + linear_range: float = 0.0 + weight_decay: float = 0.0 + use_ema: bool = False + multi_phased_distill_schedule: str = "" + pred_decay_weight: float = 0.0 + pred_decay_type: str = "" + hunyuan_teacher_disable_cfg: bool = False + + # master_weight_type + master_weight_type: str = "" + + # VSA training decay parameters + VSA_decay_rate: float = 0.01 # decay rate -> 0.02 + VSA_decay_interval_steps: int = 1 # decay interval steps -> 50 + + # LoRA training parameters + lora_rank: int | None = None + lora_alpha: int | None = None + lora_training: bool = False + + # camera dataset + json_path: str = "" + eval_json_path: str = "" + causal: bool = False + window_frames: int = 9 + wandb_key: str = "f8f8f614d325bd04540f8517b142ca942b21017a" + wandb_entity: str = "zhyzhy" + action: bool = False + i2v_rate: float = 0.0 + use_mem: bool = False + memory_nframes: int = 4 + + # Path arguments (to avoid hardcoded absolute paths) + vae_path: str = "" # VAE model path + random_pose_path: str = "" # Random pose json file path + neg_prompt_path: str = "" # Negative prompt pt file path + neg_byt5_prompt_path: str = "" # Negative byt5 prompt pt file path + generated_videos_dir: str = ( + "" # Directory to save generated videos during training + ) + + # Reward model arguments + camera_estimator: str = ( + "dav3" # Camera pose estimator: "dav3" or "worldmirror" + ) + cache_dir: str = "" # Cache directory for downloading models + + # distillation args + generator_update_interval: int = 5 + min_timestep_ratio: float = 0.2 + max_timestep_ratio: float = 0.98 + real_score_guidance_scale: float = 3.5 + fake_score_learning_rate: float = ( + 0.0 # separate learning rate for fake_score_transformer, if 0.0, use learning_rate + ) + fake_score_lr_scheduler: str = ( + "constant" # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler + ) + training_state_checkpointing_steps: int = 0 # for resuming training + weight_only_checkpointing_steps: int = 0 # for inference + log_visualization: bool = False + # simulate generator forward to match inference + simulate_generator_forward: bool = False + + @classmethod + def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": + provided_args = clean_cli_args(args) + # Get all fields from the dataclass + attrs = [attr.name for attr in dataclasses.fields(cls)] + logger.info(provided_args) + # Create a dictionary of attribute values, with defaults for missing attributes + kwargs: dict[str, Any] = {} + for attr in attrs: + if attr == "pipeline_config": + if "HunyuanTransformer" in provided_args["cls_name"]: + continue + pipeline_config = PipelineConfig.from_kwargs(provided_args) + kwargs[attr] = pipeline_config + elif attr == "mode": + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value) + kwargs[attr] = ( + ExecutionMode.from_string(mode_value) + if isinstance(mode_value, str) + else mode_value + ) + elif attr == "workload_type": + # Convert string to WorkloadType enum + workload_type_value = getattr( + args, "workload_type", WorkloadType.T2V.value + ) + kwargs[attr] = ( + WorkloadType.from_string(workload_type_value) + if isinstance(workload_type_value, str) + else workload_type_value + ) + # Use getattr with default value from the dataclass for potentially missing attributes + else: + # Get the field to check its default value + field = dataclasses.fields(cls)[ + next( + i + for i, f in enumerate(dataclasses.fields(cls)) + if f.name == attr + ) + ] + + # Check if the attribute is provided in args + if hasattr(args, attr): + value = getattr(args, attr) + else: + # Use the field's default value + if field.default_factory is not dataclasses.MISSING: + value = field.default_factory() + elif field.default is not dataclasses.MISSING: + value = field.default + else: + # No default value, use None + value = None + + kwargs[attr] = value + + return cls(**kwargs) # type: ignore + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument( + "--data-path", type=str, required=True, help="Path to parquet files" + ) + parser.add_argument( + "--dataloader-num-workers", + type=int, + required=True, + help="Number of workers for dataloader", + ) + + # GRPO setting + parser.add_argument( + "--std-type", + type=str, + default="sample", + choices=["sample", "global", "sample_max"], + help="Standard deviation type", + ) + parser.add_argument( + "--adv-clip-max", type=float, default=5.0, help="Advantage clip max" + ) + parser.add_argument( + "--sampling-steps", + type=int, + required=True, + help="Number of sampling steps", + ) + parser.add_argument( + "--sampling-batch-size", + type=int, + required=True, + help="Number of sampling batch size", + ) + parser.add_argument( + "--grpo-generation-num", + type=int, + required=True, + help="GRPO generation number", + ) + parser.add_argument( + "--bestofn", type=int, required=True, help="GRPO best of n" + ) + parser.add_argument( + "--train-timestep-fraction", + type=float, + required=True, + help="GRPO train timestep fraction", + ) + parser.add_argument( + "--single-chunk-size", + type=int, + default=4, + help="GRPO single chunk size", + ) + parser.add_argument( + "--action-reward-type", + type=str, + default="action", + help="Whether to use fine action reward", + ) + parser.add_argument( + "--hpsv3-reward-weight", + type=float, + default=0.2, + help="Hpsv3 reward weight", + ) + parser.add_argument( + "--hpsv3-quality-reward-weight", + type=float, + default=0.2, + help="Hpsv3 quality reward weight", + ) + parser.add_argument( + "--action-reward-weight", + type=float, + default=1.0, + help="Action reward weight", + ) + parser.add_argument( + "--hpsv3-quality-drift-reward-weight", + type=float, + default=0.2, + help="Hpsv3 quality drift reward weight", + ) + + # Chunk selection parameters + parser.add_argument( + "--chunk-selection-strategy", + type=str, + default=TrainingArgs.chunk_selection_strategy, + choices=["min2max", "max2min"], + help="Chunk selection strategy: min2max, max2min, or random", + ) + parser.add_argument( + "--min-chunk-id", + type=int, + default=TrainingArgs.min_chunk_id, + help="Minimum chunk ID for chunk selection", + ) + parser.add_argument( + "--max-chunk-id", + type=int, + default=TrainingArgs.max_chunk_id, + help="Maximum chunk ID for chunk selection", + ) + + # EMA parameters + parser.add_argument( + "--ema-min-decay", + type=float, + default=TrainingArgs.ema_min_decay, + help="Minimum decay rate for EMA", + ) + parser.add_argument( + "--ema-max-decay", + type=float, + default=TrainingArgs.ema_max_decay, + help="Maximum decay rate for EMA", + ) + parser.add_argument( + "--ema-step-decay", + type=float, + default=TrainingArgs.ema_step_decay, + help="Step decay rate for EMA", + ) + parser.add_argument( + "--ema-ckpt-decay", + type=float, + default=TrainingArgs.ema_ckpt_decay, + help="Checkpoint decay rate for EMA", + ) + + # Training batch and model configuration + parser.add_argument( + "--train-batch-size", + type=int, + required=True, + help="Training batch size", + ) + parser.add_argument( + "--group-frame", + action=StoreBoolean, + help="Whether to group frames during training", + ) + parser.add_argument( + "--group-resolution", + action=StoreBoolean, + help="Whether to group resolutions during training", + ) + + # Model paths + parser.add_argument( + "--pretrained-model-name-or-path", + type=str, + required=True, + help="Path to pretrained model or model name", + ) + parser.add_argument( + "--dit-model-name-or-path", + type=str, + required=False, + help="Path to DiT model or model name", + ) + parser.add_argument( + "--cache-dir", type=str, help="Directory to cache models" + ) + + # Diffusion settings + parser.add_argument( + "--ema-decay", type=float, default=0.999, help="EMA decay rate" + ) + parser.add_argument( + "--ema-start-step", type=int, default=0, help="Step to start EMA" + ) + parser.add_argument( + "--training-cfg-rate", + type=float, + help="Classifier-free guidance scale", + ) + parser.add_argument( + "--precondition-outputs", + action=StoreBoolean, + help="Whether to precondition the outputs of the model", + ) + + # Validation and logging + parser.add_argument( + "--validation-dataset-file", + type=str, + help="Path to unprocessed validation dataset", + ) + parser.add_argument( + "--validation-preprocessed-path", + type=str, + help="Path to processed validation dataset", + ) + parser.add_argument( + "--validation-sampling-steps", + type=str, + help="Validation sampling steps", + ) + parser.add_argument( + "--validation-guidance-scale", + type=str, + help="Validation guidance scale", + ) + parser.add_argument( + "--validation-steps", type=float, help="Number of validation steps" + ) + parser.add_argument( + "--log-validation", + action=StoreBoolean, + help="Whether to log validation results", + ) + parser.add_argument( + "--tracker-project-name", type=str, help="Project name for tracking" + ) + parser.add_argument( + "--wandb-run-name", type=str, help="Run name for wandb" + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Seed for deterministic training", + ) + + # Output configuration + parser.add_argument( + "--output-dir", + type=str, + required=True, + help="Output directory for checkpoints and logs", + ) + parser.add_argument( + "--checkpoints-total-limit", + type=int, + help="Maximum number of checkpoints to keep", + ) + parser.add_argument( + "--checkpointing-steps", type=int, help="Steps between checkpoints" + ) + parser.add_argument( + "--training-state-checkpointing-steps", + type=int, + help="Steps between training state checkpoints (for resuming training)", + ) + parser.add_argument( + "--weight-only-checkpointing-steps", + type=int, + help="Steps between weight-only checkpoints (for inference)", + ) + parser.add_argument( + "--resume-from-checkpoint", + type=str, + help="Path to checkpoint to resume from", + ) + parser.add_argument( + "--logging-dir", type=str, help="Directory for logging" + ) + + # Training configuration + parser.add_argument( + "--num-train-epochs", type=int, help="Number of training epochs" + ) + parser.add_argument( + "--max-train-steps", + type=int, + help="Maximum number of training steps", + ) + parser.add_argument( + "--gradient-accumulation-steps", + type=int, + help="Number of steps to accumulate gradients", + ) + parser.add_argument( + "--learning-rate", type=float, required=True, help="Learning rate" + ) + parser.add_argument( + "--scale-lr", + action=StoreBoolean, + help="Whether to scale learning rate", + ) + parser.add_argument( + "--lr-scheduler", + type=str, + default="constant", + help="Learning rate scheduler type", + ) + parser.add_argument( + "--lr-warmup-steps", + type=int, + default=10, + help="Number of warmup steps for learning rate", + ) + parser.add_argument( + "--max-grad-norm", type=float, help="Maximum gradient norm" + ) + parser.add_argument( + "--enable-gradient-checkpointing-type", + type=str, + choices=["full", "ops", "block_skip"], + default=None, + help="Gradient checkpointing type", + ) + parser.add_argument( + "--selective-checkpointing", + type=float, + help="Selective checkpointing threshold", + ) + parser.add_argument( + "--mixed-precision", type=str, help="Mixed precision training type" + ) + parser.add_argument( + "--train-sp-batch-size", + type=int, + help="Training spatial parallelism batch size", + ) + + parser.add_argument( + "--fsdp-sharding-strategy", type=str, help="FSDP sharding strategy" + ) + + parser.add_argument( + "--weighting_scheme", + type=str, + default="uniform", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"], + ) + parser.add_argument( + "--logit_mean", + type=float, + default=0.0, + help="mean to use when using the `'logit_normal'` weighting scheme.", + ) + parser.add_argument( + "--logit_std", + type=float, + default=1.0, + help="std to use when using the `'logit_normal'` weighting scheme.", + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + + # Additional training parameters + parser.add_argument( + "--num-euler-timesteps", type=int, help="Number of Euler timesteps" + ) + parser.add_argument( + "--lr-num-cycles", type=int, help="Number of learning rate cycles" + ) + parser.add_argument( + "--lr-power", type=float, help="Learning rate power" + ) + parser.add_argument( + "--min-lr-ratio", + type=float, + default=TrainingArgs.min_lr_ratio, + help="Minimum learning rate ratio for cosine_with_min_lr scheduler", + ) + parser.add_argument( + "--not-apply-cfg-solver", + action=StoreBoolean, + help="Whether to not apply CFG solver", + ) + parser.add_argument( + "--distill-cfg", type=float, help="Distillation CFG scale" + ) + parser.add_argument("--scheduler-type", type=str, help="Scheduler type") + parser.add_argument( + "--linear-quadratic-threshold", + type=float, + help="Linear quadratic threshold", + ) + parser.add_argument("--linear-range", type=float, help="Linear range") + parser.add_argument("--weight-decay", type=float, help="Weight decay") + parser.add_argument( + "--use-ema", action=StoreBoolean, help="Whether to use EMA" + ) + parser.add_argument( + "--multi-phased-distill-schedule", + type=str, + help="Multi-phased distillation schedule", + ) + parser.add_argument( + "--pred-decay-weight", type=float, help="Prediction decay weight" + ) + parser.add_argument( + "--pred-decay-type", type=str, help="Prediction decay type" + ) + parser.add_argument( + "--hunyuan-teacher-disable-cfg", + action=StoreBoolean, + help="Whether to disable CFG for Hunyuan teacher", + ) + parser.add_argument( + "--master-weight-type", type=str, help="Master weight type" + ) + + # VSA parameters for training with dense to sparse adaption + parser.add_argument( + "--VSA-decay-rate", # decay rate, how much sparsity you want to decay each step + type=float, + default=TrainingArgs.VSA_decay_rate, + help="VSA decay rate", + ) + parser.add_argument( + "--VSA-decay-interval-steps", # how many steps for training with current sparsity + type=int, + default=TrainingArgs.VSA_decay_interval_steps, + help="VSA decay interval steps", + ) + parser.add_argument( + "--lora-training", + action=StoreBoolean, + help="Whether to use LoRA training", + ) + parser.add_argument("--lora-rank", type=int, help="LoRA rank") + parser.add_argument("--lora-alpha", type=int, help="LoRA alpha") + + # Distillation arguments + parser.add_argument( + "--generator-update-interval", + type=int, + default=TrainingArgs.generator_update_interval, + help="Ratio of student updates to critic updates.", + ) + parser.add_argument( + "--min-timestep-ratio", + type=float, + default=TrainingArgs.min_timestep_ratio, + help="Minimum step ratio", + ) + parser.add_argument( + "--max-timestep-ratio", + type=float, + default=TrainingArgs.max_timestep_ratio, + help="Maximum step ratio", + ) + parser.add_argument( + "--real-score-guidance-scale", + type=float, + default=TrainingArgs.real_score_guidance_scale, + help="Teacher guidance scale", + ) + parser.add_argument( + "--fake-score-learning-rate", + type=float, + default=TrainingArgs.fake_score_learning_rate, + help="Learning rate for fake score transformer", + ) + parser.add_argument( + "--fake-score-lr-scheduler", + type=str, + default=TrainingArgs.fake_score_lr_scheduler, + help="Learning rate scheduler for fake score transformer", + ) + parser.add_argument( + "--log-visualization", + action=StoreBoolean, + help="Whether to log visualization", + ) + parser.add_argument( + "--simulate-generator-forward", + action=StoreBoolean, + help="Whether to simulate generator forward to match inference", + ) + + # camera + parser.add_argument( + "--json-path", type=str, default="", help="camera json path" + ) + parser.add_argument( + "--eval-json-path", + type=str, + default="", + help="eval camera json path", + ) + parser.add_argument( + "--causal", action=StoreBoolean, help="training ar model" + ) + parser.add_argument( + "--use-mem", action=StoreBoolean, help="use memory for training" + ) + parser.add_argument( + "--window-frames", type=int, default=9, help="window size" + ) + parser.add_argument("--wandb-key", type=str, default="", help="key") + parser.add_argument( + "--wandb-entity", type=str, default="", help="entity" + ) + parser.add_argument( + "--action", action=StoreBoolean, help="training ar model" + ) + parser.add_argument( + "--i2v-rate", + type=float, + default=0.0, + help="training rating for i2v", + ) + # Path arguments + parser.add_argument( + "--vae-path", type=str, default="", help="VAE model path" + ) + parser.add_argument( + "--random-pose-path", + type=str, + default="", + help="Random pose json file path", + ) + parser.add_argument( + "--neg-prompt-path", + type=str, + default="", + help="Negative prompt pt file path", + ) + parser.add_argument( + "--neg-byt5-prompt-path", + type=str, + default="", + help="Negative byt5 prompt pt file path", + ) + parser.add_argument( + "--generated-videos-dir", + type=str, + default="", + help="Directory to save generated videos during training", + ) + parser.add_argument( + "--camera-estimator", + type=str, + default="dav3", + choices=["dav3", "worldmirror"], + help="Camera pose estimator model: dav3 or worldmirror", + ) + + return parser + + +def parse_int_list(value: str) -> list[int]: + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] diff --git a/worldcompass/fastvideo/forward_context.py b/worldcompass/fastvideo/forward_context.py new file mode 100644 index 0000000000000000000000000000000000000000..4f8a40c4ba6d3684121d4e5fa0bb41a0199db603 --- /dev/null +++ b/worldcompass/fastvideo/forward_context.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/forward_context.py + +import time +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from fastvideo.logger import init_logger + +if TYPE_CHECKING: + from fastvideo.attention import AttentionMetadata + from fastvideo.pipelines import ForwardBatch + +logger = init_logger(__name__) + +# TODO(will): check if this is needed +# track_batchsize: bool = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL >= 0 +track_batchsize: bool = False +last_logging_time: float = 0 +forward_start_time: float = 0 +# batchsize_logging_interval: float = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL +batchsize_logging_interval: float = 1000 +batchsize_forward_time: defaultdict = defaultdict(list) + + +# +@dataclass +class ForwardContext: + current_timestep: int + # TODO(will): check this arg + # copy from vllm_config.compilation_config.static_forward_context + # attn_layers: Dict[str, Any] + # TODO: extend to support per-layer dynamic forward context + attn_metadata: "AttentionMetadata" # set dynamically for each forward pass + forward_batch: Optional["ForwardBatch"] = None + + +_forward_context: Optional["ForwardContext"] = None + + +def get_forward_context() -> "ForwardContext": + """Get the current forward context.""" + assert _forward_context is not None, ( + "Forward context is not set. " + "Please use `set_forward_context` to set the forward context." + ) + return _forward_context + + +# TODO(will): finalize the interface +@contextmanager +def set_forward_context( + current_timestep, + attn_metadata, + forward_batch: Optional["ForwardBatch"] = None, +): + """A context manager that stores the current forward context, can be attention metadata, etc. + Here we can inject common logic for every model forward pass.""" + global forward_start_time + need_to_track_batchsize = track_batchsize and attn_metadata is not None + if need_to_track_batchsize: + forward_start_time = time.perf_counter() + global _forward_context + prev_context = _forward_context + _forward_context = ForwardContext( + current_timestep=current_timestep, + attn_metadata=attn_metadata, + forward_batch=forward_batch, + ) + + try: + yield + finally: + global last_logging_time, batchsize_logging_interval + if need_to_track_batchsize: + if hasattr(attn_metadata, "num_prefill_tokens"): + # for v0 attention backends + batchsize = ( + attn_metadata.num_prefill_tokens + + attn_metadata.num_decode_tokens + ) + else: + # for v1 attention backends + batchsize = attn_metadata.num_input_tokens + now = time.perf_counter() + # time measurement is in milliseconds + batchsize_forward_time[batchsize].append( + (now - forward_start_time) * 1000 + ) + if now - last_logging_time > batchsize_logging_interval: + last_logging_time = now + forward_stats = [] + for bs, times in batchsize_forward_time.items(): + if len(times) <= 1: + # can be cudagraph / profiling run + continue + medium = torch.quantile(torch.tensor(times), q=0.5).item() + medium = round(medium, 2) + forward_stats.append((bs, len(times), medium)) + forward_stats.sort(key=lambda x: x[1], reverse=True) + if forward_stats: + logger.info( + ( + "Batchsize forward time stats " + "(batchsize, count, median_time(ms)): %s" + ), + forward_stats, + ) + _forward_context = prev_context diff --git a/worldcompass/fastvideo/layers/__init__.py b/worldcompass/fastvideo/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/layers/activation.py b/worldcompass/fastvideo/layers/activation.py new file mode 100644 index 0000000000000000000000000000000000000000..4baf20fb4203ac521289ceb00137a5e227e8d5e4 --- /dev/null +++ b/worldcompass/fastvideo/layers/activation.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py +"""Custom activation functions.""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# TODO (will): remove this dependency +from fastvideo.layers.custom_op import CustomOp + + +@CustomOp.register("silu_and_mul") +class SiluAndMul(CustomOp): + """An activation function for SwiGLU. + + The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + def __init__(self) -> None: + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + +@CustomOp.register("gelu_and_mul") +class GeluAndMul(CustomOp): + """An activation function for GeGLU. + + The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d) + return: (batch_size, seq_len, d) or (num_tokens, d) + """ + + def __init__(self, approximate: str = "none"): + super().__init__() + self.approximate = approximate + if approximate not in ("none", "tanh"): + raise ValueError(f"Unknown approximate mode: {approximate}") + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:] + + def extra_repr(self) -> str: + return f"approximate={repr(self.approximate)}" + + +@CustomOp.register("gelu_new") +class NewGELU(CustomOp): + def __init__(self): + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + c = math.sqrt(2.0 / math.pi) + return ( + 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0)))) + ) + + +@CustomOp.register("quick_gelu") +class QuickGELU(CustomOp): + # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90 + def __init__(self): + super().__init__() + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + return x * torch.sigmoid(1.702 * x) + + +_ACTIVATION_REGISTRY = { + "gelu": nn.GELU, + "gelu_new": NewGELU, + "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"), + "relu": nn.ReLU, + "silu": nn.SiLU, + "quick_gelu": QuickGELU, +} + + +def get_act_fn(act_fn_name: str) -> nn.Module: + """Get an activation function by name.""" + act_fn_name = act_fn_name.lower() + if act_fn_name not in _ACTIVATION_REGISTRY: + raise ValueError( + f"Activation function {act_fn_name!r} is not supported." + ) + + return _ACTIVATION_REGISTRY[act_fn_name]() + + +_ACTIVATION_AND_MUL_REGISTRY = { + "gelu": GeluAndMul, + "silu": SiluAndMul, +} + + +def get_act_and_mul_fn(act_fn_name: str) -> nn.Module: + """Get an activation-and-mul (i.e. SiluAndMul) function by name.""" + act_fn_name = act_fn_name.lower() + if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY: + raise ValueError( + f"Activation function {act_fn_name!r} is not supported." + ) + + return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]() diff --git a/worldcompass/fastvideo/layers/custom_op.py b/worldcompass/fastvideo/layers/custom_op.py new file mode 100644 index 0000000000000000000000000000000000000000..3e18fdd728c85b8fc6b2ff87b5d102245b730075 --- /dev/null +++ b/worldcompass/fastvideo/layers/custom_op.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py + +from collections.abc import Callable +from typing import Any + +import torch.nn as nn + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class CustomOp(nn.Module): + """Base class for custom ops. + + Dispatches the forward method to the appropriate backend. + """ + + def __init__(self) -> None: + super().__init__() + self._forward_method = self.dispatch_forward() + + def forward(self, *args, **kwargs) -> Any: + return self._forward_method(*args, **kwargs) + + def forward_native(self, *args, **kwargs) -> Any: + """PyTorch-native implementation of the forward method. + + This method is optional. If implemented, it can be used with compilers such as torch.compile + or PyTorch XLA. Also, it can be used for testing purposes. + """ + raise NotImplementedError + + def forward_cuda(self, *args, **kwargs) -> Any: + raise NotImplementedError + + def forward_cpu(self, *args, **kwargs) -> Any: + # By default, we assume that CPU ops are compatible with CUDA ops. + return self.forward_cuda(*args, **kwargs) + + def forward_tpu(self, *args, **kwargs) -> Any: + # By default, we assume that TPU ops are compatible with the + # PyTorch-native implementation. + # NOTE(woosuk): This is a placeholder for future extensions. + return self.forward_native(*args, **kwargs) + + def forward_oot(self, *args, **kwargs) -> Any: + # By default, we assume that OOT ops are compatible with the + # PyTorch-native implementation. + return self.forward_native(*args, **kwargs) + + def dispatch_forward(self) -> Callable: + # FIXME(will): for now, we always use the native implementation, since + # forward_cuda is using vllm's custom ops and it doesn't support + # backwards. We should add our own custom ops that support backwards. + return self.forward_native + # NOTE(woosuk): Here we assume that vLLM was built for only one + # specific backend. Currently, we do not support dynamic dispatching. + enabled = self.enabled() + + if not enabled: + return self.forward_native + + return self.forward_cuda + + @classmethod + def enabled(cls) -> bool: + # since we are not using Inductor, we always return True + return True + + @staticmethod + def default_on() -> bool: + """On by default if level < CompilationLevel.PIECEWISE Specifying 'all' or 'none' in + custom_op takes precedence.""" + raise NotImplementedError + + # Dictionary of all custom ops (classes, indexed by registered name). + # To check if an op with a name is enabled, call .enabled() on the class. + # Examples: + # - MyOp.enabled() + # - op_registry["my_op"].enabled() + op_registry: dict[str, type["CustomOp"]] = {} + + # Decorator to register custom ops. + @classmethod + def register(cls, name: str) -> Callable: + + def decorator(op_cls): + assert name not in cls.op_registry, f"Duplicate op name: {name}" + op_cls.name = name + cls.op_registry[name] = op_cls + return op_cls + + return decorator diff --git a/worldcompass/fastvideo/layers/layernorm.py b/worldcompass/fastvideo/layers/layernorm.py new file mode 100644 index 0000000000000000000000000000000000000000..efb9eafe8ca2307139d0ec5404baf22bcc179708 --- /dev/null +++ b/worldcompass/fastvideo/layers/layernorm.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py +"""Custom normalization layers.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.layers.custom_op import CustomOp +from fastvideo.platforms import current_platform + + +@CustomOp.register("rms_norm") +class RMSNorm(CustomOp): + """Root mean square normalization. + + Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to + https://arxiv.org/abs/1910.07467 + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + var_hidden_size: int | None = None, + has_weight: bool = True, + ) -> None: + super().__init__() + + self.hidden_size = hidden_size + self.variance_epsilon = eps + self.variance_size_override = ( + None if var_hidden_size == hidden_size else var_hidden_size + ) + self.has_weight = has_weight + + self.weight = ( + torch.ones(hidden_size) + if current_platform.is_cuda_alike() + else torch.ones(hidden_size, dtype=dtype) + ) + if self.has_weight: + self.weight = nn.Parameter(self.weight) + + # if we do fully_shard(model.layer_norm), and we call layer_form.forward_native(input) instead of layer_norm(input), + # we need to call model.layer_norm.register_fsdp_forward_method(model, "forward_native") to make sure fsdp2 hooks are triggered + # for mixed precision and cpu offloading + + # the even better way might be fully_shard(model.layer_norm, mp_policy=, cpu_offloading=), and call model.layer_norm(input). everything should work out of the box + # because fsdp2 hooks will be triggered with model.layer_norm.__call__ + def forward_native( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """PyTorch-native implementation equivalent to forward().""" + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + hidden_size = x.shape[-1] + if hidden_size != self.hidden_size: + raise ValueError( + "Expected hidden_size to be " + f"{self.hidden_size}, but found: {hidden_size}" + ) + + if self.variance_size_override is None: + x_var = x + else: + if hidden_size < self.variance_size_override: + raise ValueError( + "Expected hidden_size to be at least " + f"{self.variance_size_override}, but found: {hidden_size}" + ) + + x_var = x[:, :, : self.variance_size_override] + + variance = x_var.pow(2).mean(dim=-1, keepdim=True) + + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = x.to(orig_dtype) + if self.has_weight: + x = x * self.weight + if residual is None: + return x + else: + return x, residual + + def extra_repr(self) -> str: + s = f"hidden_size={self.weight.data.size(0)}" + s += f", eps={self.variance_epsilon}" + return s + + +class ScaleResidual(nn.Module): + """Applies gated residual connection.""" + + def __init__(self, prefix: str = ""): + super().__init__() + + def forward( + self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + """Apply gated residual connection.""" + return residual + x * gate + + +# adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py +# NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using +# FSDP's MixedPrecisionPolicy +class FP32LayerNorm(nn.LayerNorm): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + return F.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +class ScaleResidualLayerNormScaleShift(nn.Module): + """ + Fused operation that combines: + 1. Gated residual connection + 2. LayerNorm + 3. Scale and shift operations + + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + if norm_type == "rms": + self.norm = RMSNorm( + hidden_size, has_weight=elementwise_affine, eps=eps, dtype=dtype + ) + elif norm_type == "layer": + if compute_dtype == torch.float32: + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps + ) + else: + self.norm = nn.LayerNorm( + hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype, + ) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward( + self, + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply gated residual connection, followed by layernorm and scale/shift in a single fused + operation. + + Returns: + Tuple containing: + - normalized and modulated output + - residual value (value after residual connection + but before normalization) + """ + # Apply residual connection with gating + residual_output = residual + x * gate + # Apply normalization + normalized = self.norm(residual_output) + # Apply scale and shift + modulated = normalized * (1.0 + scale) + shift + return modulated, residual_output + + +class LayerNormScaleShift(nn.Module): + """Fused operation that combines LayerNorm with scale and shift operations. + + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.compute_dtype = compute_dtype + if norm_type == "rms": + self.norm = RMSNorm( + hidden_size, has_weight=elementwise_affine, eps=eps + ) + elif norm_type == "layer": + if self.compute_dtype == torch.float32: + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps + ) + else: + self.norm = nn.LayerNorm( + hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype, + ) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward( + self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor: + """Apply ln followed by scale and shift in a single fused operation.""" + normalized = self.norm(x) + if self.compute_dtype == torch.float32: + return (normalized.float() * (1.0 + scale) + shift).to(x.dtype) + else: + return normalized * (1.0 + scale) + shift diff --git a/worldcompass/fastvideo/layers/linear.py b/worldcompass/fastvideo/layers/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..c4cc10e3691a98656d1444776499814fb4f2a9b4 --- /dev/null +++ b/worldcompass/fastvideo/layers/linear.py @@ -0,0 +1,1105 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py + +from abc import abstractmethod + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from fastvideo.distributed import ( + divide, + get_tp_rank, + get_tp_world_size, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from fastvideo.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from fastvideo.logger import init_logger + +# yapf: disable +from fastvideo.models.parameter import (BasevLLMParameter, + BlockQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, + PerTensorScaleParameter, + RowvLLMParameter) +# yapf: enable +from fastvideo.models.utils import set_weight_attrs + +logger = init_logger(__name__) + +WEIGHT_LOADER_V2_SUPPORTED = [ + "CompressedTensorsLinearMethod", + "AWQMarlinLinearMethod", + "AWQLinearMethod", + "GPTQMarlinLinearMethod", + "Fp8LinearMethod", + "MarlinLinearMethod", + "QQQLinearMethod", + "GPTQMarlin24LinearMethod", + "TPUInt8LinearMethod", + "GPTQLinearMethod", + "FBGEMMFp8LinearMethod", + "ModelOptFp8LinearMethod", + "IPEXAWQLinearMethod", + "IPEXGPTQLinearMethod", + "HQQMarlinMethod", + "QuarkLinearMethod", +] + + +def adjust_scalar_to_fused_array( + param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str | int +) -> tuple[torch.Tensor, torch.Tensor]: + """For fused modules (QKV and MLP) we have an array of length N that holds 1 scale for each + "logical" matrix. + + So the param is an array of length N. The loaded_weight corresponds to one of the shards on + disk. Here, we slice the param based on the shard_id for loading. + """ + qkv_idxs = {"q": 0, "k": 1, "v": 2} + + if isinstance(shard_id, str): + shard_id = qkv_idxs[shard_id] + elif not isinstance(shard_id, int): + raise ValueError(f"Unknown Shard Id {shard_id}") + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + return param[shard_id], loaded_weight + + +class LinearMethodBase(QuantizeMethodBase): + """Base class for different (maybe quantized) linear methods.""" + + @abstractmethod + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + """Create weights for a linear layer. The weights will be set as attributes of the layer. + + Args: + layer: The layer that is using the LinearMethodBase factory. + input_size_per_partition: Size of the weight input dim on rank X. + output_partition_sizes: Sizes of the output dim of each logical + weight on rank X. E.g., output_partition_sizes for QKVLinear + is a list contains the width of Wq, Wk, Wv on rank X. + input_size: Size of the input dim of the weight across all ranks. + output_size: Size of the output dim of the weight across all ranks. + params_dtype: Datatype of the parameters. + """ + raise NotImplementedError + + @abstractmethod + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + + Expects create_weights to have been called before on the layer. + """ + raise NotImplementedError + + +class UnquantizedLinearMethod(LinearMethodBase): + """Linear method without quantization.""" + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output = ( + F.linear(x, layer.weight, bias) + if torch.cuda.is_available() or bias is None + else F.linear(x, layer.weight, bias.to(x.dtype)) + ) # NOTE: this line assumes that we are using amp when using cuda and is needed to account for the fact that amp isn't supported in mps + return output + + +class LinearBase(torch.nn.Module): + """Base linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__( + self, + input_size: int, + output_size: int, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + + # Keep input parameters + self.input_size = input_size + self.output_size = output_size + self.skip_bias_add = skip_bias_add + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + self.quant_config = quant_config + self.prefix = prefix + if quant_config is None: + self.quant_method: QuantizeMethodBase | None = ( + UnquantizedLinearMethod() + ) + else: + self.quant_method = quant_config.get_quant_method( + self, prefix=prefix + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: + raise NotImplementedError + + +class ReplicatedLinear(LinearBase): + """Replicated linear layer. + + Args: + input_size: input dimension of the linear layer. + output_size: output dimension of the linear layer. + bias: If true, add bias. + skip_bias_add: If true, skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix=prefix, + ) + + # All the linear layer supports quant method. + assert self.quant_method is not None + self.quant_method.create_weights( + self, + self.input_size, + [self.output_size], + self.input_size, + self.output_size, + self.params_dtype, + weight_loader=self.weight_loader, + ) + + if bias: + self.bias = Parameter( + torch.empty( + self.output_size, + dtype=self.params_dtype, + ) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader( + self, param: Parameter, loaded_weight: torch.Tensor + ) -> None: + # If the weight on disk does not have a shape, give it one + # (such scales for AutoFp8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param.size() == loaded_weight.size(), ( + f"Tried to load weights of size {loaded_weight.size()}" + f"to a parameter of size {param.size()}" + ) + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: + bias = self.bias if not self.skip_bias_add else None + assert self.quant_method is not None + output = self.quant_method.apply(self, x, bias) + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + return s + + +class ColumnParallelLinear(LinearBase): + """Linear layer with column parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its second dimension as A = [A_1, ..., A_p]. + + Args: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make Y available + to all GPUs, otherwise, every GPU will have its output + which is Y_i = XA_i + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + output_sizes: list of output sizes packed into one output, like for QKV + the list would be size 3. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + output_sizes: list[int] | None = None, + prefix: str = "", + ): + # Divide the weight matrix along the last dimension. + self.tp_size = get_tp_world_size() + self.input_size_per_partition = input_size + self.output_size_per_partition = divide(output_size, self.tp_size) + self.output_partition_sizes = [self.output_size_per_partition] + # If QKV or MergedColumn, use output size of each partition. + if hasattr(self, "output_sizes"): + self.output_partition_sizes = [ + divide(output_size, self.tp_size) + for output_size in self.output_sizes + ] + + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix, + ) + + self.gather_output = gather_output + + if output_sizes is None: + output_sizes = [output_size] + + assert self.quant_method is not None + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ + in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ), + ) + if bias: + self.bias = Parameter( + torch.empty( + self.output_size_per_partition, + dtype=params_dtype, + ) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader( + self, param: Parameter, loaded_weight: torch.Tensor + ) -> None: + tp_rank = get_tp_rank() + output_dim = getattr(param, "output_dim", None) + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + is_sharded_weight = is_sharded_weight + + param_data = param.data + if output_dim is not None and not is_sharded_weight: + shard_size = param_data.shape[output_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2( + self, param: Parameter, loaded_weight: torch.Tensor + ) -> None: + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + param.load_column_parallel_weight(loaded_weight=loaded_weight) + + def forward( + self, input_: torch.Tensor + ) -> tuple[torch.Tensor, Parameter | None]: + bias = self.bias if not self.skip_bias_add else None + + # Matrix multiply. + assert self.quant_method is not None + output_parallel = self.quant_method.apply(self, input_, bias) + if self.gather_output: + # All-gather across the partitions. + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + def extra_repr(self) -> str: + s = f"in_features={self.input_size}" + s += f", output_features={self.output_size_per_partition}" + s += f", bias={self.bias is not None}" + s += f", tp_size={get_tp_world_size()}" + s += f", gather_output={self.gather_output}" + return s + + +class MergedColumnParallelLinear(ColumnParallelLinear): + """Packed linear layers with column parallelism. + + Similar to ColumnParallelLinear, but the weight matrix is concatenated + along the output dimension. When the weight matrix is loaded, the + different partitions are sharded separately. + + Args: + input_size: input dimension of the linear layer. + output_sizes: list of output dimensions of the linear layer. + bias: If true, add bias. + gather_output: If true, call all-gather on output and make the output + available to all GPUs, otherwise, every GPU will have + its own output. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + input_size: int, + output_sizes: list[int], + bias: bool = True, + gather_output: bool = False, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + self.output_sizes = output_sizes + tp_size = get_tp_world_size() + assert all(output_size % tp_size == 0 for output_size in output_sizes) + super().__init__( + input_size=input_size, + output_size=sum(output_sizes), + bias=bias, + gather_output=gather_output, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + ) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None, + ) -> None: + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + # Special case for per-tensor scale to load scalar into fused array. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (mlp). + # (e.g., Phi-3's gate_up_proj). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0 + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + current_shard_offset = 0 + shard_offsets: list[tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + for shard_id, shard_offset, shard_size in shard_offsets: + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size + ) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + assert loaded_shard_id < len(self.output_sizes) + tp_rank = get_tp_rank() + tp_size = get_tp_world_size() + if output_dim is not None: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size + shard_size = self.output_sizes[loaded_shard_id] // tp_size + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + start_idx = tp_rank * shard_size + if not is_sharded_weight: + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + # Special case for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_offset = loaded_shard_id * shard_size + param_data = param_data.narrow(0, shard_offset, shard_size) + + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id + ) + + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "MergedColumnParallelLinear, assume the weight is " + "the same for all partitions." + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def _load_fused_module_from_checkpoint( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ) -> None: + """Handle special case for models where MLP layers are already fused on disk. In this case, + we have no shard id. This function determmines the shard id by splitting these layers and + then calls the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + + current_shard_offset = 0 + shard_offsets: list[tuple[int, int, int]] = [] + for i, output_size in enumerate(self.output_sizes): + shard_offsets.append((i, current_shard_offset, output_size)) + current_shard_offset += output_size + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if ( + isinstance(param, PackedColumnParameter | PackedvLLMParameter) + and param.packed_dim == param.output_dim + ): + shard_size, shard_offset = ( + param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset + ) + ) + + loaded_weight_shard = loaded_weight.narrow( + param.output_dim, shard_offset, shard_size + ) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None, + ) -> None: + if loaded_shard_id is None: + if isinstance(param, PerTensorScaleParameter): + param.load_merged_column_weight( + loaded_weight=loaded_weight, shard_id=0 + ) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_merged_column_weight(loaded_weight=loaded_weight) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id < len(self.output_sizes) + + tp_size = get_tp_world_size() + + if isinstance(param, BlockQuantScaleParameter): + raise NotImplementedError("FP8 is not implemented yet") + # FIXME(will): add fp8 support + # from vllm.model_executor.layers.quantization.fp8 import ( + # Fp8LinearMethod, Fp8MoEMethod) + # assert self.quant_method is not None + # assert isinstance(self.quant_method, + # (Fp8LinearMethod, Fp8MoEMethod)) + # weight_block_size = self.quant_method.quant_config.weight_block_size + # assert weight_block_size is not None + # block_n, _ = weight_block_size[0], weight_block_size[1] + # shard_offset = ( + # (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // + # block_n) // tp_size + # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) // + # block_n // tp_size) + else: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size + shard_size = self.output_sizes[loaded_shard_id] // tp_size + + param.load_merged_column_weight( + loaded_weight=loaded_weight, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size, + ) + + +class QKVParallelLinear(ColumnParallelLinear): + """Linear layers for the attention's QKV transformation. + + Linear layers for the linear transformation of the query, key, and value + vectors in the attention layer. The weight matrix is concatenated along + the output dimension. The layer is parallelized along the head dimension. + When the number of key/value heads is smaller than the number of query + heads (e.g., multi-query/grouped-query attention), the key/value head may + be replicated while the query heads are partitioned. + + Args: + hidden_size: input hidden state size of the transformer. + head_size: size of each attention head. + total_num_heads: total number of attention query heads. + total_num_kv_heads: total number of attention key/value heads. If + None, assume total_num_kv_heads = total_num_heads. + bias: If true, add bias. + skip_bias_add: This was added to enable performance optimizations where + bias can be fused with other element-wise operations. we + skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict, including all parents + (e.g. model.layers.0.qkv_proj) + """ + + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int | None = None, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + self.hidden_size = hidden_size + self.head_size = head_size + self.total_num_heads = total_num_heads + if total_num_kv_heads is None: + total_num_kv_heads = total_num_heads + self.total_num_kv_heads = total_num_kv_heads + # Divide the weight matrix along the last dimension. + tp_size = get_tp_world_size() + self.num_heads = divide(self.total_num_heads, tp_size) + if tp_size >= self.total_num_kv_heads: + self.num_kv_heads = 1 + self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) + else: + self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) + self.num_kv_head_replicas = 1 + input_size = self.hidden_size + output_size = ( + (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size + ) + self.output_sizes = [ + self.num_heads * self.head_size * tp_size, # q_proj + self.num_kv_heads * self.head_size * tp_size, # k_proj + self.num_kv_heads * self.head_size * tp_size, # v_proj + ] + + super().__init__( + input_size=input_size, + output_size=output_size, + bias=bias, + gather_output=False, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + ) + + def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: + shard_offset_mapping = { + "q": 0, + "k": self.num_heads * self.head_size, + "v": (self.num_heads + self.num_kv_heads) * self.head_size, + "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size, + } + return shard_offset_mapping.get(loaded_shard_id) + + def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None: + shard_size_mapping = { + "q": self.num_heads * self.head_size, + "k": self.num_kv_heads * self.head_size, + "v": self.num_kv_heads * self.head_size, + } + return shard_size_mapping.get(loaded_shard_id) + + def _load_fused_module_from_checkpoint( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ): + """Handle special case for models where QKV layers are already fused on disk. In this case, + we have no shard id. This function determmines the shard id by splitting these layers and + then calls the weight loader using the shard id. + + An example of a model with these fused layers: + https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + """ + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ( + "k", + self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ( + "v", + (self.total_num_heads + self.total_num_kv_heads) + * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ] + + for shard_id, shard_offset, shard_size in shard_offsets: + # Special case for Quantization. + # If quantized, we need to adjust the offset and size to account + # for the packing. + if ( + isinstance(param, PackedColumnParameter | PackedvLLMParameter) + and param.packed_dim == param.output_dim + ): + shard_size, shard_offset = ( + param.adjust_shard_indexes_for_packing( + shard_size=shard_size, shard_offset=shard_offset + ) + ) + + loaded_weight_shard = loaded_weight.narrow( + param.output_dim, shard_offset, shard_size + ) + self.weight_loader_v2(param, loaded_weight_shard, shard_id) + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None, + ): + if loaded_shard_id is None: # special case for certain models + if isinstance(param, PerTensorScaleParameter): + param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0) + return + elif type(param) in (RowvLLMParameter, BasevLLMParameter): + param.load_qkv_weight(loaded_weight=loaded_weight) + return + # TODO: @dsikka - move to parameter.py + self._load_fused_module_from_checkpoint(param, loaded_weight) + return + + assert loaded_shard_id in ["q", "k", "v"] + + shard_offset = self._get_shard_offset_mapping(loaded_shard_id) + shard_size = self._get_shard_size_mapping(loaded_shard_id) + + param.load_qkv_weight( + loaded_weight=loaded_weight, + num_heads=self.num_kv_head_replicas, + shard_id=loaded_shard_id, + shard_offset=shard_offset, + shard_size=shard_size, + ) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None, + ): + + param_data = param.data + output_dim = getattr(param, "output_dim", None) + # Special case for AQLM codebooks. + is_metadata = getattr(param, "is_metadata", False) + + # Special case for per-tensor scales in fused case. + needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) + + if loaded_shard_id is None: + # Loaded weight is already fused on disk (qkv). + # (e.g., Phi-3's qkv_proj). + if output_dim is None: + if needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, 0 + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + return + shard_offsets = [ + # (shard_id, shard_offset, shard_size) + ("q", 0, self.total_num_heads * self.head_size), + ( + "k", + self.total_num_heads * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ( + "v", + (self.total_num_heads + self.total_num_kv_heads) + * self.head_size, + self.total_num_kv_heads * self.head_size, + ), + ] + + for shard_id, shard_offset, shard_size in shard_offsets: + loaded_weight_shard = loaded_weight.narrow( + output_dim, shard_offset, shard_size + ) + self.weight_loader(param, loaded_weight_shard, shard_id) + return + + tp_rank = get_tp_rank() + assert loaded_shard_id in ["q", "k", "v"] + + # If output dim is defined, use the default loading process. + if output_dim is not None: + if loaded_shard_id == "q": + shard_offset = 0 + shard_size = self.num_heads * self.head_size + elif loaded_shard_id == "k": + shard_offset = self.num_heads * self.head_size + shard_size = self.num_kv_heads * self.head_size + elif loaded_shard_id == "v": + shard_offset = ( + self.num_heads + self.num_kv_heads + ) * self.head_size + shard_size = self.num_kv_heads * self.head_size + + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + shard_idx = 0 + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + if loaded_shard_id == "q": + shard_idx = tp_rank + else: + shard_idx = tp_rank // self.num_kv_head_replicas + start_idx = shard_idx * shard_size + + if not is_sharded_weight: + loaded_weight = loaded_weight.narrow( + output_dim, start_idx, shard_size + ) + + # Special case for for AQLM codebooks. + elif is_metadata: + # metadata indicates fixed size concatenated along dim 0 + shard_size = loaded_weight.shape[0] + shard_index = ["q", "k", "v"].index(loaded_shard_id) + param_data = param_data.narrow( + 0, shard_index * shard_size, shard_size + ) + # Special case for per-tensor scales in fused case. + elif needs_scalar_to_array: + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id + ) + else: + ignore_warning = getattr(param, "ignore_warning", False) + if not ignore_warning: + logger.warning( + "Loading a weight without `output_dim` attribute in " + "QKVParallelLinear, assume the weight is the same " + "for all partitions." + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class RowParallelLinear(LinearBase): + """Linear layer with row parallelism. + + The linear layer is defined as Y = XA + b. A is parallelized along + its first dimension and X along its second dimension as: + - - + | A_1 | + | . | + A = | . | X = [X_1, ..., X_p] + | . | + | A_p | + - - + Arguments: + input_size: first dimension of matrix A. + output_size: second dimension of matrix A. + bias: If true, add bias. Note that bias is not parallelized. + input_is_parallel: If true, we assume that the input is already + split across the GPUs and we do not split + again. + skip_bias_add: This was added to enable performance optimization where + bias can be fused with other element-wise operations. + We skip adding bias but instead return it. + params_dtype: Data type for the parameters. + quant_config: Quantization configure. + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + input_is_parallel: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + reduce_results: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + # Divide the weight matrix along the first dimension. + self.tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.input_size_per_partition = divide(input_size, self.tp_size) + self.output_size_per_partition = output_size + self.output_partition_sizes = [output_size] + + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix, + ) + + self.input_is_parallel = input_is_parallel + self.reduce_results = reduce_results + + assert self.quant_method is not None + self.quant_method.create_weights( + layer=self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ + in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ), + ) + if not reduce_results and (bias and not skip_bias_add): + raise ValueError( + "When not reduce the results, adding bias to the " + "results can lead to incorrect results" + ) + + if bias: + self.bias = Parameter( + torch.empty(self.output_size, dtype=params_dtype) + ) + set_weight_attrs( + self.bias, + { + "output_dim": 0, + "weight_loader": self.weight_loader, + }, + ) + else: + self.register_parameter("bias", None) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + tp_rank = get_tp_rank() + input_dim = getattr(param, "input_dim", None) + is_sharded_weight = getattr(param, "is_sharded_weight", False) + # bitsandbytes loads the weights of the specific portion + # no need to narrow + is_sharded_weight = is_sharded_weight + + param_data = param.data + if input_dim is not None and not is_sharded_weight: + shard_size = param_data.shape[input_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow( + input_dim, start_idx, shard_size + ) + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def weight_loader_v2( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ): + + # Special case for loading scales off disk, which often do not + # have a shape (such as in the case of AutoFP8). + if len(loaded_weight.shape) == 0: + assert loaded_weight.numel() == 1 + loaded_weight = loaded_weight.reshape(1) + + param.load_row_parallel_weight(loaded_weight=loaded_weight) + + def forward(self, input_) -> tuple[torch.Tensor, Parameter | None]: + if self.input_is_parallel: + input_parallel = input_ + else: + tp_rank = get_tp_rank() + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.tp_size + ) + input_parallel = splitted_input[tp_rank].contiguous() + + # Matrix multiply. + assert self.quant_method is not None + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in TP>1 case) + bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias + output_parallel = self.quant_method.apply( + self, input_parallel, bias=bias_ + ) + if self.reduce_results and self.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + + output_bias = self.bias if self.skip_bias_add else None + + return output, output_bias + + def extra_repr(self) -> str: + s = f"input_features={self.input_size_per_partition}" + s += f", output_features={self.output_size}" + s += f", bias={self.bias is not None}" + s += f", tp_size={self.tp_size}" + s += f", reduce_results={self.reduce_results}" + return s diff --git a/worldcompass/fastvideo/layers/lora/linear.py b/worldcompass/fastvideo/layers/lora/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..2f472d197506cc44e08e69b3b9b6027bf133f804 --- /dev/null +++ b/worldcompass/fastvideo/layers/lora/linear.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: Apache-2.0 +# Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py + +import math + +import torch +from torch import nn +from torch.distributed._composable.fsdp import ( + CPUOffloadPolicy, + OffloadPolicy, + fully_shard, +) +from torch.distributed.tensor import DTensor + +from fastvideo.distributed import ( + get_local_torch_device, + get_tp_rank, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from fastvideo.layers.linear import ( + ColumnParallelLinear, + LinearBase, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.utils import get_mixed_precision_state + +torch._dynamo.config.recompile_limit = 16 + + +class BaseLayerWithLoRA(nn.Module): + def __init__( + self, + base_layer: nn.Module, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ): + super().__init__() + self.base_layer: nn.Module = base_layer + + self.merged: bool = False + self.cpu_weight = base_layer.weight.to("cpu") + # indicates adapter weights don't contain this layer + # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging) + self.disable_lora: bool = False + self.lora_rank = lora_rank + self.lora_alpha = lora_alpha + self.training_mode = training_mode + self.lora_path: str | None = None + + if training_mode: + assert ( + self.lora_rank is not None + ), "LoRA rank must be set for training mode" + if self.lora_rank is None or self.lora_alpha is None: + self.lora_alpha = lora_rank + self.base_layer.requires_grad_(False) + in_dim = self.base_layer.weight.shape[1] + out_dim = self.base_layer.weight.shape[0] + self.lora_A = nn.Parameter( + torch.zeros( + self.lora_rank, + in_dim, + device=self.base_layer.weight.device, + dtype=self.base_layer.weight.dtype, + ) + ) + self.lora_B = nn.Parameter( + torch.zeros( + out_dim, + self.lora_rank, + device=self.base_layer.weight.device, + dtype=self.base_layer.weight.dtype, + ) + ) + torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + torch.nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5)) + else: + self.lora_A = None + self.lora_B = None + + @torch.compile() + def forward(self, x: torch.Tensor) -> torch.Tensor: + lora_A = self.lora_A + lora_B = self.lora_B + if isinstance(self.lora_B, DTensor): + lora_B = self.lora_B.to_local() + lora_A = self.lora_A.to_local() + + if (self.training_mode or not self.merged) and not self.disable_lora: + delta = x @ ( + self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) + @ self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) + ) + if self.lora_alpha != self.lora_rank: + delta = delta * ( + self.lora_alpha / self.lora_rank # type: ignore + ) # type: ignore + out, output_bias = self.base_layer(x) + return out + delta, output_bias + else: + out, output_bias = self.base_layer(x) + return out.to(x), output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + return B + + def set_lora_weights( + self, + A: torch.Tensor, + B: torch.Tensor, + training_mode: bool = False, + lora_path: str | None = None, + ) -> None: + self.lora_A = torch.nn.Parameter( + A + ) # share storage with weights in the pipeline + self.lora_B = torch.nn.Parameter(B) + self.disable_lora = False + if not training_mode: + self.merge_lora_weights() + self.lora_path = lora_path + + @torch.no_grad() + def merge_lora_weights(self) -> None: + if self.disable_lora: + return + + if self.merged: + self.unmerge_lora_weights() + assert ( + self.lora_A is not None and self.lora_B is not None + ), "LoRA weights not set. Please set them first." + if isinstance(self.base_layer.weight, DTensor): + mesh = self.base_layer.weight.data.device_mesh + unsharded_base_layer = ReplicatedLinear( + input_size=self.base_layer.input_size, + output_size=self.base_layer.output_size, + bias=getattr(self.base_layer, "bias", None) is not None, + skip_bias_add=self.base_layer.skip_bias_add, + params_dtype=self.base_layer.params_dtype, + quant_config=self.base_layer.quant_config, + prefix=self.base_layer.prefix, + ) + # Using offload param is on CPU, so current_device is for "CPU -> GPU -> merge -> CPU" + current_device = self.base_layer.weight.data.device + data = self.base_layer.weight.data.to( + get_local_torch_device() + ).full_tensor() + data += self.slice_lora_b_weights(self.lora_B).to( + data + ) @ self.slice_lora_a_weights(self.lora_A).to(data) + unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + if isinstance(getattr(self.base_layer, "bias", None), DTensor): + unsharded_base_layer.bias = nn.Parameter( + self.base_layer.bias.to( + get_local_torch_device(), non_blocking=True + ) + .full_tensor() + .to(current_device) + ) + + offload_policy = ( + CPUOffloadPolicy() + if "cpu" in str(current_device) + else OffloadPolicy() + ) + mp_policy = get_mixed_precision_state().mp_policy + + self.base_layer = fully_shard( + unsharded_base_layer, + mesh=mesh, + mp_policy=mp_policy, + offload_policy=offload_policy, + ) + else: + current_device = self.base_layer.weight.data.device + data = self.base_layer.weight.data.to(get_local_torch_device()) + data += self.slice_lora_b_weights( + self.lora_B.to(data) + ) @ self.slice_lora_a_weights(self.lora_A.to(data)) + self.base_layer.weight.data = data.to( + current_device, non_blocking=True + ) + + self.merged = True + + @torch.no_grad() + # @torch.compile(dynamic=True) + def unmerge_lora_weights(self) -> None: + if self.disable_lora: + return + + if not self.merged: + raise ValueError( + "LoRA weights not merged. Please merge them first before unmerging." + ) + + # avoid precision loss + if isinstance(self.base_layer.weight, DTensor): + device = self.base_layer.weight.data.device + self.base_layer.weight = nn.Parameter( + self.cpu_weight.to(device, non_blocking=True) + ) + else: + self.base_layer.weight.data = self.cpu_weight.data.to( + self.base_layer.weight, non_blocking=True + ) + + self.merged = False + + +class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): + """Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation). + + Note: The current version does not yet implement the LoRA functionality. + This class behaves exactly the same as the base VocabParallelEmbedding. + Future versions will integrate LoRA functionality to support efficient parameter fine-tuning. + """ + + def __init__( + self, + base_layer: VocabParallelEmbedding, + ) -> None: + super().__init__(base_layer) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + raise NotImplementedError( + "We don't support VocabParallelEmbeddingWithLoRA yet." + ) + + +class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): + def __init__( + self, + base_layer: ColumnParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + # duplicate the logic in ColumnParallelLinear + bias = ( + self.base_layer.bias if not self.base_layer.skip_bias_add else None + ) + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_, bias + ) + if self.base_layer.gather_output: + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + output_bias = ( + self.base_layer.bias if self.base_layer.skip_bias_add else None + ) + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + shard_size = self.base_layer.output_partition_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + B = B[start_idx:end_idx, :] + return B + + +class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + def __init__( + self, + base_layer: MergedColumnParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A.to(self.base_layer.weight) + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + # Since the outputs for both gate and up are identical, we use a random one. + shard_size = self.base_layer.output_partition_sizes[0] + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + return B[:, start_idx:end_idx, :] + + +class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): + def __init__( + self, + base_layer: QKVParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + return A + + def slice_lora_b_weights( + self, B: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + tp_rank = get_tp_rank() + B_q, B_kv = B + base_layer = self.base_layer + q_proj_shard_size = base_layer.q_proj_shard_size + kv_proj_shard_size = base_layer.kv_proj_shard_size + num_kv_head_replicas = base_layer.num_kv_head_replicas + + q_start_idx = q_proj_shard_size * tp_rank + q_end_idx = q_start_idx + q_proj_shard_size + + kv_shard_id = tp_rank // num_kv_head_replicas + kv_start_idx = kv_proj_shard_size * kv_shard_id + kv_end_idx = kv_start_idx + kv_proj_shard_size + + return ( + B_q[q_start_idx:q_end_idx, :], + B_kv[:, kv_start_idx:kv_end_idx, :], + ) + + +class RowParallelLinearWithLoRA(BaseLayerWithLoRA): + def __init__( + self, + base_layer: RowParallelLinear, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, + ) -> None: + super().__init__(base_layer, lora_rank, lora_alpha, training_mode) + + def forward(self, input_: torch.Tensor): + # duplicate the logic in RowParallelLinear + if self.base_layer.input_is_parallel: + input_parallel = input_ + else: + tp_rank = get_tp_rank() + splitted_input = split_tensor_along_last_dim( + input_, num_partitions=self.base_layer.tp_size + ) + input_parallel = splitted_input[tp_rank].contiguous() + output_parallel = self.base_layer.quant_method.apply( + self.base_layer, input_parallel + ) + + if self.set_lora: + output_parallel = self.apply_lora(output_parallel, input_parallel) + + if self.base_layer.reduce_results and self.base_layer.tp_size > 1: + output_ = tensor_model_parallel_all_reduce(output_parallel) + else: + output_ = output_parallel + + if not self.base_layer.skip_bias_add: + output = ( + output_ + self.base_layer.bias + if self.base_layer.bias is not None + else output_ + ) + output_bias = None + else: + output = output_ + output_bias = self.base_layer.bias + return output, output_bias + + def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: + tp_rank = get_tp_rank() + shard_size = self.base_layer.input_size_per_partition + start_idx = tp_rank * shard_size + end_idx = (tp_rank + 1) * shard_size + A = A[:, start_idx:end_idx].contiguous() + return A + + def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: + return B + + +def get_lora_layer( + layer: nn.Module, + lora_rank: int | None = None, + lora_alpha: int | None = None, + training_mode: bool = False, +) -> BaseLayerWithLoRA | None: + supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = { + # the order matters + # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, + QKVParallelLinear: QKVParallelLinearWithLoRA, + MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA, + ColumnParallelLinear: ColumnParallelLinearWithLoRA, + RowParallelLinear: RowParallelLinearWithLoRA, + ReplicatedLinear: BaseLayerWithLoRA, + } + for src_layer_type, lora_layer_type in supported_layer_types.items(): + if isinstance( + layer, src_layer_type + ): # pylint: disable=unidiomatic-typecheck + ret = lora_layer_type( + layer, + lora_rank=lora_rank, + lora_alpha=lora_alpha, + training_mode=training_mode, + ) + return ret + return None + + +# source: https://github.com/vllm-project/vllm/blob/93b38bea5dd03e1b140ca997dfaadef86f8f1855/vllm/lora/utils.py#L9 +def replace_submodule( + model: nn.Module, module_name: str, new_module: nn.Module +) -> nn.Module: + """Replace a submodule in a model with a new module.""" + parent = model.get_submodule(".".join(module_name.split(".")[:-1])) + target_name = module_name.split(".")[-1] + setattr(parent, target_name, new_module) + return new_module diff --git a/worldcompass/fastvideo/layers/mlp.py b/worldcompass/fastvideo/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..10de38c29062e454682a47639651e78fcfb82a80 --- /dev/null +++ b/worldcompass/fastvideo/layers/mlp.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn + +from fastvideo.layers.activation import get_act_fn +from fastvideo.layers.linear import ReplicatedLinear + + +class MLP(nn.Module): + """MLP for DiT blocks, NO gated linear units.""" + + def __init__( + self, + input_dim: int, + mlp_hidden_dim: int, + output_dim: int | None = None, + bias: bool = True, + act_type: str = "gelu_pytorch_tanh", + dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.fc_in = ReplicatedLinear( + input_dim, + mlp_hidden_dim, # For activation func like SiLU that need 2x width + bias=bias, + params_dtype=dtype, + ) + + self.act = get_act_fn(act_type) + if output_dim is None: + output_dim = input_dim + self.fc_out = ReplicatedLinear( + mlp_hidden_dim, output_dim, bias=bias, params_dtype=dtype + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc_in(x) + x = self.act(x) + x, _ = self.fc_out(x) + return x diff --git a/worldcompass/fastvideo/layers/quantization/__init__.py b/worldcompass/fastvideo/layers/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8956306a2f80ca876d2be96d7fb373a3db574547 --- /dev/null +++ b/worldcompass/fastvideo/layers/quantization/__init__.py @@ -0,0 +1,67 @@ +from typing import Literal, get_args + +from fastvideo.layers.quantization.base_config import QuantizationConfig + +QuantizationMethods = Literal[None] + +QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) + +# The customized quantization methods which will be added to this dict. +_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {} + + +def register_quantization_config(quantization: str): + """Register a customized vllm quantization config. + + When a quantization method is not supported by vllm, you can register a customized + quantization config to support it. + + Args: + quantization (str): The quantization method name. + + Examples: + >>> from fastvideo.layers.quantization import register_quantization_config + >>> from fastvideo.layers.quantization import get_quantization_config + >>> from fastvideo.layers.quantization.base_config import QuantizationConfig + >>> + >>> @register_quantization_config("my_quant") + ... class MyQuantConfig(QuantizationConfig): + ... pass + >>> + >>> get_quantization_config("my_quant") + + """ # noqa: E501 + + def _wrapper(quant_config_cls): + if quantization in QUANTIZATION_METHODS: + raise ValueError( + f"The quantization method `{quantization}` is already exists." + ) + if not issubclass(quant_config_cls, QuantizationConfig): + raise ValueError( + "The quantization config must be a subclass of `QuantizationConfig`." + ) + _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls + QUANTIZATION_METHODS.append(quantization) + return quant_config_cls + + return _wrapper + + +def get_quantization_config(quantization: str) -> type[QuantizationConfig]: + if quantization not in QUANTIZATION_METHODS: + raise ValueError(f"Invalid quantization method: {quantization}") + + method_to_config: dict[str, type[QuantizationConfig]] = {} + # Update the `method_to_config` with customized quantization methods. + method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) + + return method_to_config[quantization] + + +all = [ + "QuantizationMethods", + "QuantizationConfig", + "get_quantization_config", + "QUANTIZATION_METHODS", +] diff --git a/worldcompass/fastvideo/layers/quantization/base_config.py b/worldcompass/fastvideo/layers/quantization/base_config.py new file mode 100644 index 0000000000000000000000000000000000000000..2a1b74c676540107dc96ed21b5e08bc77039514e --- /dev/null +++ b/worldcompass/fastvideo/layers/quantization/base_config.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py + +import inspect +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +if TYPE_CHECKING: + from fastvideo.layers.quantization import QuantizationMethods +else: + QuantizationMethods = str + + +class QuantizeMethodBase(ABC): + """Base class for different quantized methods.""" + + @abstractmethod + def create_weights( + self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs + ): + """Create weights for a layer. + + The weights will be set as attributes of the layer. + """ + raise NotImplementedError + + @abstractmethod + def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + + Expects create_weights to have been called before on the layer. + """ + raise NotImplementedError + + # Not required functions + def embedding( + self, layer: torch.nn.Module, *args, **kwargs + ) -> torch.Tensor: + """Gather embeddings in the layer based on indices in the input tensor. + + Expects create_weights to have been called before on the layer. + """ + raise NotImplementedError + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Process the weight after loading. + + This can be used for example, to transpose weights for computation. + """ + return + + +def method_has_implemented_embedding( + method_class: type[QuantizeMethodBase], +) -> bool: + """Not all quant methods have embedding implemented, so we need to check that it exists for our + given method. + + We check this by making sure the function has been changed from the base implementation. + """ + base_embedding = inspect.getattr_static( + QuantizeMethodBase, "embedding", None + ) + class_embedding = inspect.getattr_static(method_class, "embedding", None) + + return class_embedding is not None and class_embedding is not base_embedding + + +class QuantizationConfig(ABC): + """Base class for quantization configs.""" + + def __init__(self): + super().__init__() + # mapping is updated by models as they initialize + self.packed_modules_mapping: dict[str, list[str]] = dict() + + @abstractmethod + def get_name(self) -> QuantizationMethods: + """Name of the quantization method.""" + raise NotImplementedError + + @abstractmethod + def get_supported_act_dtypes(self) -> list[torch.dtype]: + """List of supported activation dtypes.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """Minimum GPU capability to support the quantization method. + + E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA + kernels used by the quantization method. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_config_filenames() -> list[str]: + """List of filenames to search for in the model directory.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig": + """Create a config class from the model's quantization config.""" + raise NotImplementedError + + @classmethod + def override_quantization_method( + cls, hf_quant_cfg, user_quant + ) -> QuantizationMethods | None: + """Detects if this quantization method can support a given checkpoint format by overriding + the user specified quantization method -- this method should only be overwritten by + subclasses in exceptional circumstances.""" + return None + + @staticmethod + def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any: + """Get a value from the model's quantization config.""" + for key in keys: + if key in config: + return config[key] + raise ValueError( + f"Cannot find any of {keys} in the model's quantization config." + ) + + @staticmethod + def get_from_keys_or( + config: dict[str, Any], keys: list[str], default: Any + ) -> Any: + """Get a optional value from the model's quantization config.""" + try: + return QuantizationConfig.get_from_keys(config, keys) + except ValueError: + return default + + @abstractmethod + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> QuantizeMethodBase | None: + """Get the quantize method to use for the quantized layer. + + Args: + layer: The layer for the quant method. + prefix: The full name of the layer in the state dict + Returns: + The quantize method. None if the given layer doesn't support quant + method. + """ + raise NotImplementedError + + def get_cache_scale(self, name: str) -> str | None: + return None diff --git a/worldcompass/fastvideo/layers/rotary_embedding.py b/worldcompass/fastvideo/layers/rotary_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..24866b366251c36c31543480399dfd57b0832b6b --- /dev/null +++ b/worldcompass/fastvideo/layers/rotary_embedding.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Rotary Positional Embeddings.""" + +from typing import Any + +import torch + +from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.layers.custom_op import CustomOp + + +def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + +def _apply_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + is_neox_style: bool, +) -> torch.Tensor: + """ + Args: + x: [num_tokens, num_heads, head_size] + cos: [num_tokens, head_size // 2] + sin: [num_tokens, head_size // 2] + is_neox_style: Whether to use the Neox-style or GPT-J-style rotary + positional embeddings. + """ + # cos = cos.unsqueeze(-2).to(x.dtype) + # sin = sin.unsqueeze(-2).to(x.dtype) + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + if is_neox_style: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = (x1.float() * cos - x2.float() * sin).type_as(x) + o2 = (x2.float() * cos + x1.float() * sin).type_as(x) + if is_neox_style: + return torch.cat((o1, o2), dim=-1) + else: + return torch.stack((o1, o2), dim=-1).flatten(-2) + + +@CustomOp.register("rotary_embedding") +class RotaryEmbedding(CustomOp): + """Original rotary positional embedding.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int | float, + is_neox_style: bool, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.head_size = head_size + self.rotary_dim = rotary_dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.is_neox_style = is_neox_style + self.dtype = dtype + + cache = self._compute_cos_sin_cache() + cache = cache.to(dtype) + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) + + def _compute_inv_freq(self, base: int | float) -> torch.Tensor: + """Compute the inverse frequency.""" + # NOTE(woosuk): To exactly match the HF implementation, we need to + # use CPU to compute the cache and then move it to GPU. However, we + # create the cache on GPU for faster initialization. This may cause + # a slight numerical difference between the HF implementation and ours. + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) + / self.rotary_dim + ) + ) + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" + inv_freq = self._compute_inv_freq(self.base) + t = torch.arange(self.max_position_embeddings, dtype=torch.float) + + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """A PyTorch-native implementation of forward().""" + if offsets is not None: + positions = positions + offsets + positions = positions.flatten() + num_tokens = positions.shape[0] + cos_sin = self.cos_sin_cache.index_select(0, positions) + cos, sin = cos_sin.chunk(2, dim=-1) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + return s + + +def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]: + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd( + start: int | tuple[int, ...], *args: int | tuple[int, ...], dim: int = 2 +) -> torch.Tensor: + """Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = tuple(stop[i] - start[i] for i in range(dim)) + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +def get_1d_rotary_pos_embed( + dim: int, + pos: torch.FloatTensor | int, + theta: float = 10000.0, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Precompute the frequency tensor for complex exponential (cis) with given dimensions. (Note: + `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0. + + Returns: + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor ** (dim / (dim - 2)) + + freqs = 1.0 / ( + theta ** (torch.arange(0, dim, 2)[: (dim // 2)].to(dtype) / dim) + ) # [D/2] + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + freqs_cos = freqs.cos() # [S, D/2] + freqs_sin = freqs.sin() # [S, D/2] + return freqs_cos, freqs_sin + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + theta_rescale_factor: float | list[float] = 1.0, + interpolation_factor: float | list[float] = 1.0, + shard_dim: int = 0, + sp_rank: int = 0, + sp_world_size: int = 1, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + Supports sequence parallelism by allowing sharding of a specific dimension. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + interpolation_factor (float): Factor to scale positions. Defaults to 1.0. + shard_dim (int): Which dimension to shard for sequence parallelism. Defaults to 0. + sp_rank (int): Rank in the sequence parallel group. Defaults to 0. + sp_world_size (int): World size of the sequence parallel group. Defaults to 1. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D/2] + """ + # Get the full grid + full_grid = get_meshgrid_nd( + start, *args, dim=len(rope_dim_list) + ) # [3, W, H, D] / [2, W, H] + + # Shard the grid if using sequence parallelism (sp_world_size > 1) + assert shard_dim < len( + rope_dim_list + ), f"shard_dim {shard_dim} must be less than number of dimensions {len(rope_dim_list)}" + if sp_world_size > 1: + # Get the shape of the full grid + grid_shape = list(full_grid.shape[1:]) + + # Ensure the dimension to shard is divisible by sp_world_size + assert grid_shape[shard_dim] % sp_world_size == 0, ( + f"Dimension {shard_dim} with size {grid_shape[shard_dim]} is not divisible " + f"by sequence parallel world size {sp_world_size}" + ) + + # Compute the start and end indices for this rank's shard + shard_size = grid_shape[shard_dim] // sp_world_size + start_idx = sp_rank * shard_size + end_idx = (sp_rank + 1) * shard_size + + # Create slicing indices for each dimension + slice_indices = [slice(None) for _ in range(len(grid_shape))] + slice_indices[shard_dim] = slice(start_idx, end_idx) + + # Shard the grid + # Update grid shape for the sharded dimension + grid_shape[shard_dim] = grid_shape[shard_dim] // sp_world_size + grid = torch.empty( + (len(rope_dim_list),) + tuple(grid_shape), dtype=full_grid.dtype + ) + for i in range(len(rope_dim_list)): + grid[i] = full_grid[i][tuple(slice_indices)] + else: + grid = full_grid + + if isinstance(theta_rescale_factor, int | float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif ( + isinstance(theta_rescale_factor, list) + and len(theta_rescale_factor) == 1 + ): + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int | float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif ( + isinstance(interpolation_factor, list) + and len(interpolation_factor) == 1 + ): + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + dtype=dtype, + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + + +def get_rotary_pos_embed( + rope_sizes, + hidden_size, + heads_num, + rope_dim_list, + rope_theta, + theta_rescale_factor=1.0, + interpolation_factor=1.0, + shard_dim: int = 0, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate rotary positional embeddings for the given sizes. + + Args: + rope_sizes: Tuple of dimensions (t, h, w) + hidden_size: Hidden dimension size + heads_num: Number of attention heads + rope_dim_list: List of dimensions for each axis, or None + rope_theta: Base for frequency calculations + theta_rescale_factor: Rescale factor for theta. Defaults to 1.0 + interpolation_factor: Factor to scale positions. Defaults to 1.0 + shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0. + + Returns: + Tuple of (cos, sin) tensors for rotary embeddings + """ + + target_ndim = 3 + head_dim = hidden_size // heads_num + + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + + assert ( + sum(rope_dim_list) == head_dim + ), "sum(rope_dim_list) should equal to head_dim of attention layer" + + # Get SP info + sp_group = get_sp_group() + sp_rank = sp_group.rank_in_group + sp_world_size = sp_group.world_size + + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=rope_theta, + theta_rescale_factor=theta_rescale_factor, + interpolation_factor=interpolation_factor, + shard_dim=shard_dim, + sp_rank=sp_rank, + sp_world_size=sp_world_size, + dtype=dtype, + ) + return freqs_cos, freqs_sin + + +_ROPE_DICT: dict[tuple, RotaryEmbedding] = {} + + +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: int | float, + is_neox_style: bool = True, + rope_scaling: dict[str, Any] | None = None, + dtype: torch.dtype | None = None, + partial_rotary_factor: float = 1.0, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = ( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + rope_scaling_args, + dtype, + ) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + if rope_scaling is None: + rotary_emb = RotaryEmbedding( + head_size, rotary_dim, max_position, base, is_neox_style, dtype + ) + else: + raise ValueError(f"Unknown RoPE scaling {rope_scaling}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb diff --git a/worldcompass/fastvideo/layers/utils.py b/worldcompass/fastvideo/layers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3ea2dce0328cb6ad337dc31deced8cc34bd1d6e7 --- /dev/null +++ b/worldcompass/fastvideo/layers/utils.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py +"""Utility methods for model layers.""" + +import torch + + +def get_token_bin_counts_and_mask( + tokens: torch.Tensor, + vocab_size: int, + num_seqs: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # Compute the bin counts for the tokens. + # vocab_size + 1 for padding. + bin_counts = torch.zeros( + (num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device + ) + bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) + bin_counts = bin_counts[:, :vocab_size] + mask = bin_counts > 0 + + return bin_counts, mask diff --git a/worldcompass/fastvideo/layers/visual_embedding.py b/worldcompass/fastvideo/layers/visual_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..ee31ac24c773988d19f7b7744e23578a9c5d866d --- /dev/null +++ b/worldcompass/fastvideo/layers/visual_embedding.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch +import torch.nn as nn + +from fastvideo.layers.activation import get_act_fn +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding. + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in + https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + prefix: str = "", + ): + super().__init__() + # Convert patch_size to 2-tuple + if isinstance(patch_size, list | tuple): + if len(patch_size) == 1: + patch_size = (patch_size[0], patch_size[0]) + else: + patch_size = (patch_size, patch_size) + + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + dtype=dtype, + ) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TimestepEmbedder(nn.Module): + """Embeds scalar timesteps into vector representations.""" + + def __init__( + self, + hidden_size, + act_layer="silu", + frequency_embedding_size=256, + max_period=10000, + dtype=None, + freq_dtype=torch.float32, + prefix: str = "", + ): + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + + self.mlp = MLP( + frequency_embedding_size, + hidden_size, + hidden_size, + act_type=act_layer, + dtype=dtype, + ) + self.freq_dtype = freq_dtype + + def forward(self, t: torch.Tensor) -> torch.Tensor: + t_freq = timestep_embedding( + t, + self.frequency_embedding_size, + self.max_period, + dtype=self.freq_dtype, + ).to(self.mlp.fc_in.weight.dtype) + # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +def timestep_embedding( + t: torch.Tensor, + dim: int, + max_period: int = 10000, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create sinusoidal timestep embeddings. + + Args: + t: Tensor of shape [B] with timesteps + dim: Embedding dimension + max_period: Controls the minimum frequency of the embeddings + + Returns: + Tensor of shape [B, dim] with embeddings + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=dtype) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + +class ModulateProjection(nn.Module): + """Modulation layer for DiT blocks.""" + + def __init__( + self, + hidden_size: int, + factor: int = 2, + act_layer: str = "silu", + dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.factor = factor + self.hidden_size = hidden_size + self.linear = ReplicatedLinear( + hidden_size, hidden_size * factor, bias=True, params_dtype=dtype + ) + self.act = get_act_fn(act_layer) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.act(x) + x, _ = self.linear(x) + return x + + +def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor: + """Convert patched representation back to image space. + + Args: + x: Tensor of shape [B, T*H*W, C*P_t*P_h*P_w] + t, h, w: Temporal and spatial dimensions + + Returns: + Unpatchified tensor of shape [B, C, T*P_t, H*P_h, W*P_w] + """ + assert x.ndim == 3, f"x.ndim: {x.ndim}" + assert len(patch_size) == 3, f"patch_size: {patch_size}" + assert ( + t * h * w == x.shape[1] + ), f"t * h * w: {t * h * w}, x.shape[1]: {x.shape[1]}" + c = channels + pt, ph, pw = patch_size + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs diff --git a/worldcompass/fastvideo/layers/vocab_parallel_embedding.py b/worldcompass/fastvideo/layers/vocab_parallel_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..14f569a476c337c7198ca69f1afd5ed869c4ce11 --- /dev/null +++ b/worldcompass/fastvideo/layers/vocab_parallel_embedding.py @@ -0,0 +1,506 @@ +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter, UninitializedParameter + +from fastvideo.distributed import ( + divide, + get_tp_rank, + get_tp_world_size, + tensor_model_parallel_all_reduce, +) +from fastvideo.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + method_has_implemented_embedding, +) +from fastvideo.models.parameter import BasevLLMParameter +from fastvideo.models.utils import set_weight_attrs +from fastvideo.platforms import current_platform + +DEFAULT_VOCAB_PADDING_SIZE = 64 + + +class UnquantizedEmbeddingMethod(QuantizeMethodBase): + """Unquantized method for embeddings.""" + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Create weights for embedding layer.""" + + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return F.linear(x, layer.weight, bias) + + def embedding( + self, layer: torch.nn.Module, input_: torch.Tensor + ) -> torch.Tensor: + return F.embedding(input_, layer.weight) + + +def pad_vocab_size( + vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE +) -> int: + """Pad the vocab size to the given value.""" + return ((vocab_size + pad_to - 1) // pad_to) * pad_to + + +def vocab_range_from_per_partition_vocab_size( + per_partition_vocab_size: int, rank: int, offset: int = 0 +) -> Sequence[int]: + index_f = rank * per_partition_vocab_size + index_l = index_f + per_partition_vocab_size + return index_f + offset, index_l + offset + + +def vocab_range_from_global_vocab_size( + global_vocab_size: int, rank: int, world_size: int, offset: int = 0 +) -> Sequence[int]: + per_partition_vocab_size = divide(global_vocab_size, world_size) + return vocab_range_from_per_partition_vocab_size( + per_partition_vocab_size, rank, offset=offset + ) + + +@dataclass +class VocabParallelEmbeddingShardIndices: + """Indices for a shard of a vocab parallel embedding.""" + + padded_org_vocab_start_index: int + padded_org_vocab_end_index: int + padded_added_vocab_start_index: int + padded_added_vocab_end_index: int + + org_vocab_start_index: int + org_vocab_end_index: int + added_vocab_start_index: int + added_vocab_end_index: int + + @property + def num_org_elements(self) -> int: + return self.org_vocab_end_index - self.org_vocab_start_index + + @property + def num_added_elements(self) -> int: + return self.added_vocab_end_index - self.added_vocab_start_index + + @property + def num_org_elements_padded(self) -> int: + return ( + self.padded_org_vocab_end_index - self.padded_org_vocab_start_index + ) + + @property + def num_added_elements_padded(self) -> int: + return ( + self.padded_added_vocab_end_index + - self.padded_added_vocab_start_index + ) + + @property + def num_org_vocab_padding(self) -> int: + return self.num_org_elements_padded - self.num_org_elements + + @property + def num_added_vocab_padding(self) -> int: + return self.num_added_elements_padded - self.num_added_elements + + @property + def num_elements_padded(self) -> int: + return self.num_org_elements_padded + self.num_added_elements_padded + + def __post_init__(self): + # sanity checks + assert ( + self.padded_org_vocab_start_index <= self.padded_org_vocab_end_index + ) + assert ( + self.padded_added_vocab_start_index + <= self.padded_added_vocab_end_index + ) + + assert self.org_vocab_start_index <= self.org_vocab_end_index + assert self.added_vocab_start_index <= self.added_vocab_end_index + + assert self.org_vocab_start_index <= self.padded_org_vocab_start_index + assert ( + self.added_vocab_start_index <= self.padded_added_vocab_start_index + ) + assert self.org_vocab_end_index <= self.padded_org_vocab_end_index + assert self.added_vocab_end_index <= self.padded_added_vocab_end_index + + assert self.num_org_elements <= self.num_org_elements_padded + assert self.num_added_elements <= self.num_added_elements_padded + + +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) +def get_masked_input_and_mask( + input_: torch.Tensor, + org_vocab_start_index: int, + org_vocab_end_index: int, + num_org_vocab_padding: int, + added_vocab_start_index: int, + added_vocab_end_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # torch.compile will fuse all of the pointwise ops below + # into a single kernel, making it very fast + org_vocab_mask = (input_ >= org_vocab_start_index) & ( + input_ < org_vocab_end_index + ) + added_vocab_mask = (input_ >= added_vocab_start_index) & ( + input_ < added_vocab_end_index + ) + added_offset = ( + added_vocab_start_index + - (org_vocab_end_index - org_vocab_start_index) + - num_org_vocab_padding + ) + valid_offset = (org_vocab_start_index * org_vocab_mask) + ( + added_offset * added_vocab_mask + ) + vocab_mask = org_vocab_mask | added_vocab_mask + input_ = vocab_mask * (input_ - valid_offset) + return input_, ~vocab_mask + + +class VocabParallelEmbedding(torch.nn.Module): + """Embedding parallelized in the vocabulary dimension. + + Adapted from torch.nn.Embedding, note that we pad the vocabulary size to + make sure it is divisible by the number of model parallel GPUs. + + In order to support various loading methods, we ensure that LoRA-added + embeddings are always at the end of TP-sharded tensors. In other words, + we shard base embeddings and LoRA embeddings separately (both padded), + and place them in the same tensor. + In this example, we will have the original vocab size = 1010, + added vocab size = 16 and padding to 64. Therefore, the total + vocab size with padding will be 1088 (because we first pad 1010 to + 1024, add 16, and then pad to 1088). + Therefore, the tensor format looks like the following: + TP1, rank 0 (no sharding): + |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| + corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 | + + TP2, rank 0: + |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| + corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | + TP2, rank 1: + |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| + corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | + index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 | + + Args: + num_embeddings: vocabulary size. + embedding_dim: size of hidden state. + params_dtype: type of the parameters. + org_num_embeddings: original vocabulary size (without LoRA). + padding_size: padding size for the vocabulary. + quant_config: quant config for the layer + prefix: full name of the layer in the state dict + """ # noqa: E501 + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + params_dtype: torch.dtype | None = None, + org_num_embeddings: int | None = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + + # Keep the input dimensions. + tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.num_embeddings = num_embeddings + self.padding_size = padding_size + self.org_vocab_size = org_num_embeddings or num_embeddings + num_added_embeddings = num_embeddings - self.org_vocab_size + self.org_vocab_size_padded = pad_vocab_size( + self.org_vocab_size, self.padding_size + ) + self.num_embeddings_padded = pad_vocab_size( + self.org_vocab_size_padded + num_added_embeddings, self.padding_size + ) + assert self.org_vocab_size_padded <= self.num_embeddings_padded + + self.shard_indices = self._get_indices( + self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, + tp_rank, + self.tp_size, + ) + self.embedding_dim = embedding_dim + + quant_method = None + if quant_config is not None: + quant_method = quant_config.get_quant_method(self, prefix=prefix) + if quant_method is None: + quant_method = UnquantizedEmbeddingMethod() + + # If we are making an embedding layer, then our quantization linear + # method must implement the embedding operation. If we are another + # layer type like ParallelLMHead, this is not important. + is_embedding_layer = type(self.__class__) is VocabParallelEmbedding + quant_method_implements_embedding = method_has_implemented_embedding( + type(quant_method) + ) + if is_embedding_layer and not quant_method_implements_embedding: + raise NotImplementedError( + f"The class {type(quant_method).__name__} must implement " + "the 'embedding' method, see UnquantizedEmbeddingMethod." + ) + + self.quant_method: QuantizeMethodBase = quant_method + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + # Divide the weight matrix along the vocaburaly dimension. + self.num_added_embeddings = self.num_embeddings - self.org_vocab_size + self.num_embeddings_per_partition = divide( + self.num_embeddings_padded, self.tp_size + ) + assert ( + self.shard_indices.num_elements_padded + == self.num_embeddings_per_partition + ) + self.num_org_embeddings_per_partition = ( + self.shard_indices.org_vocab_end_index + - self.shard_indices.org_vocab_start_index + ) + self.num_added_embeddings_per_partition = ( + self.shard_indices.added_vocab_end_index + - self.shard_indices.added_vocab_start_index + ) + + self.quant_method.create_weights( + self, + self.embedding_dim, + [self.num_embeddings_per_partition], + self.embedding_dim, + self.num_embeddings_padded, + params_dtype=params_dtype, + weight_loader=self.weight_loader, + ) + + @classmethod + def _get_indices( + cls, + vocab_size_padded: int, + org_vocab_size_padded: int, + vocab_size: int, + org_vocab_size: int, + tp_rank: int, + tp_size: int, + ) -> VocabParallelEmbeddingShardIndices: + """Get start and end indices for vocab parallel embedding, following the layout outlined in + the class docstring, based on the given tp_rank and tp_size.""" + num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded + padded_org_vocab_start_index, padded_org_vocab_end_index = ( + vocab_range_from_global_vocab_size( + org_vocab_size_padded, tp_rank, tp_size + ) + ) + padded_added_vocab_start_index, padded_added_vocab_end_index = ( + vocab_range_from_global_vocab_size( + num_added_embeddings_padded, + tp_rank, + tp_size, + offset=org_vocab_size, + ) + ) + # remove padding + org_vocab_start_index = min( + padded_org_vocab_start_index, org_vocab_size + ) + org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size) + added_vocab_start_index = min( + padded_added_vocab_start_index, vocab_size + ) + added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size) + return VocabParallelEmbeddingShardIndices( + padded_org_vocab_start_index, + padded_org_vocab_end_index, + padded_added_vocab_start_index, + padded_added_vocab_end_index, + org_vocab_start_index, + org_vocab_end_index, + added_vocab_start_index, + added_vocab_end_index, + ) + + def get_sharded_to_full_mapping(self) -> list[int] | None: + """Get a mapping that can be used to reindex the gathered logits for sampling. + + During sampling, we gather logits from all ranks. The relationship of index->token_id will + follow the same format as outlined in the class docstring. However, after the gather, we + want to reindex the final logits tensor to map index->token_id one-to-one (the index is + always equal the token_id it corresponds to). The indices returned by this method allow us + to do that. + """ + if self.tp_size < 2: + return None + + base_embeddings: list[int] = [] + added_embeddings: list[int] = [] + padding: list[int] = [] + for tp_rank in range(self.tp_size): + shard_indices = self._get_indices( + self.num_embeddings_padded, + self.org_vocab_size_padded, + self.num_embeddings, + self.org_vocab_size, + tp_rank, + self.tp_size, + ) + range_start = self.num_embeddings_per_partition * tp_rank + range_end = self.num_embeddings_per_partition * (tp_rank + 1) + base_embeddings.extend( + range(range_start, range_start + shard_indices.num_org_elements) + ) + padding.extend( + range( + range_start + shard_indices.num_org_elements, + range_start + shard_indices.num_org_elements_padded, + ) + ) + added_embeddings.extend( + range( + range_start + shard_indices.num_org_elements_padded, + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements, + ) + ) + padding.extend( + range( + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements, + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded, + ) + ) + assert ( + range_start + + shard_indices.num_org_elements_padded + + shard_indices.num_added_elements_padded + == range_end + ) + ret = base_embeddings + added_embeddings + padding + assert len(ret) == self.num_embeddings_padded + return ret + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + output_dim = getattr(param, "output_dim", None) + packed_dim = getattr(param, "packed_dim", None) + + # If the parameter is a gguf weight, then load it directly. + if getattr(param, "is_gguf_weight_type", None): + param.data.copy_(loaded_weight) + param.weight_type = loaded_weight.item() + return + elif isinstance(param, UninitializedParameter): + shape = list(loaded_weight.shape) + if output_dim is not None: + shape[output_dim] = self.num_embeddings_per_partition + param.materialize(tuple(shape), dtype=loaded_weight.dtype) + + # If parameter does not have output dim, then it should + # be copied onto all gpus (e.g. g_idx for act_order gptq). + if output_dim is None: + assert param.data.shape == loaded_weight.shape + param.data.copy_(loaded_weight) + return + + # Shard indexes for loading the weight + start_idx = self.shard_indices.org_vocab_start_index + shard_size = self.shard_indices.org_vocab_end_index - start_idx + + # If param packed on the same dim we are sharding on, then + # need to adjust offsets of loaded weight by pack_factor. + if packed_dim is not None and packed_dim == output_dim: + packed_factor = ( + param.packed_factor + if isinstance(param, BasevLLMParameter) + else param.pack_factor + ) + assert loaded_weight.shape[output_dim] == ( + self.org_vocab_size // param.packed_factor + ) + start_idx = start_idx // packed_factor + shard_size = shard_size // packed_factor + else: + assert loaded_weight.shape[output_dim] == self.org_vocab_size + + # Copy the data. Select chunk corresponding to current shard. + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) + + param[: loaded_weight.shape[0]].data.copy_(loaded_weight) + param[loaded_weight.shape[0] :].data.fill_(0) + + def forward(self, input_): + if self.tp_size > 1: + # Build the mask. + masked_input, input_mask = get_masked_input_and_mask( + input_, + self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + self.shard_indices.num_org_vocab_padding, + self.shard_indices.added_vocab_start_index, + self.shard_indices.added_vocab_end_index, + ) + else: + masked_input = input_ + # Get the embeddings. + output_parallel = self.quant_method.embedding(self, masked_input.long()) + # Mask the output embedding. + if self.tp_size > 1: + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + # Reduce across all the model parallel GPUs. + output = tensor_model_parallel_all_reduce(output_parallel) + return output + + def extra_repr(self) -> str: + s = f"num_embeddings={self.num_embeddings_per_partition}" + s += f", embedding_dim={self.embedding_dim}" + s += f", org_vocab_size={self.org_vocab_size}" + s += f", num_embeddings_padded={self.num_embeddings_padded}" + s += f", tp_size={self.tp_size}" + return s diff --git a/worldcompass/fastvideo/logger.py b/worldcompass/fastvideo/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..2ad7546bc449619684d8ccebbe56ba02c23ba8f2 --- /dev/null +++ b/worldcompass/fastvideo/logger.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py +"""Logging configuration for fastvideo.""" + +import datetime +import json +import logging +import os +import sys +from functools import lru_cache, partial +from logging import Logger +from logging.config import dictConfig +from os import path +from types import MethodType +from typing import Any, cast + +import fastvideo.envs as envs + +FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING +FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH +FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL +FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX + +RED = "\033[91m" +GREEN = "\033[92m" +RESET = "\033[0;0m" + +_warned_local_main_process = False +_warned_main_process = False + +_FORMAT = ( + f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s " + "[%(filename)s:%(lineno)d] %(message)s" +) +_DATE_FORMAT = "%m-%d %H:%M:%S" + +DEFAULT_LOGGING_CONFIG = { + "formatters": { + "fastvideo": { + "class": "fastvideo.logging_utils.NewLineFormatter", + "datefmt": _DATE_FORMAT, + "format": _FORMAT, + }, + }, + "handlers": { + "fastvideo": { + "class": "logging.StreamHandler", + "formatter": "fastvideo", + "level": FASTVIDEO_LOGGING_LEVEL, + "stream": "ext://sys.stdout", + }, + }, + "loggers": { + "fastvideo": { + "handlers": ["fastvideo"], + "level": "DEBUG", + "propagate": False, + }, + }, + "root": { + "handlers": ["fastvideo"], + "level": "DEBUG", + }, + "version": 1, + "disable_existing_loggers": False, +} + + +@lru_cache +def _print_info_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.info(msg, stacklevel=2) + + +@lru_cache +def _print_warning_once(logger: Logger, msg: str) -> None: + # Set the stacklevel to 2 to print the original caller's line info + logger.warning(msg, stacklevel=2) + + +# TODO(will): add env variable to control this process-aware logging behavior +def _info( + logger: Logger, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, +) -> None: + """Process-aware INFO level logging function. + + This function controls logging behavior based on the process rank, allowing for + selective logging from specific processes in a distributed environment. + + Args: + logger: The logger instance to use for logging + msg: The message format string to log + *args: Format string arguments + main_process_only: If True, only log if this is the global main process (RANK=0) + local_main_process_only: If True, only log if this is the local main process (LOCAL_RANK=0) + **kwargs: Additional keyword arguments to pass to the logger.log method + - stacklevel: Defaults to 2 to show the original caller's location + + Note: + - When both main_process_only and local_main_process_only are True, + the message will be logged only if both conditions are met + - When both are False, the message will be logged from all processes + - By default, only logs from processes with LOCAL_RANK=0 + """ + try: + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + except Exception: + local_rank = 0 + rank = 0 + + is_main_process = rank == 0 + is_local_main_process = local_rank == 0 + + if (main_process_only and is_main_process) or ( + local_main_process_only and is_local_main_process + ): + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + global _warned_local_main_process, _warned_main_process + + if not _warned_local_main_process and local_main_process_only: + logger.warning( + "%s is_local_main_process is set to True, logging only from the local main process.%s", + GREEN, + RESET, + ) + _warned_local_main_process = True + if not _warned_main_process and main_process_only: + logger.warning( + "%s is_main_process_only is set to True, logging only from the main process.%s", + GREEN, + RESET, + ) + _warned_main_process = True + + if not main_process_only and not local_main_process_only: + logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + + +class _FastvideoLogger(Logger): + """ + Note: + This class is just to provide type information. + We actually patch the methods directly on the :class:`logging.Logger` + instance to avoid conflicting with other libraries such as + `intel_extension_for_pytorch.utils._logger`. + """ + + def info_once(self, msg: str) -> None: + """As :meth:`info`, but subsequent calls with the same message are silently dropped.""" + _print_info_once(self, msg) + + def warning_once(self, msg: str) -> None: + """As :meth:`warning`, but subsequent calls with the same message are silently dropped.""" + _print_warning_once(self, msg) + + def info( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, + ) -> None: + _info( + self, + msg, + *args, + main_process_only=main_process_only, + local_main_process_only=local_main_process_only, + **kwargs, + ) + + +def _configure_fastvideo_root_logger() -> None: + logging_config = dict[str, Any]() + + if not FASTVIDEO_CONFIGURE_LOGGING and FASTVIDEO_LOGGING_CONFIG_PATH: + raise RuntimeError( + "FASTVIDEO_CONFIGURE_LOGGING evaluated to false, but " + "FASTVIDEO_LOGGING_CONFIG_PATH was given. FASTVIDEO_LOGGING_CONFIG_PATH " + "implies FASTVIDEO_CONFIGURE_LOGGING. Please enable " + "FASTVIDEO_CONFIGURE_LOGGING or unset FASTVIDEO_LOGGING_CONFIG_PATH." + ) + + if FASTVIDEO_CONFIGURE_LOGGING: + logging_config = DEFAULT_LOGGING_CONFIG + + if FASTVIDEO_LOGGING_CONFIG_PATH: + if not path.exists(FASTVIDEO_LOGGING_CONFIG_PATH): + raise RuntimeError( + "Could not load logging config. File does not exist: %s", + FASTVIDEO_LOGGING_CONFIG_PATH, + ) + with open(FASTVIDEO_LOGGING_CONFIG_PATH, encoding="utf-8") as file: + custom_config = json.loads(file.read()) + + if not isinstance(custom_config, dict): + raise ValueError( + "Invalid logging config. Expected Dict, got %s.", + type(custom_config).__name__, + ) + logging_config = custom_config + + for formatter in logging_config.get("formatters", {}).values(): + # This provides backwards compatibility after #10134. + if formatter.get("class") == "fastvideo.logging.NewLineFormatter": + formatter["class"] = "fastvideo.logging_utils.NewLineFormatter" + + if logging_config: + dictConfig(logging_config) + + +def init_logger(name: str) -> _FastvideoLogger: + """The main purpose of this function is to ensure that loggers are retrieved in such a way that + we can be sure the root fastvideo logger has already been configured.""" + + logger = logging.getLogger(name) + + methods_to_patch = { + "info_once": _print_info_once, + "warning_once": _print_warning_once, + "info": _info, + } + + for method_name, method in methods_to_patch.items(): + setattr(logger, method_name, MethodType(method, logger)) # type: ignore[arg-type] + + return cast(_FastvideoLogger, logger) + + +# The root logger is initialized when the module is imported. +# This is thread-safe as the module is only imported once, +# guaranteed by the Python GIL. +_configure_fastvideo_root_logger() + +logger = init_logger(__name__) + + +def _trace_calls(log_path, root_dir, frame, event, arg=None): + if event in ["call", "return"]: + # Extract the filename, line number, function name, and the code object + filename = frame.f_code.co_filename + lineno = frame.f_lineno + func_name = frame.f_code.co_name + if not filename.startswith(root_dir): + # only log the functions in the fastvideo root_dir + return + # Log every function call or return + try: + last_frame = frame.f_back + if last_frame is not None: + last_filename = last_frame.f_code.co_filename + last_lineno = last_frame.f_lineno + last_func_name = last_frame.f_code.co_name + else: + # initial frame + last_filename = "" + last_lineno = 0 + last_func_name = "" + with open(log_path, "a") as f: + ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") + if event == "call": + f.write( + f"{ts} Call to" + f" {func_name} in {filename}:{lineno}" + f" from {last_func_name} in {last_filename}:" + f"{last_lineno}\n" + ) + else: + f.write( + f"{ts} Return from" + f" {func_name} in {filename}:{lineno}" + f" to {last_func_name} in {last_filename}:" + f"{last_lineno}\n" + ) + except NameError: + # modules are deleted during shutdown + pass + return partial(_trace_calls, log_path, root_dir) + + +def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): + """Enable tracing of every function call in code under `root_dir`. This is useful for debugging + hangs or crashes. `log_file_path` is the path to the log file. `root_dir` is the root directory + of the code to trace. If None, it is the fastvideo root directory. + + Note that this call is thread-level, any threads calling this function will have the trace + enabled. Other threads will not be affected. + """ + logger.warning( + "FASTVIDEO_TRACE_FUNCTION is enabled. It will record every" + " function executed by Python. This will slow down the code. It " + "is suggested to be used for debugging hang or crashes only." + ) + logger.info("Trace frame log is saved to %s", log_file_path) + if root_dir is None: + # by default, this is the fastvideo root directory + root_dir = os.path.dirname(os.path.dirname(__file__)) + sys.settrace(partial(_trace_calls, log_file_path, root_dir)) diff --git a/worldcompass/fastvideo/logging_utils/__init__.py b/worldcompass/fastvideo/logging_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba38f69e3ca3280c004f2285b34b5fd6ad1668c3 --- /dev/null +++ b/worldcompass/fastvideo/logging_utils/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.logging_utils.formatter import NewLineFormatter +from fastvideo.logging_utils.formatter import setup_for_distributed + +__all__ = [ + "NewLineFormatter", + "setup_for_distributed", +] diff --git a/worldcompass/fastvideo/logging_utils/formatter.py b/worldcompass/fastvideo/logging_utils/formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..78e96ce0be37e3fac42294a93767c03481d81b96 --- /dev/null +++ b/worldcompass/fastvideo/logging_utils/formatter.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logging_utils/formatter.py + +import logging + + +class NewLineFormatter(logging.Formatter): + """Adds logging prefix to newlines to align multi-line messages.""" + + def __init__(self, fmt, datefmt=None, style="%"): + logging.Formatter.__init__(self, fmt, datefmt, style) + + def format(self, record): + msg = logging.Formatter.format(self, record) + if record.message != "": + parts = msg.split(record.message) + msg = msg.replace("\n", "\r\n" + parts[0]) + return msg + + +def setup_for_distributed(is_master): + """This function disables printing when not in master process.""" + import builtins as __builtin__ + + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print diff --git a/worldcompass/fastvideo/models/__init__.py b/worldcompass/fastvideo/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/models/encoders/base.py b/worldcompass/fastvideo/models/encoders/base.py new file mode 100644 index 0000000000000000000000000000000000000000..0600b07ce3694290a8fc695166396efbfff8e30e --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/base.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +from abc import ABC, abstractmethod +from dataclasses import field + +import torch +from torch import nn + +from fastvideo.configs.models.encoders import ( + BaseEncoderOutput, + ImageEncoderConfig, + TextEncoderConfig, +) +from fastvideo.platforms import AttentionBackendEnum + + +class TextEncoder(nn.Module, ABC): + _fsdp_shard_conditions: list = field(default_factory=lambda: []) + _stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=list + ) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + TextEncoderConfig()._supported_attention_backends + ) + + def __init__(self, config: TextEncoderConfig) -> None: + super().__init__() + self.config = config + self._fsdp_shard_conditions = config._fsdp_shard_conditions + self._stacked_params_mapping = config.arch_config.stacked_params_mapping + if not self.supported_attention_backends: + raise ValueError( + f"Subclass {self.__class__.__name__} must define _supported_attention_backends" + ) + + @abstractmethod + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + pass + + @property + def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: + return self._supported_attention_backends + + +class ImageEncoder(nn.Module, ABC): + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + ImageEncoderConfig()._supported_attention_backends + ) + + def __init__(self, config: ImageEncoderConfig) -> None: + super().__init__() + self.config = config + if not self.supported_attention_backends: + raise ValueError( + f"Subclass {self.__class__.__name__} must define _supported_attention_backends" + ) + + @abstractmethod + def forward( + self, pixel_values: torch.Tensor, **kwargs + ) -> BaseEncoderOutput: + pass + + @property + def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: + return self._supported_attention_backends diff --git a/worldcompass/fastvideo/models/encoders/bert.py b/worldcompass/fastvideo/models/encoders/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..b52b9434933d90f9d995d35ccce391ca4de8be85 --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/bert.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +import os + +import torch +import torch.nn as nn +from transformers import BertModel, BertTokenizer + + +class HunyuanClip(nn.Module): + """ + Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py + hunyuan's clip used BertModel and BertTokenizer, so we copy it. + """ + + def __init__(self, model_dir, max_length=77): + super().__init__() + + self.max_length = max_length + self.tokenizer = BertTokenizer.from_pretrained( + os.path.join(model_dir, "tokenizer") + ) + self.text_encoder = BertModel.from_pretrained( + os.path.join(model_dir, "clip_text_encoder") + ) + + @torch.no_grad + def forward(self, prompts, with_mask=True): + self.device = next(self.text_encoder.parameters()).device + text_inputs = self.tokenizer( + prompts, + padding="max_length", + max_length=self.max_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + prompt_embeds = self.text_encoder( + text_inputs.input_ids.to(self.device), + attention_mask=( + text_inputs.attention_mask.to(self.device) + if with_mask + else None + ), + ) + return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output diff --git a/worldcompass/fastvideo/models/encoders/clip.py b/worldcompass/fastvideo/models/encoders/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..ee7ee58198b1ec337f6bd835a1ddbd163ce55f84 --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/clip.py @@ -0,0 +1,698 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/clip.py +# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py +"""Minimal implementation of CLIPVisionModel intended to be only used within a vision language +model.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn + +# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.encoders import ( + BaseEncoderOutput, + CLIPTextConfig, + CLIPVisionConfig, +) +from fastvideo.distributed import divide, get_tp_world_size +from fastvideo.layers.activation import get_act_fn +from fastvideo.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.logger import init_logger +from fastvideo.models.encoders.base import ImageEncoder, TextEncoder +from fastvideo.models.encoders.vision import resolve_visual_encoder_outputs + +# TODO: support quantization +# from vllm.model_executor.layers.quantization import QuantizationConfig +from fastvideo.models.loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + + +# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 # noqa +class CLIPVisionEmbeddings(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + assert self.image_size % self.patch_size == 0 + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding( + self.num_positions, self.embed_dim + ) + self.register_buffer( + "position_ids", + torch.arange(self.num_positions).expand((1, -1)), + persistent=False, + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size = pixel_values.shape[0] + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding( + pixel_values.to(dtype=target_dtype) + ) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.position_embedding(self.position_ids) + + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding( + config.max_position_embeddings, embed_dim + ) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False, + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + if input_ids is not None: + seq_length = input_ids.shape[-1] + elif inputs_embeds is not None: + seq_length = inputs_embeds.shape[-2] + else: + raise ValueError( + "Either input_ids or inputs_embeds must be provided." + ) + + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +class CLIPAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper.""" + + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + "embed_dim must be divisible by num_heads " + f"(got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.embed_dim, + head_size=self.head_dim, + total_num_heads=self.num_heads, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.out_proj = RowParallelLinear( + input_size=self.embed_dim, + output_size=self.embed_dim, + quant_config=quant_config, + prefix=f"{prefix}.out_proj", + ) + + self.tp_size = get_tp_world_size() + self.num_heads_per_partition = divide(self.num_heads, self.tp_size) + + self.attn = LocalAttention( + self.num_heads_per_partition, + self.head_dim, + self.num_heads_per_partition, + softmax_scale=self.scale, + causal=True, + supported_attention_backends=config._supported_attention_backends, + ) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return ( + tensor.view(bsz, seq_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .contiguous() + ) + + def forward( + self, + hidden_states: torch.Tensor, + ): + """Input shape: Batch x Time x Channel""" + + qkv_states, _ = self.qkv_proj(hidden_states) + query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) + # use flash_attn_func + query_states = query_states.reshape( + query_states.shape[0], + query_states.shape[1], + self.num_heads_per_partition, + self.head_dim, + ) + key_states = key_states.reshape( + key_states.shape[0], + key_states.shape[1], + self.num_heads_per_partition, + self.head_dim, + ) + value_states = value_states.reshape( + value_states.shape[0], + value_states.shape[1], + self.num_heads_per_partition, + self.head_dim, + ) + attn_output = self.attn(query_states, key_states, value_states) + attn_output = attn_output.reshape( + attn_output.shape[0], + attn_output.shape[1], + self.num_heads_per_partition * self.head_dim, + ) + attn_output, _ = self.out_proj(attn_output) + + return attn_output, None + + +class CLIPMLP(nn.Module): + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.activation_fn = get_act_fn(config.hidden_act) + self.fc1 = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc1", + ) + self.fc2 = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc2", + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states, _ = self.fc2(hidden_states) + + return hidden_states + + +class CLIPEncoderLayer(nn.Module): + def __init__( + self, + config: CLIPTextConfig | CLIPVisionConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = CLIPAttention( + config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.layer_norm1 = nn.LayerNorm( + config.hidden_size, eps=config.layer_norm_eps + ) + self.mlp = CLIPMLP( + config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + self.layer_norm2 = nn.LayerNorm( + config.hidden_size, eps=config.layer_norm_eps + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn(hidden_states=hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class CLIPEncoder(nn.Module): + """Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each + layer is a [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__( + self, + config: CLIPVisionConfig | CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + + if num_hidden_layers_override is None: + num_hidden_layers = config.num_hidden_layers + else: + num_hidden_layers = num_hidden_layers_override + self.layers = nn.ModuleList( + [ + CLIPEncoderLayer( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_idx}", + ) + for layer_idx in range(num_hidden_layers) + ] + ) + + def forward( + self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool + ) -> torch.Tensor | list[torch.Tensor]: + hidden_states_pool = [inputs_embeds] + hidden_states = inputs_embeds + + for encoder_layer in self.layers: + hidden_states = encoder_layer(hidden_states) + if return_all_hidden_states: + hidden_states_pool.append(hidden_states) + # If we have multiple feature sample layers, we return all hidden + # states in order and grab the ones we need by index. + if return_all_hidden_states: + return hidden_states_pool + return [hidden_states] + + +class CLIPTextTransformer(nn.Module): + def __init__( + self, + config: CLIPTextConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + prefix: str = "", + ): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPTextEmbeddings(config) + + self.encoder = CLIPEncoder( + config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=prefix, + ) + + self.final_layer_norm = nn.LayerNorm( + embed_dim, eps=config.layer_norm_eps + ) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + ) -> BaseEncoderOutput: + r""" + Returns: + + """ + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings( + input_ids=input_ids, position_ids=position_ids + ) + + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + # causal_attention_mask = _create_4d_causal_attention_mask( + # input_shape, hidden_states.dtype, device=hidden_states.device + # ) + + # # expand attention_mask + # if attention_mask is not None and not self._use_flash_attention_2: + # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") + # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + # attention_mask=attention_mask, + # causal_attention_mask=causal_attention_mask, + # output_attentions=output_attentions, + return_all_hidden_states=output_hidden_states, + # return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[-1] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange( + last_hidden_state.shape[0], device=last_hidden_state.device + ), + input_ids.to( + dtype=torch.int, device=last_hidden_state.device + ).argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange( + last_hidden_state.shape[0], device=last_hidden_state.device + ), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + ( + input_ids.to( + dtype=torch.int, device=last_hidden_state.device + ) + == self.eos_token_id + ) + .int() + .argmax(dim=-1), + ] + + return BaseEncoderOutput( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs, + # attentions=encoder_outputs.attentions, + ) + + +class CLIPTextModel(TextEncoder): + def __init__( + self, + config: CLIPTextConfig, + ) -> None: + super().__init__(config) + self.text_model = CLIPTextTransformer( + config=config, + quant_config=config.quant_config, + prefix=config.prefix, + ) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + + outputs: BaseEncoderOutput = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=output_hidden_states, + ) + return outputs + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + + # Define mapping for stacked parameters + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + # Handle q_proj, k_proj, v_proj -> qkv_proj mapping + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name in name: + # Replace the weight name with the parameter name + model_param_name = name.replace(weight_name, param_name) + + if model_param_name in params_dict: + param = params_dict[model_param_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(model_param_name) + break + else: + # Use default weight loader for all other parameters + if name in params_dict: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + + +class CLIPVisionTransformer(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + quant_config: QuantizationConfig | None = None, + num_hidden_layers_override: int | None = None, + require_post_norm: bool | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + + # NOTE: This typo of "layrnorm" is not fixed on purpose to match + # the original transformers code and name of the model weights. + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.encoder = CLIPEncoder( + config=config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=f"{prefix}.encoder", + ) + + num_hidden_layers = config.num_hidden_layers + if len(self.encoder.layers) > config.num_hidden_layers: + raise ValueError( + f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers." + ) + + # If possible, skip post_layernorm to conserve memory + if require_post_norm is None: + require_post_norm = len(self.encoder.layers) == num_hidden_layers + + if require_post_norm: + self.post_layernorm = nn.LayerNorm( + embed_dim, eps=config.layer_norm_eps + ) + else: + self.post_layernorm = None + + def forward( + self, + pixel_values: torch.Tensor, + feature_sample_layers: list[int] | None = None, + ) -> torch.Tensor: + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + return_all_hidden_states = feature_sample_layers is not None + + # Produces either the last layer output or all of the hidden states, + # depending on if we have feature_sample_layers or not + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + return_all_hidden_states=return_all_hidden_states, + ) + + if not return_all_hidden_states: + encoder_outputs = encoder_outputs[0] + + # Handle post-norm (if applicable) and stacks feature layers if needed + encoder_outputs = resolve_visual_encoder_outputs( + encoder_outputs, + feature_sample_layers, + self.post_layernorm, + self.config.num_hidden_layers, + ) + + return encoder_outputs + + +class CLIPVisionModel(ImageEncoder): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + + def __init__(self, config: CLIPVisionConfig) -> None: + super().__init__(config) + self.vision_model = CLIPVisionTransformer( + config=config, + quant_config=config.quant_config, + num_hidden_layers_override=config.num_hidden_layers_override, + require_post_norm=config.require_post_norm, + prefix=f"{config.prefix}.vision_model", + ) + + def forward( + self, + pixel_values: torch.Tensor, + feature_sample_layers: list[int] | None = None, + **kwargs, + ) -> BaseEncoderOutput: + last_hidden_state = self.vision_model( + pixel_values, feature_sample_layers + ) + return BaseEncoderOutput(last_hidden_state=last_hidden_state) + + @property + def device(self): + return next(self.parameters()).device + + # (TODO) Add prefix argument for filtering out weights to be loaded + # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986 + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + layer_count = len(self.vision_model.encoder.layers) + + for name, loaded_weight in weights: + if name.startswith("visual_projection"): + continue + # post_layernorm is not needed in CLIPVisionModel + if ( + name.startswith("vision_model.post_layernorm") + and self.vision_model.post_layernorm is None + ): + continue + + # omit layers when num_hidden_layers_override is set + if name.startswith("vision_model.encoder.layers"): + layer_idx = int(name.split(".")[3]) + if layer_idx >= layer_count: + continue + + for ( + param_name, + weight_name, + shard_id, + ) in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/worldcompass/fastvideo/models/encoders/llama.py b/worldcompass/fastvideo/models/encoders/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..af6c00cd39712d126312fb5ee21dad588dc0b870 --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/llama.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/llama.py + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only LLaMA model compatible with HuggingFace weights.""" + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +# from vllm.model_executor.layers.quantization import QuantizationConfig +from fastvideo.attention import LocalAttention + +# from ..utils import (extract_layer_index) +from fastvideo.configs.models.encoders import BaseEncoderOutput, LlamaConfig +from fastvideo.distributed import get_tp_world_size +from fastvideo.layers.activation import SiluAndMul +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.layers.rotary_embedding import get_rope +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) + + +class LlamaMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + # output_size=intermediate_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class LlamaAttention(nn.Module): + def __init__( + self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + bias_o_proj: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + # layer_idx = extract_layer_index(prefix) + self.hidden_size = hidden_size + tp_size = get_tp_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + # Phi models introduced a partial_rotary_factor parameter in the config + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + is_neox_style = True + is_gguf = ( + quant_config + and hasattr(quant_config, "get_name") + and quant_config.get_name() == "gguf" + ) + if is_gguf and config.model_type == "llama": + is_neox_style = False + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=int(rope_theta), + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + ) + + self.attn = LocalAttention( + self.num_heads, + self.head_dim, + self.num_kv_heads, + softmax_scale=self.scaling, + causal=True, + supported_attention_backends=config._supported_attention_backends, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + # attn_output = self.attn(q, k, v) + # use flash_attn_func + # TODO (Attn abstraction and backend) + # from flash_attn import flash_attn_func + # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) + batch_size = q.shape[0] + seq_len = q.shape[1] + q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + # import pdb; pdb.set_trace() + # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) + attn_output = self.attn(q, k, v) + attn_output = attn_output.reshape( + batch_size, seq_len, self.num_heads * self.head_dim + ) + + output, _ = self.o_proj(attn_output) + return output + + +class LlamaDecoderLayer(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + max_position_embeddings = getattr( + config, "max_position_embeddings", 8192 + ) + # Support abacusai/Smaug-72B-v0.1 with attention_bias + # Support internlm/internlm-7b with bias + attention_bias = getattr(config, "attention_bias", False) or getattr( + config, "bias", False + ) + bias_o_proj = attention_bias + # support internlm/internlm3-8b with qkv_bias + if hasattr(config, "qkv_bias"): + attention_bias = config.qkv_bias + + self.self_attn = LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr( + config, "num_key_value_heads", config.num_attention_heads + ), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + bias_o_proj=bias_o_proj, + prefix=f"{prefix}.self_attn", + ) + self.mlp = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual + ) + + hidden_states = self.self_attn( + positions=positions, hidden_states=hidden_states + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class LlamaModel(TextEncoder): + def __init__( + self, + config: LlamaConfig, + ): + super().__init__(config) + + self.config = config + self.quant_config = self.config.quant_config + if config.lora_config is not None: + max_loras = 1 + lora_vocab_size = 1 + if hasattr(config.lora_config, "max_loras"): + max_loras = config.lora_config.max_loras + if hasattr(config.lora_config, "lora_extra_vocab_size"): + lora_vocab_size = config.lora_config.lora_extra_vocab_size + lora_vocab = lora_vocab_size * max_loras + else: + lora_vocab = 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=config.quant_config, + ) + + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer( + config=config, + quant_config=config.quant_config, + prefix=f"{config.prefix}.layers.{i}", + ) + for i in range(config.num_hidden_layers) + ] + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + + if position_ids is None: + position_ids = torch.arange( + 0, hidden_states.shape[1], device=hidden_states.device + ).unsqueeze(0) + + all_hidden_states: tuple[Any, ...] | None = ( + () if output_hidden_states else None + ) + for layer in self.layers: + if all_hidden_states is not None: + # TODO + all_hidden_states += ( + (hidden_states,) + if residual is None + else (hidden_states + residual,) + ) + hidden_states, residual = layer( + position_ids, hidden_states, residual + ) + + hidden_states, _ = self.norm(hidden_states, residual) + + # add hidden states from the last decoder layer + if all_hidden_states is not None: + all_hidden_states += (hidden_states,) + + # TODO(will): maybe unify the output format with other models and use + # our own class + output = BaseEncoderOutput( + last_hidden_state=hidden_states, + # past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + # attentions=all_self_attns, + ) + + return output + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if ( + "rotary_emb.cos_cached" in name + or "rotary_emb.sin_cached" in name + ): + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + # if (self.quant_config is not None and + # (scale_name := self.quant_config.get_cache_scale(name))): + # # Loading kv cache quantization scales + # param = params_dict[scale_name] + # weight_loader = getattr(param, "weight_loader", + # default_weight_loader) + # loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else + # loaded_weight[0]) + # weight_loader(param, loaded_weight) + # loaded_params.add(scale_name) + # continue + if "scale" in name: + # Remapping the name of FP8 kv-scale. + kv_scale_name: str | None = maybe_remap_kv_scale_name( + name, params_dict + ) + if kv_scale_name is None: + continue + else: + name = kv_scale_name + for ( + param_name, + weight_name, + shard_id, + ) in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/worldcompass/fastvideo/models/encoders/stepllm.py b/worldcompass/fastvideo/models/encoders/stepllm.py new file mode 100644 index 0000000000000000000000000000000000000000..303b119157fbbe3d43eca964f170aa965f288c29 --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/stepllm.py @@ -0,0 +1,626 @@ +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import os +from functools import wraps + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.modeling_utils import PretrainedConfig, PreTrainedModel + +from fastvideo.models.dits.stepvideo import StepVideoRMSNorm + + +class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): + def __init__(self, device=None): + self.device = device + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if getattr(func, "__module__", None) == "torch.nn.init": + if "tensor" in kwargs: + return kwargs["tensor"] + else: + return args[0] + if ( + self.device is not None + and func in torch.utils._device._device_constructors() + and kwargs.get("device") is None + ): + kwargs["device"] = self.device + return func(*args, **kwargs) + + +def with_empty_init(func): + + @wraps(func) + def wrapper(*args, **kwargs): + with EmptyInitOnDevice("cpu"): + return func(*args, **kwargs) + + return wrapper + + +class LLaMaEmbedding(nn.Module): + """Language model embeddings. + + Arguments: + hidden_size: hidden size + vocab_size: vocabulary size + max_sequence_length: maximum size of sequence. This + is used for positional embedding + embedding_dropout_prob: dropout probability for embeddings + init_method: weight initialization method + num_tokentypes: size of the token-type embeddings. 0 value + will ignore this embedding + """ + + def __init__( + self, + cfg, + ): + super().__init__() + self.hidden_size = cfg.hidden_size + self.params_dtype = cfg.params_dtype + self.fp32_residual_connection = cfg.fp32_residual_connection + self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 + self.word_embeddings = torch.nn.Embedding( + cfg.padded_vocab_size, + self.hidden_size, + ) + self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) + + def forward(self, input_ids): + # Embeddings. + if self.embedding_weights_in_fp32: + self.word_embeddings = self.word_embeddings.to(torch.float32) + embeddings = self.word_embeddings(input_ids) + if self.embedding_weights_in_fp32: + embeddings = embeddings.to(self.params_dtype) + self.word_embeddings = self.word_embeddings.to(self.params_dtype) + + # Data format change to avoid explicit transposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + + # Dropout. + embeddings = self.embedding_dropout(embeddings) + + return embeddings + + +class StepChatTokenizer: + """Step Chat Tokenizer.""" + + def __init__( + self, + model_file, + name="StepChatTokenizer", + bot_token="<|BOT|>", # Begin of Turn + eot_token="<|EOT|>", # End of Turn + call_start_token="<|CALL_START|>", # Call Start + call_end_token="<|CALL_END|>", # Call End + think_start_token="<|THINK_START|>", # Think Start + think_end_token="<|THINK_END|>", # Think End + mask_start_token="<|MASK_1e69f|>", # Mask start + mask_end_token="<|UNMASK_1e69f|>", # Mask end + ): + import sentencepiece + + self._tokenizer = sentencepiece.SentencePieceProcessor( + model_file=model_file + ) + + self._vocab = {} + self._inv_vocab = {} + + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + for idx in range(self._tokenizer.get_piece_size()): + text = self._tokenizer.id_to_piece(idx) + self._inv_vocab[idx] = text + self._vocab[text] = idx + + if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown( + idx + ): + self._special_tokens[text] = idx + self._inv_special_tokens[idx] = text + + self._unk_id = self._tokenizer.unk_id() + self._bos_id = self._tokenizer.bos_id() + self._eos_id = self._tokenizer.eos_id() + + for token in [ + bot_token, + eot_token, + call_start_token, + call_end_token, + think_start_token, + think_end_token, + ]: + assert ( + token in self._vocab + ), f"Token '{token}' not found in tokenizer" + assert ( + token in self._special_tokens + ), f"Token '{token}' is not a special token" + + for token in [mask_start_token, mask_end_token]: + assert ( + token in self._vocab + ), f"Token '{token}' not found in tokenizer" + + self._bot_id = self._tokenizer.piece_to_id(bot_token) + self._eot_id = self._tokenizer.piece_to_id(eot_token) + self._call_start_id = self._tokenizer.piece_to_id(call_start_token) + self._call_end_id = self._tokenizer.piece_to_id(call_end_token) + self._think_start_id = self._tokenizer.piece_to_id(think_start_token) + self._think_end_id = self._tokenizer.piece_to_id(think_end_token) + self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) + self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) + + self._underline_id = self._tokenizer.piece_to_id("\u2581") + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def vocab_size(self): + return self._tokenizer.vocab_size() + + def tokenize(self, text: str) -> list[int]: + return self._tokenizer.encode_as_ids(text) + + def detokenize(self, token_ids: list[int]) -> str: + return self._tokenizer.decode_ids(token_ids) + + +class Tokens: + def __init__( + self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len + ) -> None: + self.input_ids = input_ids + self.attention_mask = attention_mask + self.cu_input_ids = cu_input_ids + self.cu_seqlens = cu_seqlens + self.max_seq_len = max_seq_len + + def to(self, device): + self.input_ids = self.input_ids.to(device) + self.attention_mask = self.attention_mask.to(device) + self.cu_input_ids = self.cu_input_ids.to(device) + self.cu_seqlens = self.cu_seqlens.to(device) + return self + + +class Wrapped_StepChatTokenizer(StepChatTokenizer): + def __call__( + self, + text, + max_length=320, + padding="max_length", + truncation=True, + return_tensors="pt", + ): + # [bos, ..., eos, pad, pad, ..., pad] + self.BOS = 1 + self.EOS = 2 + self.PAD = 2 + out_tokens = [] + attn_mask = [] + if len(text) == 0: + part_tokens = [self.BOS] + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + else: + for part in text: + part_tokens = self.tokenize(part) + part_tokens = part_tokens[ + : (max_length - 2) + ] # leave 2 space for bos and eos + part_tokens = [self.BOS] + part_tokens + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append( + [1] * valid_size + [0] * (max_length - valid_size) + ) + + out_tokens = torch.tensor(out_tokens, dtype=torch.long) + attn_mask = torch.tensor(attn_mask, dtype=torch.long) + + # padding y based on tp size + padded_len = 0 + padded_flag = False + if padded_len > 0: + padded_flag = True + if padded_flag: + pad_tokens = torch.tensor( + [[self.PAD] * max_length], device=out_tokens.device + ) + pad_attn_mask = torch.tensor( + [[1] * padded_len + [0] * (max_length - padded_len)], + device=attn_mask.device, + ) + out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) + attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) + + # cu_seqlens + cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) + seqlen = attn_mask.sum(dim=1).tolist() + cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to( + device=out_tokens.device, dtype=torch.int32 + ) + max_seq_len = max(seqlen) + return Tokens( + out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len + ) + + +def flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=True, + return_attn_probs=False, + tp_group_rank=0, + tp_group_size=1, +): + softmax_scale = ( + q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale + ) + return torch.ops.Optimus.fwd( + q, + k, + v, + None, + dropout_p, + softmax_scale, + causal, + return_attn_probs, + None, + tp_group_rank, + tp_group_size, + )[0] + + +class FlashSelfAttention(torch.nn.Module): + def __init__( + self, + attention_dropout=0.0, + ): + super().__init__() + self.dropout_p = attention_dropout + + def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): + if cu_seqlens is None: + output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) + else: + raise ValueError("cu_seqlens is not supported!") + + return output + + +def safediv(n, d): + q, r = divmod(n, d) + assert r == 0 + return q + + +class MultiQueryAttention(nn.Module): + def __init__(self, cfg, layer_id=None): + super().__init__() + + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.max_seq_len = cfg.seq_length + self.use_flash_attention = cfg.use_flash_attn + assert self.use_flash_attention, "FlashAttention is required!" + + self.n_groups = cfg.num_attention_groups + self.tp_size = 1 + self.n_local_heads = cfg.num_attention_heads + self.n_local_groups = self.n_groups + + self.wqkv = nn.Linear( + cfg.hidden_size, + cfg.hidden_size + self.head_dim * 2 * self.n_groups, + bias=False, + ) + self.wo = nn.Linear( + cfg.hidden_size, + cfg.hidden_size, + bias=False, + ) + + # assert self.use_flash_attention, 'non-Flash attention not supported yet.' + self.core_attention = FlashSelfAttention( + attention_dropout=cfg.attention_dropout + ) + # self.core_attention = LocalAttention( + # num_heads = self.n_local_heads, + # head_size = self.head_dim, + # # num_kv_heads = self.n_local_groups, + # casual = True, + # supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA], # RIVER TODO + # ) + self.layer_id = layer_id + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + seqlen, bsz, dim = x.shape + xqkv = self.wqkv(x) + + xq, xkv = torch.split( + xqkv, + ( + dim // self.tp_size, + self.head_dim * 2 * self.n_groups // self.tp_size, + ), + dim=-1, + ) + + # gather on 1st dimension + xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) + xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) + xk, xv = xkv.chunk(2, -1) + + # rotary embedding + flash attn + xq = rearrange(xq, "s b h d -> b s h d") + xk = rearrange(xk, "s b h d -> b s h d") + xv = rearrange(xv, "s b h d -> b s h d") + + # q_per_kv = self.n_local_heads // self.n_local_groups + # if q_per_kv > 1: + # b, s, h, d = xk.size() + # if h == 1: + # xk = xk.expand(b, s, q_per_kv, d) + # xv = xv.expand(b, s, q_per_kv, d) + # else: + # ''' To cover the cases where h > 1, we have + # the following implementation, which is equivalent to: + # xk = xk.repeat_interleave(q_per_kv, dim=-2) + # xv = xv.repeat_interleave(q_per_kv, dim=-2) + # but can avoid calling aten::item() that involves cpu. + # ''' + # idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() + # xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + # xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + if self.use_flash_attention: + output = self.core_attention(xq, xk, xv) + # reduce-scatter only support first dimension now + output = rearrange(output, "b s h d -> s b (h d)").contiguous() + else: + xq, xk, xv = [ + rearrange(x, "b s ... -> s b ...").contiguous() + for x in (xq, xk, xv) + ] + output = self.core_attention(xq, xk, xv) # , mask) + output = self.wo(output) + return output + + +class FeedForward(nn.Module): + def __init__( + self, + cfg, + dim: int, + hidden_dim: int, + layer_id: int, + multiple_of: int = 256, + ): + super().__init__() + + hidden_dim = multiple_of * ( + (hidden_dim + multiple_of - 1) // multiple_of + ) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.swiglu = swiglu + + self.w1 = nn.Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=False, + ) + + def forward(self, x): + x = self.swiglu(self.w1(x)) + output = self.w2(x) + return output + + +class TransformerBlock(nn.Module): + def __init__(self, cfg, layer_id: int): + super().__init__() + + self.n_heads = cfg.num_attention_heads + self.dim = cfg.hidden_size + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.attention = MultiQueryAttention( + cfg, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + cfg, + dim=cfg.hidden_size, + hidden_dim=cfg.ffn_hidden_size, + layer_id=layer_id, + ) + self.layer_id = layer_id + self.attention_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + self.ffn_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + residual = self.attention.forward( + self.attention_norm(x), mask, cu_seqlens, max_seq_len + ) + h = x + residual + ffn_res = self.feed_forward.forward(self.ffn_norm(h)) + out = h + ffn_res + return out + + +class Transformer(nn.Module): + def __init__( + self, + config, + max_seq_size=8192, + ): + super().__init__() + self.num_layers = config.num_layers + self.layers = self._build_layers(config) + + def _build_layers(self, config): + layers = torch.nn.ModuleList() + for layer_id in range(self.num_layers): + layers.append( + TransformerBlock( + config, + layer_id=layer_id + 1, + ) + ) + return layers + + def forward( + self, + hidden_states, + attention_mask, + cu_seqlens=None, + max_seq_len=None, + ): + + if max_seq_len is not None and not isinstance( + max_seq_len, torch.Tensor + ): + max_seq_len = torch.tensor( + max_seq_len, dtype=torch.int32, device="cpu" + ) + + for lid, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + attention_mask, + cu_seqlens, + max_seq_len, + ) + return hidden_states + + +class Step1Model(PreTrainedModel): + config_class = PretrainedConfig + + @with_empty_init + def __init__( + self, + config, + ): + super().__init__(config) + self.tok_embeddings = LLaMaEmbedding(config) + self.transformer = Transformer(config) + + def forward( + self, + input_ids=None, + attention_mask=None, + ): + + hidden_states = self.tok_embeddings(input_ids) + + hidden_states = self.transformer( + hidden_states, + attention_mask, + ) + return hidden_states + + +class STEP1TextEncoder(torch.nn.Module): + def __init__(self, model_dir, max_length=320): + super().__init__() + self.max_length = max_length + self.text_tokenizer = Wrapped_StepChatTokenizer( + os.path.join(model_dir, "step1_chat_tokenizer.model") + ) + text_encoder = Step1Model.from_pretrained(model_dir) + self.text_encoder = text_encoder.eval().to(torch.bfloat16) + + @torch.no_grad + def forward(self, prompts, with_mask=True, max_length=None): + self.device = next(self.text_encoder.parameters()).device + + with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): + if type(prompts) is str: + prompts = [prompts] + txt_tokens = self.text_tokenizer( + prompts, + max_length=max_length or self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + y = self.text_encoder( + txt_tokens.input_ids.to(self.device), + attention_mask=( + txt_tokens.attention_mask.to(self.device) + if with_mask + else None + ), + ) + y_mask = txt_tokens.attention_mask + return y.transpose(0, 1), y_mask diff --git a/worldcompass/fastvideo/models/encoders/vision.py b/worldcompass/fastvideo/models/encoders/vision.py new file mode 100644 index 0000000000000000000000000000000000000000..da33450613ee90e2330bd33e5e18e0845d54b6cc --- /dev/null +++ b/worldcompass/fastvideo/models/encoders/vision.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py + +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +import torch +from transformers import PretrainedConfig + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +_C = TypeVar("_C", bound=PretrainedConfig) + + +class VisionEncoderInfo(ABC, Generic[_C]): + def __init__(self, vision_config: _C) -> None: + super().__init__() + + self.vision_config = vision_config + + @abstractmethod + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + raise NotImplementedError + + @abstractmethod + def get_max_image_tokens(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_image_size(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_patch_size(self) -> int: + raise NotImplementedError + + @abstractmethod + def get_patch_grid_length(self) -> int: + raise NotImplementedError + + +def resolve_visual_encoder_outputs( + encoder_outputs: torch.Tensor | list[torch.Tensor], + feature_sample_layers: list[int] | None, + post_layer_norm: torch.nn.LayerNorm | None, + max_possible_layers: int, +) -> torch.Tensor: + """Given the outputs a visual encoder module that may correspond to the output of the last + layer, or a list of hidden states to be stacked, handle post normalization and resolve it into a + single output tensor. + + Args: + encoder_outputs: Output of encoder's last layer or all hidden states. + feature_sample_layers: Optional layer indices to grab from the encoder + outputs; if provided, encoder outputs must be a list. + post_layer_norm: Post norm to apply to the output of the encoder. + max_possible_layers: Total layers in the fully loaded visual encoder. + """ + if feature_sample_layers is None: + if post_layer_norm is not None: + return post_layer_norm(encoder_outputs) + return encoder_outputs + + # Get the hidden states corresponding to the layer indices. + # Negative values are relative to the full visual encoder, + # so offset them depending on how many layers were loaded. + # NOTE: this assumes that encoder_outputs is a list containing + # the inputs to the visual encoder, followed by the hidden states + # of each layer. + num_loaded_layers = len(encoder_outputs) - 1 + offset = max_possible_layers - num_loaded_layers + hs_pool = [ + ( + encoder_outputs[layer_idx] + if layer_idx >= 0 + else encoder_outputs[layer_idx + offset] + ) + for layer_idx in feature_sample_layers + ] + + # Apply post-norm on the final hidden state if we are using it + uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1) + if post_layer_norm is not None and uses_last_layer: + hs_pool[-1] = post_layer_norm(encoder_outputs) + return torch.cat(hs_pool, dim=-1) diff --git a/worldcompass/fastvideo/models/hf_transformer_utils.py b/worldcompass/fastvideo/models/hf_transformer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bc0309acf9ff829113d48b594e8489e58afb7852 --- /dev/null +++ b/worldcompass/fastvideo/models/hf_transformer_utils.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py + +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for Huggingface Transformers.""" + +import contextlib +import json +import os +from pathlib import Path +from typing import Any + +from huggingface_hub import snapshot_download +from transformers import AutoConfig, PretrainedConfig +from transformers.models.auto.modeling_auto import ( + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, +) + +_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { + # ChatGLMConfig.model_type: ChatGLMConfig, + # DbrxConfig.model_type: DbrxConfig, + # ExaoneConfig.model_type: ExaoneConfig, + # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig, +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf(model_path: str): + if os.path.exists(model_path): + return model_path + + return snapshot_download( + model_path, allow_patterns=["*.json", "*.bin", "*.model"] + ) + + +def get_hf_config( + model: str, + trust_remote_code: bool, + revision: str | None = None, + model_override_args: dict | None = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + raise NotImplementedError("GGUF models are not supported.") + + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + if config.model_type in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[config.model_type] + config = config_class.from_pretrained(model, revision=revision) + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + config._name_or_path = model + if model_override_args: + config.update(model_override_args) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError( + f"Can't get gguf config for {config.model_type}." + ) + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +def get_diffusers_config( + model: str, +) -> dict[str, Any]: + """Gets a configuration for the given diffusers model. + + Args: + model: The model name or path. + fastvideo_args: Optional inference arguments to override in the config. + + Returns: + The loaded configuration. + """ + config_name = "config.json" + if "scheduler" in model: + config_name = "scheduler_config.json" + # Check if the model path exists + if os.path.exists(model): + config_file = os.path.join(model, config_name) + if os.path.exists(config_file): + try: + # Load the config directly from the file + with open(config_file) as f: + config_dict: dict[str, Any] = json.load(f) + if "_diffusers_version" in config_dict: + config_dict.pop("_diffusers_version") + # TODO(will): apply any overrides from inference args + return config_dict + except Exception as e: + raise RuntimeError( + f"Failed to load diffusers config from {config_file}: {e}" + ) from e + raise RuntimeError(f"Config file not found at {config_file}") + else: + raise RuntimeError(f"Diffusers config file not found at {model}") + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def attach_additional_stop_token_ids(tokenizer): + # Special handling for stop token <|eom_id|> generated by llama 3 tool use. + if "<|eom_id|>" in tokenizer.get_added_vocab(): + tokenizer.additional_stop_token_ids = set( + [tokenizer.get_added_vocab()["<|eom_id|>"]] + ) + else: + tokenizer.additional_stop_token_ids = None + + +def check_gguf_file(model: str | os.PathLike) -> bool: + """Check if the file is a GGUF model.""" + model = Path(model) + if not model.is_file(): + return False + elif model.suffix == ".gguf": + return True + + with open(model, "rb") as f: + header = f.read(4) + return header == b"GGUF" diff --git a/worldcompass/fastvideo/models/hyvideo/__init__.py b/worldcompass/fastvideo/models/hyvideo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..48dd2bdffb8e3ac13331c6cc6d3bb9966cd4321d --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/__init__.py @@ -0,0 +1,46 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import socket +from .commons import get_gpu_memory + +if "TOKENIZERS_PARALLELISM" not in os.environ: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +def find_free_port(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("localhost", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def __initialize_default_distributed_environment(): + if "RANK" not in os.environ: + os.environ["RANK"] = "0" + if "WORLD_SIZE" not in os.environ: + os.environ["WORLD_SIZE"] = "1" + if "LOCAL_RANK" not in os.environ: + os.environ["LOCAL_RANK"] = "0" + if "MASTER_ADDR" not in os.environ: + os.environ["MASTER_ADDR"] = "localhost" + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = str(find_free_port()) + + +__initialize_default_distributed_environment() diff --git a/worldcompass/fastvideo/models/hyvideo/commons/__init__.py b/worldcompass/fastvideo/models/hyvideo/commons/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..158a7c65dafd483d36ae00199f3ec4b4266bcffd --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/commons/__init__.py @@ -0,0 +1,272 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import torch +from itertools import repeat +from contextlib import contextmanager +from torch import nn +import collections.abc + + +def _ntuple(n): + """Create a function that converts input to n-tuple.""" + + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + + return parse + + +# Convenience functions for common tuple sizes +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +# Default generation pipeline configurations +PIPELINE_CONFIGS = { + "480p_t2v": { + "guidance_scale": 6.0, + "embedded_guidance_scale": None, + "flow_shift": 5.0, + }, + "480p_i2v": { + "guidance_scale": 6.0, + "embedded_guidance_scale": None, + "flow_shift": 5.0, + }, + "720p_t2v": { + "guidance_scale": 6.0, + "embedded_guidance_scale": None, + "flow_shift": 9.0, + }, + "720p_i2v": { + "guidance_scale": 6.0, + "embedded_guidance_scale": None, + "flow_shift": 7.0, + }, + "480p_t2v_distilled": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 5.0, + }, + "480p_i2v_distilled": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 5.0, + }, + "720p_t2v_distilled": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 9.0, + }, + "720p_i2v_distilled": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 7.0, + }, + "720p_t2v_distilled_sparse": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 7.0, + }, + "720p_i2v_distilled_sparse": { + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "flow_shift": 9.0, + }, +} + +# Default super-resolution pipeline configurations +SR_PIPELINE_CONFIGS = { + "720p_sr_distilled": { + "flow_shift": 2.0, + "base_resolution": "480p", + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "num_inference_steps": 6, + }, + "1080p_sr_distilled": { + "flow_shift": 2.0, + "base_resolution": "720p", + "guidance_scale": 1.0, + "embedded_guidance_scale": None, + "num_inference_steps": 8, + }, +} + +TRANSFORMER_VERSION_TO_SR_VERSION = { + "480p_t2v": "720p_sr_distilled", + "720p_t2v": "1080p_sr_distilled", + "480p_i2v": "720p_sr_distilled", + "720p_i2v": "1080p_sr_distilled", + "480p_t2v_distilled": "720p_sr_distilled", + "720p_t2v_distilled": "1080p_sr_distilled", + "480p_i2v_distilled": "720p_sr_distilled", + "720p_i2v_distilled": "1080p_sr_distilled", + "480p_t2v_distilled_sparse": "720p_sr_distilled", + "720p_t2v_distilled_sparse": "1080p_sr_distilled", + "480p_i2v_distilled_sparse": "720p_sr_distilled", + "720p_i2v_distilled_sparse": "1080p_sr_distilled", +} + + +def is_flash2_available(): + try: + from flash_attn import flash_attn_varlen_qkvpacked_func + + return True + except Exception: + return False + + +def is_flash3_available(): + try: + from flash_attn_interface import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, + ) # noqa: F401 + + return True + except Exception: + return False + + +def is_flash_available(): + return is_flash2_available() or is_flash3_available() + + +def is_sparse_attn_supported(): + return "nvidia h" in torch.cuda.get_device_properties(0).name.lower() + + +def is_sparse_attn_available(): + if not is_sparse_attn_supported(): + return False + try: + from flex_block_attn import flex_block_attn_func # noqa: F401 + + return True + except Exception: + return False + + +def maybe_fallback_attn_mode(attn_mode, infer_state=None, block_idx=None): + """Determine the final attention mode based on configuration and availability. + + Args: + attn_mode: Requested attention mode + infer_state: Inference configuration object (optional) + block_idx: Current block index (optional) + + Returns: + Final attention mode to use + """ + import warnings + + # Check for sageattn and flex-block-attn conflict + enable_sageattn = False + if infer_state is not None: + enable_sageattn = ( + infer_state.enable_sageattn + and block_idx in infer_state.sage_blocks_range + ) + + assert not (enable_sageattn and attn_mode == "flex-block-attn"), ( + "SageAttention cannot be used with flex-block-attn mode. " + "Please disable enable_sageattn or use a different attention mode." + ) + + # Use SageAttention if configured + if enable_sageattn: + attn_mode = "sageattn" + return attn_mode + + # Handle flash attention modes + if attn_mode == "flash": + if is_flash3_available(): + attn_mode = "flash3" + elif is_flash2_available(): + attn_mode = "flash2" + else: + warnings.warn( + "flash is not available. Falling back to torch attention." + ) + attn_mode = "torch" + elif attn_mode == "flash3": + if not is_flash3_available(): + warnings.warn( + "flash3 is not available. Falling back to torch attention." + ) + attn_mode = "torch" + elif attn_mode == "flash2": + if not is_flash2_available(): + warnings.warn( + "flash2 is not available. Falling back to torch attention." + ) + attn_mode = "torch" + if attn_mode in ("flex-block-attn"): + from fastvideo.models.hyvideo.commons import is_sparse_attn_available + + if not is_sparse_attn_available(): + raise ValueError( + f"{attn_mode} is not available for your GPU or flex-block-attn is not properly installed." + ) + return attn_mode + + +@contextmanager +def auto_offload_model(models, device, enabled=True): + from diffusers.hooks.group_offloading import _is_group_offload_enabled + + if enabled: + if isinstance(models, nn.Module): + models = [models] + for model in models: + if model is not None: + model.to(device) + yield + if enabled: + for model in models: + if model is not None: + model.to(torch.device("cpu")) + + +def get_gpu_memory(device=None): + if not torch.cuda.is_available(): + return 0 + device = device if device is not None else torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + if hasattr(torch.cuda, "get_per_process_memory_fraction"): + memory_fraction = torch.cuda.get_per_process_memory_fraction() + else: + memory_fraction = 1.0 + return props.total_memory * memory_fraction + + +def get_rank(): + return int(os.environ.get("RANK", "0")) diff --git a/worldcompass/fastvideo/models/hyvideo/commons/infer_state.py b/worldcompass/fastvideo/models/hyvideo/commons/infer_state.py new file mode 100644 index 0000000000000000000000000000000000000000..46aafc17a62cea3d15e7e1c7b61ea2e8fbfb182b --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/commons/infer_state.py @@ -0,0 +1,55 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional +from dataclasses import dataclass + + +@dataclass +class InferState: + enable_sageattn: bool = False # whether to use SageAttention + sage_blocks_range: Optional[range] = ( + None # block range to use SageAttention + ) + enable_torch_compile: bool = False # whether to use torch compile + + +__infer_state = None + + +def parse_range(value): + if "-" in value: + start, end = map(int, value.split("-")) + return list(range(start, end + 1)) + else: + return [int(x) for x in value.split(",")] + + +def initialize_infer_state(args): + global __infer_state + sage_blocks_range = parse_range(args.sage_blocks_range) + # Map CLI argument use_sageattn to internal enable_sageattn field + use_sageattn = getattr(args, "use_sageattn", False) + __infer_state = InferState( + enable_sageattn=use_sageattn, + sage_blocks_range=sage_blocks_range, + enable_torch_compile=args.enable_torch_compile, + ) + return __infer_state + + +def get_infer_state(): + return __infer_state diff --git a/worldcompass/fastvideo/models/hyvideo/commons/parallel_states.py b/worldcompass/fastvideo/models/hyvideo/commons/parallel_states.py new file mode 100644 index 0000000000000000000000000000000000000000..4905b2c78a9449a44779fcf501c065f73ce0dc5d --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/commons/parallel_states.py @@ -0,0 +1,88 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from dataclasses import dataclass + +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh + + +@dataclass +class ParallelDims: + sp: int = 1 + world_size: int = -1 + + def __post_init__(self): + if self.world_size == -1: + if dist.is_initialized(): + self.world_size = dist.get_world_size() + else: + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + self.build_mesh("cuda") + + def build_mesh(self, device_type): + assert ( + self.world_size % self.sp == 0 + ), "world_size must be divisible by sp" + mesh = init_device_mesh( + device_type, + [self.world_size // self.sp, self.sp], + mesh_dim_names=["dp", "sp"], + ) + self.world_mesh = mesh + return mesh + + @property + def sp_enabled(self): + return self.sp > 1 + + @property + def sp_group(self): + return self.world_mesh["sp"].get_group() + + @property + def sp_mesh(self): + return self.world_mesh["sp"] + + @property + def sp_rank(self): + if self.sp_enabled: + return self.world_mesh["sp"].get_local_rank() + else: + return dist.get_rank() + + @property + def dp_enabled(self): + return self.sp > 1 + + +__parallel_dims = None + + +def initialize_parallel_state( + sp: int = 1, +): + global __parallel_dims + __parallel_dims = ParallelDims(sp=sp) + return __parallel_dims + + +def get_parallel_state(): + if __parallel_dims is None: + # create default parallel states (without enabling any parallelism) + initialize_parallel_state() + return __parallel_dims diff --git a/worldcompass/fastvideo/models/hyvideo/models/__init__.py b/worldcompass/fastvideo/models/hyvideo/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6a429aae5c9bed641d117863ba11ac863103c860 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/__init__.py @@ -0,0 +1,15 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. diff --git a/worldcompass/fastvideo/models/hyvideo/models/autoencoders/__init__.py b/worldcompass/fastvideo/models/hyvideo/models/autoencoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6a429aae5c9bed641d117863ba11ac863103c860 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/autoencoders/__init__.py @@ -0,0 +1,15 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. diff --git a/worldcompass/fastvideo/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py b/worldcompass/fastvideo/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..8866abbfdd5ffaccc589f5d709fa39b351450d9b --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/autoencoders/hunyuanvideo_15_vae_w_cache.py @@ -0,0 +1,1472 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import loguru +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, nn +from torch.nn import Conv3d +from torch import distributed as dist + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.autoencoders.vae import ( + BaseOutput, + DiagonalGaussianDistribution, +) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin + +from fastvideo.models.hyvideo.commons.parallel_states import get_parallel_state + +CACHE_T = 2 + + +@dataclass +class DecoderOutput(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +def swish(x: Tensor, inplace=False) -> Tensor: + """Applies the swish activation function (SiLU) with optional inplace support.""" + return F.silu(x, inplace=inplace) + + +def forward_with_checkpointing(module, *inputs, use_checkpointing=False): + """Executes the forward pass of a module, optionally using gradient checkpointing. + + Gradient checkpointing trades compute for memory by not storing intermediate + activations during the forward pass, but re-computing them during the backward pass. + + Args: + module (nn.Module): The module to run. + *inputs: The inputs to the module. + use_checkpointing (bool): If True, wraps execution in torch.utils.checkpoint. + """ + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if use_checkpointing: + # If checkpointing is enabled, wrap the custom forward function. + # use_reentrant=False is generally recommended for newer PyTorch versions + # to avoid issues with backward hooks and distributed training. + return torch.utils.checkpoint.checkpoint( + create_custom_forward(module), *inputs, use_reentrant=False + ) + else: + # Standard forward pass without checkpointing overhead + return module(*inputs) + + +# Optimized implementation of CogVideoXSafeConv3d +# https://github.com/huggingface/diffusers/blob/c9ff360966327ace3faad3807dc871a4e5447501/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py#L38 +class PatchCausalConv3d(nn.Conv3d): + r"""Causal Conv3d with efficient patch processing for large tensors.""" + + def find_split_indices(self, seq_len, part_num): + ideal_interval = seq_len / part_num + possible_indices = list(range(0, seq_len, self.stride[0])) + selected_indices = [] + + for i in range(1, part_num): + closest = min( + possible_indices, + key=lambda x: abs(x - round(i * ideal_interval)), + ) + if closest not in selected_indices: + selected_indices.append(closest) + + merged_indices = [] + prev_idx = 0 + for idx in selected_indices: + if idx - prev_idx >= self.kernel_size[0]: + merged_indices.append(idx) + prev_idx = idx + + return merged_indices + + def forward(self, input): + T = input.shape[2] # input: NCTHW + memory_count = ( + torch.prod(torch.tensor(input.shape)).item() * 2 / 1024**3 + ) + + if T > self.kernel_size[0] and memory_count > 0.6: + kernel_size = self.kernel_size[0] + part_num = int(memory_count / 2) + 1 + split_indices = self.find_split_indices(T, part_num) + input_chunks = ( + torch.tensor_split(input, split_indices, dim=2) + if len(split_indices) > 0 + else [input] + ) + if kernel_size > 1: + input_chunks = [input_chunks[0]] + [ + torch.cat( + ( + input_chunks[i - 1][:, :, -kernel_size + 1 :], + input_chunks[i], + ), + dim=2, + ) + for i in range(1, len(input_chunks)) + ] + output_chunks = [] + for input_chunk in input_chunks: + output_chunks.append(super().forward(input_chunk)) + output = torch.cat(output_chunks, dim=2) + return output + else: + return super().forward(input) + + +class RMS_norm(nn.Module): + """Root Mean Square Layer Normalization for Channel-First or Last.""" + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return ( + F.normalize(x, dim=(1 if self.channel_first else -1)) + * self.scale + * self.gamma + + self.bias + ) + + +class CausalConv3d(nn.Module): + """Causal Conv3d with configurable padding for temporal axis.""" + + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode="replicate", + disable_causal=False, + enable_patch_conv=False, + **kwargs, + ): + super().__init__() + + self.pad_mode = pad_mode + if disable_causal: + padding = ( + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + ) + else: + padding = ( + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size - 1, + 0, + ) + self.time_causal_padding = padding + + if enable_patch_conv: + self.conv = PatchCausalConv3d( + chan_in, + chan_out, + kernel_size, + stride=stride, + dilation=dilation, + **kwargs, + ) + else: + self.conv = nn.Conv3d( + chan_in, + chan_out, + kernel_size, + stride=stride, + dilation=dilation, + **kwargs, + ) + + def forward(self, x, cache_x=None): + padding = list(self.time_causal_padding) + if ( + cache_x is not None and self.time_causal_padding[4] > 0 + ): # not the first frame + cache_x = cache_x.to(x.device) + x = torch.cat( + [cache_x, x], dim=2 + ) # concat on temporal dimension (concat one) + padding[4] -= cache_x.shape[ + 2 + ] # from [1,1,1,1,2,0] to [1,1,1,1,1,0] (padding one) + x = F.pad(x, padding, mode=self.pad_mode) + return self.conv(x) + + +def prepare_causal_attention_mask( + n_frame: int, n_hw: int, dtype, device, batch_size: int = None +): + """Prepare a causal attention mask for 3D videos. + + Args: + n_frame (int): Number of frames (temporal length). + n_hw (int): Product of height and width. + dtype: Desired mask dtype. + device: Device for the mask. + batch_size (int, optional): If set, expands for batch. + + Returns: + torch.Tensor: Causal attention mask. + """ + seq_len = n_frame * n_hw + mask = torch.full( + (seq_len, seq_len), float("-inf"), dtype=dtype, device=device + ) + for i in range(seq_len): + i_frame = i // n_hw + mask[i, : (i_frame + 1) * n_hw] = 0 + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class AttnBlock(nn.Module): + """Self-attention block for 3D video tensors.""" + + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = RMS_norm(in_channels, images=False) + + self.q = Conv3d(in_channels, in_channels, kernel_size=1) + self.k = Conv3d(in_channels, in_channels, kernel_size=1) + self.v = Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, f, h, w = q.shape + q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous() + k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous() + v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous() + attention_mask = prepare_causal_attention_mask( + f, h * w, h_.dtype, h_.device, batch_size=b + ) + h_ = nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attention_mask.unsqueeze(1) + ) + + return rearrange( + h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b + ) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + """ResNet-style block for 3D video tensors.""" + + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = RMS_norm(in_channels, images=False) + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = RMS_norm(out_channels, images=False) + self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3) + if self.in_channels != self.out_channels: + self.nin_shortcut = Conv3d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = x + h = self.norm1(h) + h = swish(h, inplace=True) + + # apply the feature cacheing mechanism + if feat_cache is not None: + # Retrieve the current layer index + idx = feat_idx[0] + + # Clone the last CACHE_T frames from the current input to store for the next step. + cache_x = h[:, :, -CACHE_T:, :, :].clone() + + # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames) + # and we have a previous cache, prepend the last frame of the previous cache. + # This ensures sufficient temporal context for the convolution kernel. + if ( + cache_x.shape[2] < 2 and feat_cache[idx] is not None + ): # actually this means chunk 1 after chunk 0 + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + + # Apply the convolution layer using the input 'h' and the *previous* cached state. + h = self.conv1(h, feat_cache[idx]) + + # Update the cache for this layer with the newly prepared context ('cache_x'). + feat_cache[idx] = cache_x + + # Increment the global layer index. + feat_idx[0] += 1 + else: + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv2(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + """Hierarchical downsampling with temporal/ spatial support.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + add_temporal_downsample: bool = True, + ): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + self.conv = CausalConv3d( + in_channels, out_channels // factor, kernel_size=3 + ) + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward(self, x, feat_cache=None, feat_idx=[0]): + r1 = 2 if self.add_temporal_downsample else 1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv( + x + ) # change the channel, ready for saptial or temporal downsample + + if self.add_temporal_downsample: + if x.shape[2] == 1: + h = rearrange( + h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2 + ) + h = torch.cat([h, h], dim=1) + # compute the shortcut part + shortcut = rearrange( + x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2 + ) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view( + B, h.shape[1], self.group_size // 2, T, H, W + ).mean(dim=2) + else: + h = rearrange( + h, + "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", + r1=r1, + r2=2, + r3=2, + ) + # compute the shortcut part + shortcut = rearrange( + x, + "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", + r1=r1, + r2=2, + r3=2, + ) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view( + B, h.shape[1], self.group_size, T, H, W + ).mean(dim=2) + else: + h = rearrange( + h, + "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", + r1=r1, + r2=2, + r3=2, + ) + shortcut = rearrange( + x, + "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", + r1=r1, + r2=2, + r3=2, + ) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view( + B, h.shape[1], self.group_size, T, H, W + ).mean(dim=2) + + return h + shortcut + + +class Upsample(nn.Module): + """Hierarchical upsampling with temporal/ spatial support.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + add_temporal_upsample: bool = True, + ): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = CausalConv3d( + in_channels, out_channels * factor, kernel_size=3 + ) + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + r1 = 2 if self.add_temporal_upsample else 1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv( + x + ) # change the channel, ready for saptial or temporal downsample + + if self.add_temporal_upsample: + if first_chunk: + h = rearrange( + h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2 + ) + h = h[:, : h.shape[1] // 2] + # compute the shortcut part + shortcut = rearrange( + x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2 + ) + shortcut = shortcut.repeat_interleave( + repeats=self.repeats // 2, dim=1 + ) + else: + h = rearrange( + h, + "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", + r1=r1, + r2=2, + r3=2, + ) + # compute the shortcut part + shortcut = rearrange( + x, + "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", + r1=r1, + r2=2, + r3=2, + ) + shortcut = shortcut.repeat_interleave( + repeats=self.repeats, dim=1 + ) + else: + h = rearrange( + h, + "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", + r1=r1, + r2=2, + r3=2, + ) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange( + shortcut, + "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", + r1=r1, + r2=2, + r3=2, + ) + + return h + shortcut + + +class Encoder(nn.Module): + """Hierarchical video encoder with temporal and spatial factorization.""" + + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + downsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[-1] % (2 * z_channels) == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv3d( + in_channels, block_out_channels[0], kernel_size=3 + ) + + self.down = nn.ModuleList() + block_in = block_out_channels[0] + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks): + block.append( + ResnetBlock(in_channels=block_in, out_channels=block_out) + ) + block_in = block_out + down = nn.Module() + down.block = block + + add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_downsample = add_spatial_downsample and bool( + i_level >= np.log2(ffactor_spatial // ffactor_temporal) + ) + if add_spatial_downsample or add_temporal_downsample: + assert i_level < len(block_out_channels) - 1 + block_out = ( + block_out_channels[i_level + 1] + if downsample_match_channel + else block_in + ) + down.downsample = Downsample( + block_in, block_out, add_temporal_downsample + ) + block_in = block_out + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, out_channels=block_in + ) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, out_channels=block_in + ) + + # end + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, 2 * z_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=[0]): + """Forward pass through the encoder.""" + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() # cache x before layers + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv_in( + x, feat_cache[idx] + ) # [B, 160, T, H, W] with padding + feat_cache[idx] = cache_x + feat_idx[0] += 1 # move to next cache + else: + h = self.conv_in(x) + + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks): + h = forward_with_checkpointing( + self.down[i_level].block[i_block], + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + if hasattr(self.down[i_level], "downsample"): + h = forward_with_checkpointing( + self.down[i_level].downsample, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + + # middle + h = forward_with_checkpointing( + self.mid.block_1, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + h = forward_with_checkpointing( + self.mid.attn_1, h, use_checkpointing=use_checkpointing + ) + h = forward_with_checkpointing( + self.mid.block_2, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + + # end + group_size = self.block_out_channels[-1] // (2 * self.z_channels) + shortcut = rearrange( + h, "b (c r) f h w -> b c r f h w", r=group_size + ).mean(dim=2) + h = self.norm_out(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + + h = self.conv_out(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv_out(h) + h += shortcut + return h + + +class Decoder(nn.Module): + """Hierarchical video decoder with upsampling factories.""" + + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: Tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + upsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[0] % z_channels == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + block_in = block_out_channels[0] + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, out_channels=block_in + ) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, out_channels=block_in + ) + + # upsampling + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append( + ResnetBlock(in_channels=block_in, out_channels=block_out) + ) + block_in = block_out + up = nn.Module() + up.block = block + + add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal)) + if add_spatial_upsample or add_temporal_upsample: + assert i_level < len(block_out_channels) - 1 + block_out = ( + block_out_channels[i_level + 1] + if upsample_match_channel + else block_in + ) + up.upsample = Upsample( + block_in, block_out, add_temporal_upsample + ) + block_in = block_out + self.up.append(up) + + # end + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, z, feat_cache=None, feat_idx=[0], first_chunk=False): + """Forward pass through the decoder.""" + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + repeats = self.block_out_channels[0] // (self.z_channels) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = z[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv_in(z, feat_cache[idx]) + z.repeat_interleave( + repeats=repeats, dim=1 + ) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + # z to block_in + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # middle + h = forward_with_checkpointing( + self.mid.block_1, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + h = forward_with_checkpointing( + self.mid.attn_1, h, use_checkpointing=use_checkpointing + ) + h = forward_with_checkpointing( + self.mid.block_2, + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing( + self.up[i_level].block[i_block], + h, + feat_cache, + feat_idx, + use_checkpointing=use_checkpointing, + ) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing( + self.up[i_level].upsample, + h, + feat_cache, + feat_idx, + first_chunk, + use_checkpointing=use_checkpointing, + ) + + # end + h = self.norm_out(h) + h = swish(h, inplace=True) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = h[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + h = self.conv_out(h, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv_out(h) + return h + + +class AutoencoderKLConv3D(ModelMixin, ConfigMixin): + """KL regularized 3D Conv VAE with advanced tiling and slicing strategies.""" + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: Tuple[int, ...], + layers_per_block: int, + ffactor_spatial: int, + ffactor_temporal: int, + sample_size: int, + sample_tsize: int, + scaling_factor: float = None, + shift_factor: Optional[float] = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + ): + super().__init__() + self.ffactor_spatial = ffactor_spatial + self.ffactor_temporal = ffactor_temporal + self.scaling_factor = scaling_factor + self.shift_factor = shift_factor + + self.encoder = Encoder( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + downsample_match_channel=downsample_match_channel, + ) + self.decoder = Decoder( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + upsample_match_channel=upsample_match_channel, + ) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // ffactor_spatial + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // ffactor_temporal + self.tile_overlap_factor = 0.25 + + self._tile_parallelism_enabled = False + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": ( + sum(isinstance(m, CausalConv3d) for m in self.decoder.modules()) + if self.decoder is not None + else 0 + ), + "encoder": ( + sum(isinstance(m, CausalConv3d) for m in self.encoder.modules()) + if self.encoder is not None + else 0 + ), + } + + def set_tile_sample_min_size( + self, sample_size: int, tile_overlap_factor: float = 0.2 + ): + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // self.ffactor_spatial + self.tile_overlap_factor = tile_overlap_factor + + assert ( + self.tile_latent_min_size * self.tile_overlap_factor + ).is_integer(), "self.tile_latent_min_size multiplied by tile_overlap_factor must be an integer" + + def _set_gradient_checkpointing(self, module, value=False): + """Enable or disable gradient checkpointing on encoder and decoder.""" + if isinstance(module, (Encoder, Decoder)): + module.gradient_checkpointing = value + + def enable_temporal_tiling(self, use_tiling: bool = True): + raise RuntimeError("Temporal tiling is not supported for this VAE.") + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + self.enable_spatial_tiling(use_tiling) + + def disable_tiling(self): + self.disable_spatial_tiling() + + def enable_slicing(self): + self.use_slicing = True + + def disable_slicing(self): + self.use_slicing = False + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b horizontally into a at blend_extent region.""" + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent + ) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b vertically into a at blend_extent region.""" + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent + ) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + """Blend tensor b temporally into a at blend_extent region.""" + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( + 1 - x / blend_extent + ) + b[:, :, x, :, :] * (x / blend_extent) + return b + + def spatial_tiled_encode(self, x: torch.Tensor): + """Tiled spatial encoding for large inputs via overlapping.""" + B, C, T, H, W = x.shape + overlap_size = int( + self.tile_sample_min_size * (1 - self.tile_overlap_factor) + ) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = x[ + :, + :, + :, + i : i + self.tile_sample_min_size, + j : j + self.tile_sample_min_size, + ] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + return moments + + def temporal_tiled_encode(self, x: torch.Tensor): + """Tiled temporal encoding for large video sequences.""" + raise RuntimeError("Temporal tiling is not supported for this VAE.") + B, C, T, H, W = x.shape + overlap_size = int( + self.tile_sample_min_tsize * (1 - self.tile_overlap_factor) + ) + blend_extent = int( + self.tile_latent_min_tsize * self.tile_overlap_factor + ) + t_limit = self.tile_latent_min_tsize - blend_extent + + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_tsize + 1, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_sample_min_size + or tile.shape[-2] > self.tile_sample_min_size + ): + tile = self.spatial_tiled_encode(tile) + else: + tile = self.encoder(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, : t_limit + 1, :, :]) + moments = torch.cat(result_row, dim=-3) + return moments + + def enable_tile_parallelism(self): + self._tile_parallelism_enabled = True + + def disable_tile_parallelism(self): + self._tile_parallelism_enabled = False + + def tile_parallel_spatial_tiled_decode(self, z: torch.Tensor): + B, C, T, H, W = z.shape + overlap_size = int( + self.tile_latent_min_size * (1 - self.tile_overlap_factor) + ) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + rank = get_parallel_state().sp_rank + world_size = get_parallel_state().sp + + num_rows = math.ceil(H / overlap_size) + num_cols = math.ceil(W / overlap_size) + total_tiles = num_rows * num_cols + tiles_per_rank = math.ceil(total_tiles / world_size) + + my_linear_indices = list(range(rank, total_tiles, world_size)) + decoded_tiles = [] + decoded_metas = [] + H_out_std = self.tile_sample_min_size + W_out_std = self.tile_sample_min_size + for lin_idx in my_linear_indices: + ri = lin_idx // num_cols + rj = lin_idx % num_cols + i = ri * overlap_size + j = rj * overlap_size + tile = z[ + :, + :, + :, + i : i + self.tile_latent_min_size, + j : j + self.tile_latent_min_size, + ] + dec = self.decoder(tile) + + pad_h = max(0, H_out_std - dec.shape[-2]) + pad_w = max(0, W_out_std - dec.shape[-1]) + if pad_h > 0 or pad_w > 0: + dec = F.pad(dec, (0, pad_w, 0, pad_h, 0, 0), "constant", 0) + decoded_tiles.append(dec) + decoded_metas.append( + torch.tensor( + [ri, rj, pad_w, pad_h], device=z.device, dtype=torch.int64 + ) + ) + + while len(decoded_tiles) < tiles_per_rank: + zero_tile = torch.zeros( + [ + 1, + 3, + (T - 1) * self.ffactor_temporal + 1, + self.tile_sample_min_size, + self.tile_sample_min_size, + ], + device=dec.device, + dtype=dec.dtype, + ) + decoded_tiles.append(zero_tile) + meta_tensor = torch.tensor( + [-1, -1, self.tile_sample_min_size, self.tile_sample_min_size], + device=z.device, + dtype=torch.int64, + ) + decoded_metas.append(meta_tensor) + + decoded_tiles = torch.stack(decoded_tiles, dim=0) + decoded_metas = torch.stack(decoded_metas, dim=0) + + tiles_gather_list = [ + torch.empty_like(decoded_tiles) for _ in range(world_size) + ] + metas_gather_list = [ + torch.empty_like(decoded_metas) for _ in range(world_size) + ] + + dist.all_gather( + tiles_gather_list, + decoded_tiles, + group=get_parallel_state().sp_group, + ) + dist.all_gather( + metas_gather_list, + decoded_metas, + group=get_parallel_state().sp_group, + ) + + if rank != 0: + return torch.empty(0, device=z.device) + + rows = [[None for _ in range(num_cols)] for _ in range(num_rows)] + for r in range(world_size): + gathered_tiles_r = tiles_gather_list[ + r + ] # [tiles_per_rank, B, C, T, H, W] + gathered_metas_r = metas_gather_list[r] # [tiles_per_rank, 4] + for k in range(gathered_tiles_r.shape[0]): + ri = int(gathered_metas_r[k][0]) + rj = int(gathered_metas_r[k][1]) + if ri < 0 or rj < 0: + continue + if ri < num_rows and rj < num_cols: + # remove padding + pad_w = int(gathered_metas_r[k][2]) + pad_h = int(gathered_metas_r[k][3]) + h_end = None if pad_h == 0 else -pad_h + w_end = None if pad_w == 0 else -pad_w + rows[ri][rj] = gathered_tiles_r[k][:, :, :, :h_end, :w_end] + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if tile is None: + continue + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + return dec + + def spatial_tiled_decode(self, z: torch.Tensor): + if self._tile_parallelism_enabled: + return self.tile_parallel_spatial_tiled_decode(z) + + B, C, T, H, W = z.shape + overlap_size = int( + self.tile_latent_min_size * (1 - self.tile_overlap_factor) + ) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = z[ + :, + :, + :, + i : i + self.tile_latent_min_size, + j : j + self.tile_latent_min_size, + ] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + return dec + + def temporal_tiled_decode(self, z: torch.Tensor): + """Tiled temporal decoding for long sequence latents.""" + raise RuntimeError("Temporal tiling is not supported for this VAE.") + B, C, T, H, W = z.shape + overlap_size = int( + self.tile_latent_min_tsize * (1 - self.tile_overlap_factor) + ) + blend_extent = int( + self.tile_sample_min_tsize * self.tile_overlap_factor + ) + t_limit = self.tile_sample_min_tsize - blend_extent + assert 0 < overlap_size < self.tile_latent_min_tsize + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_tsize + 1, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_latent_min_size + or tile.shape[-2] > self.tile_latent_min_size + ): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self.decoder(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, : t_limit + 1, :, :]) + dec = torch.cat(result_row, dim=-3) + return dec + + def clear_cache(self): + # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + def encode(self, x: Tensor, return_dict: bool = True): + + def _encode(x): + """Slice encoding video latents.""" + _, _, num_frame, height, width = x.shape + + if ( + self.use_temporal_tiling + and x.shape[-3] > self.tile_sample_min_tsize + ): + return self.temporal_tiled_encode(x) + if self.use_spatial_tiling and ( + x.shape[-1] > self.tile_sample_min_size + or x.shape[-2] > self.tile_sample_min_size + ): + return self.spatial_tiled_encode(x) + + self.clear_cache() + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) # [B, C, 4, H, W] after patchify + out = torch.cat([out, out_], 2) + + self.clear_cache() + return out + + assert len(x.shape) == 5 # (B, C, T, H, W) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [_encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = _encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: Tensor, return_dict: bool = True, generator=None): + + def _decode(z): + """Slice decoding video latents.""" + _, _, num_frame, height, width = z.shape + + if ( + self.use_temporal_tiling + and z.shape[-3] > self.tile_latent_min_tsize + ): + return self.temporal_tiled_decode(z) + if self.use_spatial_tiling and ( + z.shape[-1] > self.tile_latent_min_size + or z.shape[-2] > self.tile_latent_min_size + ): + return self.spatial_tiled_decode(z) + + self.clear_cache() + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + z[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + z[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_], 2) + + self.clear_cache() + + return out + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [_decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = _decode(z) + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_posterior: bool = True, + return_dict: bool = True, + ): + """Performs the full forward pass of the Autoencoder. + + It encodes the input into a latent distribution, samples (or takes the mode) from it, + and then decodes it back to the image space. + + Args: + sample (torch.Tensor): Input image tensor (B, C, H, W). + sample_posterior (bool): If True, samples from the latent distribution (stochastic). + If False, uses the mode (deterministic). + return_posterior (bool): Whether to include the posterior distribution in the output. + return_dict (bool): If True, returns a structured DecoderOutput object; + otherwise, returns a standard tuple. + """ + # 1. Encode: Map input 'sample' to a latent distribution (Gaussian). + posterior = self.encode(sample).latent_dist + + # 2. Sample: Reparameterization trick. + # Use random sampling for training (variational) or mode for deterministic inference. + z = posterior.sample() if sample_posterior else posterior.mode() + + # 3. Decode: Reconstruct the image from the latent vector 'z'. + dec = self.decode(z).sample + + # 4. Return: Format the output based on 'return_dict' preference. + if return_dict: + return DecoderOutput(sample=dec, posterior=posterior) + else: + return (dec, posterior) diff --git a/worldcompass/fastvideo/models/hyvideo/models/text_encoders/__init__.py b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5aeb4dc5f4336ebd61d85cc6ee86762e52b427b3 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/__init__.py @@ -0,0 +1,549 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from copy import deepcopy +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from transformers import ( + AutoTokenizer, + AutoModel, +) +from transformers.utils import ModelOutput + + +def use_default(value, default): + """Utility: return value if not None, else default.""" + return value if value is not None else default + + +# Prompt templates for different models and tasks + + +__all__ = [ + "C_SCALE", + "PROMPT_TEMPLATE", + "MODEL_BASE", +] + +# =================== Constant Values ===================== +# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid +# overflow error when tensorboard logging values. +C_SCALE = 1_000_000_000_000_000 + +PROMPT_TEMPLATE_ENCODE_IMAGE_JSON = [ + { + "role": "system", + "content": "You are a helpful assistant. Describe the image by detailing the following aspects: \ + 1. The main content and theme of the image. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. The background environment, light, style and atmosphere.", + }, + {"role": "user", "content": "{}"}, +] + +PROMPT_TEMPLATE_ENCODE_VIDEO_JSON = [ + { + "role": "system", + "content": "You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video.", + }, + {"role": "user", "content": "{}"}, +] + +PROMPT_TEMPLATE = { + "li-dit-encode-image-json": { + "template": PROMPT_TEMPLATE_ENCODE_IMAGE_JSON, + "crop_start": -1, + }, # auto-calculate crop_start + "li-dit-encode-video-json": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO_JSON, + "crop_start": -1, + }, # auto-calculate crop_start +} + + +MODEL_BASE = os.getenv("MODEL_BASE", "") +TEXT_ENCODER_PATH = {} +TOKENIZER_PATH = {} + +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +def load_text_encoder( + text_encoder_type, + text_encoder_precision=None, + text_encoder_path=None, + logger=None, + device=None, +): + if text_encoder_path is None: + if text_encoder_type not in TEXT_ENCODER_PATH: + raise ValueError( + f"Unsupported text encoder type: {text_encoder_type}" + ) + text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] + + text_encoder = AutoModel.from_pretrained( + text_encoder_path, low_cpu_mem_usage=True + ) + + if hasattr(text_encoder, "language_model"): + text_encoder = text_encoder.language_model + text_encoder.final_layer_norm = text_encoder.norm + + # from_pretrained will ensure that the model is in eval mode. + if text_encoder_precision is not None: + text_encoder = text_encoder.to( + dtype=PRECISION_TO_TYPE[text_encoder_precision] + ) + + text_encoder.requires_grad_(False) + + if device is not None: + text_encoder = text_encoder.to(device) + + return text_encoder, text_encoder_path + + +def load_tokenizer( + tokenizer_type, tokenizer_path=None, padding_side="right", logger=None +): + processor = None + if tokenizer_path is None: + if tokenizer_type not in TOKENIZER_PATH: + raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") + tokenizer_path = TOKENIZER_PATH[tokenizer_type] + + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, padding_side=padding_side + ) + + return tokenizer, tokenizer_path, processor + + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list + (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + image_features: Optional[list] = None + + +class TextEncoder(nn.Module): + def __init__( + self, + text_encoder_type: str, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + tokenizer_type: Optional[str] = None, + tokenizer_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + prompt_template: Optional[dict] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + logger=None, + device=None, + ): + super().__init__() + self.text_encoder_type = text_encoder_type + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.tokenizer_type = ( + tokenizer_type if tokenizer_type is not None else text_encoder_type + ) + self.tokenizer_path = ( + tokenizer_path if tokenizer_path is not None else text_encoder_path + ) + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert ( + use_attention_mask is True + ), "Attention mask is True required when training videos." + self.prompt_template = prompt_template + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + self.logger = logger + + self.use_template = self.prompt_template is not None + if self.use_template: + assert ( + isinstance(self.prompt_template, dict) + and "template" in self.prompt_template + ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" + assert "{}" in str(self.prompt_template["template"]), ( + "`prompt_template['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template['template']}" + ) + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert ( + isinstance(self.prompt_template_video, dict) + and "template" in self.prompt_template_video + ), f"`prompt_template_video` must be a dictionary with a key 'template'," + f"got {self.prompt_template_video}" + assert "{}" in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}" + ) + + if text_encoder_type != "llm": + raise ValueError( + f"Unsupported text encoder type: {text_encoder_type}" + ) + self.output_key = output_key or "last_hidden_state" + + self.model, self.model_path = load_text_encoder( + text_encoder_type=self.text_encoder_type, + text_encoder_precision=self.precision, + text_encoder_path=self.model_path, + logger=self.logger, + device=device, + ) + + self.tokenizer, self.tokenizer_path, self.processor = load_tokenizer( + tokenizer_type=self.tokenizer_type, + tokenizer_path=self.tokenizer_path, + padding_side="right", + logger=self.logger, + ) + + # pre-calculate crop_start for image and video + if self.use_template and self.prompt_template is not None: + self.text2tokens("a photo of a cat", data_type="image") + if self.use_video_template and self.prompt_template_video is not None: + self.text2tokens("a photo of a cat", data_type="video") + + @property + def dtype(self): + return self.model.dtype + + @property + def device(self): + return self.model.device + + def __repr__(self): + return ( + f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + ) + + @staticmethod + def apply_text_to_template(text, template, prevent_empty_text=True): + """Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If Ture, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + elif isinstance(template, list): + # For JSON list template format (chat conversation) + # Create a deep copy to avoid modifying the original template + template_copy = deepcopy(template) + for item in template_copy: + if isinstance(item, dict) and "content" in item: + # Replace placeholder with text in the content field + item["content"] = item["content"].format( + text if text else (" " if prevent_empty_text else "") + ) + return template_copy + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def calculate_crop_start(self, tokenized_input): + """Automatically calculate the crop_start position based on identifying user tokens. + + Args: + tokenized_input: The output from the tokenizer containing input_ids + + Returns: + int: The position where the actual prompt content begins (after user markers) + """ + input_ids = tokenized_input["input_ids"][ + 0 + ].tolist() # Get the first example's tokens + + marker = "<|im_start|>user\n" + + # Tokenize just the marker to get its token IDs + marker_tokens = self.tokenizer(marker, add_special_tokens=False)[ + "input_ids" + ] + + # Find the end position of the marker in the input sequence + for i in range(len(input_ids) - len(marker_tokens) + 1): + if input_ids[i : i + len(marker_tokens)] == marker_tokens: + # Return the position after the marker + return i + len(marker_tokens) + + # If marker not found, try to find based on special tokens + if hasattr(self.tokenizer, "special_tokens_map"): + # Check for user token or any other special token that might indicate user input start + for ( + token_name, + token_value, + ) in self.tokenizer.special_tokens_map.items(): + if "user" in token_name.lower(): + user_token_id = self.tokenizer.convert_tokens_to_ids( + token_value + ) + if user_token_id in input_ids: + return input_ids.index(user_token_id) + 1 + + # Default fallback: return 0 (no cropping) + return 0 + + def text2tokens(self, text, data_type="image", max_length=300): + """Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = "str" + if self.use_template or self.use_video_template: + if data_type == "image": + prompt_template = self.prompt_template["template"] + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + prompt_template = self.prompt_template_video["template"] + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [ + self.apply_text_to_template(one_text, prompt_template) + for one_text in text + ] + if isinstance(text[0], list): + tokenize_input_type = "list" + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = "list" + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + # First pass: tokenize with arbitrary max_length to find crop_start + if crop_start == -1: + # Use temporary max_length for the first pass (large enough) + temp_kwargs = dict( + truncation=True, + max_length=256, # Temporary large value + padding="max_length", + return_tensors="pt", + ) + + # First tokenization pass to calculate crop_start + if tokenize_input_type == "str": + temp_tokenized = self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **temp_kwargs, + ) + elif tokenize_input_type == "list": + temp_tokenized = self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **temp_kwargs, + ) + + # Calculate the crop_start from this first pass + crop_start = self.calculate_crop_start(temp_tokenized) + + # Store the calculated crop_start for future use + if data_type == "image": + self.prompt_template["crop_start"] = crop_start + else: + self.prompt_template_video["crop_start"] = crop_start + else: + crop_start = 0 + + # Second pass: tokenize with the proper max_length using the found crop_start + kwargs = dict( + truncation=True, + max_length=max_length + (crop_start if crop_start > 0 else 0), + padding="max_length", + return_tensors="pt", + ) + + if tokenize_input_type == "str": + tokenized_output = self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **kwargs, + ) + elif tokenize_input_type == "list": + tokenized_output = self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **kwargs, + ) + else: + raise ValueError( + f"Unsupported tokenize_input_type: {tokenize_input_type}" + ) + + return tokenized_output + + def encode( + self, + batch_encoding, + use_attention_mask=None, + output_hidden_states=False, + do_sample=None, + hidden_state_skip_layer=None, + return_texts=False, + data_type="image", + device=None, + is_uncond=False, + ): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + device = self.model.device if device is None else device + use_attention_mask = use_default( + use_attention_mask, self.use_attention_mask + ) + hidden_state_skip_layer = use_default( + hidden_state_skip_layer, self.hidden_state_skip_layer + ) + do_sample = use_default(do_sample, not self.reproduce) + + attention_mask = ( + batch_encoding["attention_mask"].to(device) + if use_attention_mask + else None + ) + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + attention_mask=attention_mask, + output_hidden_states=output_hidden_states + or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[ + -(hidden_state_skip_layer + 1) + ] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm( + last_hidden_state + ) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", 0) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", 0) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = ( + attention_mask[:, crop_start:] + if use_attention_mask + else None + ) + + if output_hidden_states: + return TextEncoderModelOutput( + last_hidden_state, attention_mask, outputs.hidden_states + ) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text, max_length=self.max_length) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/__init__.py b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..78737cb0b54230de784b336755388e4b1ca4fe56 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/__init__.py @@ -0,0 +1,233 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json +import torch +import torch.nn as nn +from transformers import AutoTokenizer, T5ForConditionalGeneration + + +def load_glyph_byT5_v2(args, device): + """Loads ByT5 tokenizer and encoder model for glyph encoding. + + Args: + args (dict): Configuration dictionary containing paths and settings. + device (str or torch.device): Device to load the model onto. + + Returns: + dict: Dictionary with keys 'byt5_tokenizer', 'byt5_model', 'byt5_max_length'. + """ + byt5_tokenizer, byt5_model, byt5_max_length = create_byt5(args, device) + byt5_model = byt5_model.to(device=device) + return { + "byt5_tokenizer": byt5_tokenizer, + "byt5_model": byt5_model, + "byt5_max_length": byt5_max_length, + } + + +def create_byt5(args, device): + """Create ByT5 tokenizer and encoder, load weights if provided. + + Args: + args (dict): Configuration dictionary. + device (str or torch.device): Device to load the model onto. + + Returns: + tuple: (byt5_tokenizer, byt5_model, byt5_max_length) + """ + byt5_max_length = args["byt5_max_length"] + byt5_config = dict( + byt5_name=args["byT5_google_path"], + special_token=True, + color_special_token=True, + font_special_token=True, + color_ann_path=args["multilingual_prompt_format_color_path"], + font_ann_path=args["multilingual_prompt_format_font_path"], + multilingual=True, + ) + huggingface_cache_dir = None + byt5_model, byt5_tokenizer = load_byt5_and_byt5_tokenizer( + **byt5_config, + huggingface_cache_dir=huggingface_cache_dir, + device=device, + ) + + # Load custom checkpoint if provided + if args["byT5_ckpt_path"] is not None: + if "cuda" not in str(device): + byt5_state_dict = torch.load( + args["byT5_ckpt_path"], map_location=device + ) + else: + byt5_state_dict = torch.load( + args["byT5_ckpt_path"], map_location=device + ) + if "state_dict" in byt5_state_dict: + sd = byt5_state_dict["state_dict"] + newsd = {} + for k, v in sd.items(): + if k.startswith("module.text_tower.encoder."): + newsd[k[len("module.text_tower.encoder.") :]] = v + byt5_state_dict = newsd + byt5_model.load_state_dict(byt5_state_dict) + byt5_model.requires_grad_(False) + return byt5_tokenizer, byt5_model, byt5_max_length + + +def add_special_token( + tokenizer, + text_encoder, + add_color, + add_font, + color_ann_path, + font_ann_path, + multilingual=False, +): + """Add special tokens for color and font to tokenizer and text encoder. + + Args: + tokenizer: Huggingface tokenizer. + text_encoder: Huggingface T5 encoder. + add_color (bool): Whether to add color tokens. + add_font (bool): Whether to add font tokens. + color_ann_path (str): Path to color annotation JSON. + font_ann_path (str): Path to font annotation JSON. + multilingual (bool): Whether to use multilingual font tokens. + """ + with open(font_ann_path, "r") as f: + idx_font_dict = json.load(f) + with open(color_ann_path, "r") as f: + idx_color_dict = json.load(f) + + if multilingual: + font_token = [ + f"<{font_code[:2]}-font-{idx_font_dict[font_code]}>" + for font_code in idx_font_dict + ] + else: + font_token = [f"" for i in range(len(idx_font_dict))] + color_token = [f"" for i in range(len(idx_color_dict))] + additional_special_tokens = [] + if add_color: + additional_special_tokens += color_token + if add_font: + additional_special_tokens += font_token + + tokenizer.add_tokens(additional_special_tokens, special_tokens=True) + # Set mean_resizing=False to avoid PyTorch LAPACK dependency + text_encoder.resize_token_embeddings(len(tokenizer), mean_resizing=False) + + +def load_byt5_and_byt5_tokenizer( + byt5_name="google/byt5-small", + special_token=False, + color_special_token=False, + font_special_token=False, + color_ann_path="assets/color_idx.json", + font_ann_path="assets/font_idx_512.json", + huggingface_cache_dir=None, + multilingual=False, + device=None, +): + """Load ByT5 encoder and tokenizer from Huggingface, and add special tokens if needed. + + Args: + byt5_name (str): Model name or path. + special_token (bool): Whether to add special tokens. + color_special_token (bool): Whether to add color tokens. + font_special_token (bool): Whether to add font tokens. + color_ann_path (str): Path to color annotation JSON. + font_ann_path (str): Path to font annotation JSON. + huggingface_cache_dir (str): Huggingface cache directory. + multilingual (bool): Whether to use multilingual font tokens. + device (str or torch.device): Device to load the model onto. + + Returns: + tuple: (byt5_text_encoder, byt5_tokenizer) + """ + byt5_tokenizer = AutoTokenizer.from_pretrained( + byt5_name, + cache_dir=huggingface_cache_dir, + ) + byt5_text_encoder = T5ForConditionalGeneration.from_pretrained( + byt5_name, + cache_dir=huggingface_cache_dir, + ).get_encoder() + + if "cuda" not in str(device): + device = torch.device(device) + else: + device = torch.device(device) + byt5_text_encoder = byt5_text_encoder.to(device) + + if special_token: + add_special_token( + byt5_tokenizer, + byt5_text_encoder, + add_color=color_special_token, + add_font=font_special_token, + color_ann_path=color_ann_path, + font_ann_path=font_ann_path, + multilingual=multilingual, + ) + return byt5_text_encoder, byt5_tokenizer + + +class ByT5Mapper(nn.Module): + """ + ByT5Mapper: Maps ByT5 encoder outputs to a new space, with optional residual connection. + + Args: + in_dim (int): Input dimension (must equal out_dim if use_residual). + out_dim (int): Output dimension after second linear layer. + hidden_dim (int): Hidden dimension for intermediate layer. + out_dim1 (int): Final output dimension. + use_residual (bool): Whether to use residual connection (default: True). + """ + + def __init__( + self, in_dim, out_dim, hidden_dim, out_dim1, use_residual=True + ): + super().__init__() + if use_residual: + assert in_dim == out_dim + self.layernorm = nn.LayerNorm(in_dim) + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, out_dim) + self.fc3 = nn.Linear(out_dim, out_dim1) + self.use_residual = use_residual + self.act_fn = nn.GELU() + + def forward(self, x): + """Forward pass for ByT5Mapper. + + Args: + x (Tensor): Input tensor of shape (..., in_dim). + + Returns: + Tensor: Output tensor of shape (..., out_dim1). + """ + residual = x + x = self.layernorm(x) + x = self.fc1(x) + x = self.act_fn(x) + x = self.fc2(x) + x2 = self.act_fn(x) + x2 = self.fc3(x2) + if self.use_residual: + x2 = x2 + residual + return x2 diff --git a/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/format_prompt.py b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/format_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..3fb801214bd2ccef43ffd1143d64643c92a9d3ff --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/text_encoders/byT5/format_prompt.py @@ -0,0 +1,84 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json + + +def closest_color(requested_color): + import webcolors + + min_colors = {} + for key, name in webcolors.CSS3_HEX_TO_NAMES.items(): + r_c, g_c, b_c = webcolors.hex_to_rgb(key) + rd = (r_c - requested_color[0]) ** 2 + gd = (g_c - requested_color[1]) ** 2 + bd = (b_c - requested_color[2]) ** 2 + min_colors[(rd + gd + bd)] = name + return min_colors[min(min_colors.keys())] + + +def convert_rgb_to_names(rgb_tuple): + try: + import webcolors + + color_name = webcolors.rgb_to_name(rgb_tuple) + except ValueError: + color_name = closest_color(rgb_tuple) + return color_name + + +class MultilingualPromptFormat: + def __init__( + self, + font_path: str = "assets/glyph_sdxl_assets/multilingual_10-lang_idx.json", + color_path: str = "assets/glyph_sdxl_assets/color_idx.json", + ): + with open(font_path, "r") as f: + self.font_dict = json.load(f) + with open(color_path, "r") as f: + self.color_dict = json.load(f) + + def format_prompt(self, texts, styles): + """Text "{text}" in {color}, {type}.""" + + prompt = "" + for text, style in zip(texts, styles): + text_prompt = f'Text "{text}"' + + attr_list = [] + + # format color + if style["color"] is not None: + import webcolors + + hex_color = style["color"] + rgb_color = webcolors.hex_to_rgb(hex_color) + color_name = convert_rgb_to_names(rgb_color) + attr_list.append(f"") + + # format font + if style["font-family"] is not None: + attr_list.append( + f"<{style['font-family'][:2]}-font-{self.font_dict[style['font-family']]}>" + ) + attr_suffix = ", ".join(attr_list) + text_prompt += " in " + attr_suffix + text_prompt += ". " + else: + text_prompt += ". " + + prompt = prompt + text_prompt + return prompt diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..8231eed5caa8e364b23cf94336406bd150c80b93 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/ar_action_hunyuanvideo_1_5_transformer.py @@ -0,0 +1,1280 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. +from distutils.command.check import check +from typing import Any, List, Tuple, Optional, Union, Dict + +import os +import torch +import torch.nn as nn +from einops import rearrange +from loguru import logger + +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from .modules.activation_layers import get_activation_layer +from .modules.norm_layers import get_norm_layer +from .modules.embed_layers import ( + TimestepEmbedder, + PatchEmbed, + TextProjection, + VisionProjection, +) +from .modules.attention import ( + parallel_attention, + sequence_parallel_attention_txt, + sequence_parallel_attention_vision, +) +from .modules.posemb_layers import apply_rotary_emb, get_nd_rotary_pos_embed +from .modules.mlp_layers import ( + MLP, + MLPEmbedder, + FinalLayer, + LinearWarpforSingle, +) +from .modules.modulate_layers import ModulateDiT, modulate, apply_gate +from .modules.token_refiner import SingleTokenRefiner + +from fastvideo.models.hyvideo.utils.infer_utils import torch_compile_wrapper +from fastvideo.models.hyvideo.models.text_encoders.byT5 import ByT5Mapper +from fastvideo.models.prope.camera_rope import prope_qkv + +from fastvideo.distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_world_size, +) +from fastvideo.distributed import sequence_model_parallel_all_gather +from fastvideo.configs.models.dits import HunyuanVideoConfig + + +def is_blocks(n: str, m) -> bool: + print("is_blocks", n, flush=True) + return ( + "double_blocks" in n + and str.isdigit(n.split(".")[-1]) + or "single_blocks" in n + and str.isdigit(n.split(".")[-1]) + ) + + +class MMDoubleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + self.attn_mode = attn_mode + + self.hidden_size = hidden_size + self.qkv_bias = qkv_bias + self.factory_kwargs = factory_kwargs + + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.img_norm1 = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + self.img_attn_q = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.img_attn_k = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.img_attn_v = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.img_attn_k_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.img_attn_proj = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + self.img_norm2 = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + + self.txt_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.txt_norm1 = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + + self.txt_attn_q = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.txt_attn_k = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.txt_attn_v = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + self.txt_attn_q_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.txt_attn_k_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.txt_attn_proj = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + self.txt_norm2 = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + self.txt_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward_txt( + self, + txt: torch.Tensor, + vec_txt: torch.Tensor, + text_mask=None, + attn_param=None, + is_flash=False, + block_idx=None, + kv_cache: Optional[dict] = None, + cache_txt: bool = False, + ) -> Tuple[torch.Tensor]: + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec_txt).chunk(6, dim=-1) + + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate( + txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale + ) + txt_q = self.txt_attn_q(txt_modulated) + txt_k = self.txt_attn_k(txt_modulated) + txt_v = self.txt_attn_v(txt_modulated) + txt_q = rearrange(txt_q, "B L (H D) -> B L H D", H=self.heads_num) + txt_k = rearrange(txt_k, "B L (H D) -> B L H D", H=self.heads_num) + txt_v = rearrange(txt_v, "B L (H D) -> B L H D", H=self.heads_num) + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + + # attn_mode = 'flash' if is_flash else self.attn_mode + attn_mode = ( + "torch_causal" # for ar model, the default mode is flex_causal + ) + txt_attn, t_kv_cache = sequence_parallel_attention_txt( + (txt_q), + (txt_k), + (txt_v), + img_q_len=txt_q.shape[1], + img_kv_len=txt_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + block_idx=block_idx, + kv_cache=kv_cache, + cache_txt=cache_txt, + ) + + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp( + modulate( + self.txt_norm2(txt), + shift=txt_mod2_shift, + scale=txt_mod2_scale, + ) + ), + gate=txt_mod2_gate, + ) + return txt, t_kv_cache + + def forward_vision( + self, + img: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + attn_param=None, + block_idx=None, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + kv_cache: Optional[dict] = None, + cache_vision: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + + img_modulated = self.img_norm1(img) + img_modulated = modulate( + img_modulated, shift=img_mod1_shift, scale=img_mod1_scale + ) + + img_q = self.img_attn_q(img_modulated) + img_k = self.img_attn_k(img_modulated) + img_v = self.img_attn_v(img_modulated) + img_q = rearrange(img_q, "B L (H D) -> B L H D", H=self.heads_num) + img_k = rearrange(img_k, "B L (H D) -> B L H D", H=self.heads_num) + img_v = rearrange(img_v, "B L (H D) -> B L H D", H=self.heads_num) + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + + # 添加连续的camera pose，通过prope + img_q_prope, img_k_prope, img_v_prope, apply_fn_o = prope_qkv( + img_q.permute(0, 2, 1, 3), + img_k.permute(0, 2, 1, 3), + img_v.permute(0, 2, 1, 3), + viewmats=viewmats, + Ks=Ks, + ) # [batch, num_heads, seqlen, head_dim] + img_q_prope = img_q_prope.permute( + 0, 2, 1, 3 + ) # [batch, seqlen, num_heads, head_dim] + img_k_prope = img_k_prope.permute( + 0, 2, 1, 3 + ) # [batch, seqlen, num_heads, head_dim] + img_v_prope = img_v_prope.permute( + 0, 2, 1, 3 + ) # [batch, seqlen, num_heads, head_dim] + + if freqs_cis is not None: + img_qq, img_kk = apply_rotary_emb( + img_q, img_k, freqs_cis, head_first=False + ) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + # attn_mode = 'flash' if is_flash else self.attn_mode + attn_mode = ( + "torch_causal" # for ar model, the default mode is flex_causal + ) + img_attn, img_attn_prope, vison_kv_cache = ( + sequence_parallel_attention_vision( + (img_q, img_q_prope), + (img_k, img_k_prope), + (img_v, img_v_prope), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + attn_mode=attn_mode, + attn_param=attn_param, + block_idx=block_idx, + kv_cache=kv_cache, + cache_vision=cache_vision, + ) + ) + + img_attn_prope = rearrange( + img_attn_prope, "B L (H D) -> B H L D", H=self.heads_num + ) + img_attn_prope = apply_fn_o( + img_attn_prope + ) # [batch, num_heads, seqlen, head_dim] + img_attn_prope = rearrange(img_attn_prope, "B H L D -> B L (H D)") + + img = img + apply_gate( + self.img_attn_proj(img_attn) + + self.img_attn_prope_proj(img_attn_prope), + gate=img_mod1_gate, + ) + img = img + apply_gate( + self.img_mlp( + modulate( + self.img_norm2(img), + shift=img_mod2_shift, + scale=img_mod2_scale, + ) + ), + gate=img_mod2_gate, + ) + + return img, vison_kv_cache + + def forward( + self, + txt_branch: bool, + input_dict: dict, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if txt_branch: + return self.forward_txt(**input_dict) + else: + return self.forward_vision(**input_dict) + + +class MMSingleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.attn_mode = attn_mode + + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim**-0.5 + + self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_k = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_v = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_mlp = nn.Linear( + hidden_size, mlp_hidden_dim, **factory_kwargs + ) + self.linear2 = LinearWarpforSingle( + hidden_size + mlp_hidden_dim, + hidden_size, + bias=True, + **factory_kwargs, + ) + self.mlp_act = get_activation_layer(mlp_act_type)() + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + + self.pre_norm = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + self.modulation = ModulateDiT( + hidden_size, + factor=3, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec_txt: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask=None, + attn_param=None, + is_flash=False, + ) -> torch.Tensor: + """Forward pass for the single stream block.""" + mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) + txt_mod_shift, txt_mod_scale, txt_mod_gate = self.modulation( + vec_txt + ).chunk(3, dim=-1) + + x_pre_norm = self.pre_norm(x) + img_token, txt_token = ( + x_pre_norm[:, :-txt_len, :], + x_pre_norm[:, -txt_len:, :], + ) + + img_mod = modulate(img_token, shift=mod_shift, scale=mod_scale) + txt_mod = modulate(txt_token, shift=txt_mod_shift, scale=txt_mod_scale) + x_mod = torch.cat([img_mod, txt_mod], dim=1) + + q = self.linear1_q(x_mod) + k = self.linear1_k(x_mod) + v = self.linear1_v(x_mod) + + q = rearrange(q, "B L (H D) -> B L H D", H=self.heads_num) + k = rearrange(k, "B L (H D) -> B L H D", H=self.heads_num) + v = rearrange(v, "B L (H D) -> B L H D", H=self.heads_num) + + mlp = self.linear1_mlp(x_mod) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb( + img_q, img_k, freqs_cis, head_first=False + ) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + if is_flash: + attn_mode = "flash" + else: + attn_mode = self.attn_mode + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + ) + output = self.linear2(attn, self.mlp_act(mlp)) + img_output, txt_output = ( + output[:, :-txt_len, :], + output[:, -txt_len:, :], + ) + + img_output = apply_gate(img_output, gate=mod_gate) + txt_output = apply_gate(txt_output, gate=txt_mod_gate) + gate_output = torch.cat([img_output, txt_output], dim=1) + + return x + gate_output + + +class ARHunyuanVideo_1_5_DiffusionTransformer(ModelMixin, ConfigMixin): + """HunyuanVideo Transformer backbone. + + Args: + patch_size (list): The size of the patch. + in_channels (int): The number of input channels. + out_channels (int): The number of output channels. + hidden_size (int): The hidden size of the transformer backbone. + heads_num (int): The number of attention heads. + mlp_width_ratio (float): Width ratio for the transformer MLPs. + mlp_act_type (str): Activation type for the transformer MLPs. + mm_double_blocks_depth (int): Number of double-stream transformer blocks. + mm_single_blocks_depth (int): Number of single-stream transformer blocks. + rope_dim_list (list): Rotary embedding dim for t, h, w. + qkv_bias (bool): Use bias in qkv projection. + qk_norm (bool): Whether to use qk norm. + qk_norm_type (str): Type of qk norm. + guidance_embed (bool): Use guidance embedding for distillation. + text_projection (str): Text input projection. Default is "single_refiner". + use_attention_mask (bool): If to use attention mask. + text_states_dim (int): Text encoder output dim. + text_states_dim_2 (int): Secondary text encoder output dim. + text_pool_type (str): Type for text pooling. + rope_theta (int): Rotary embedding theta parameter. + attn_mode (str): Attention mode identifier. + attn_param (dict): Attention parameter dictionary. + glyph_byT5_v2 (bool): Use ByT5 glyph module. + vision_projection (str): Vision condition embedding mode. + vision_states_dim (int): Vision encoder states input dim. + is_reshape_temporal_channels (bool): For video VAE adaptation. + use_cond_type_embedding (bool): Use condition type embedding. + """ + + _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 4, + concat_condition: bool = True, + out_channels: int = None, + hidden_size: int = 3072, + heads_num: int = 24, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + mm_double_blocks_depth: int = 20, + mm_single_blocks_depth: int = 40, + rope_dim_list: list = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + guidance_embed: bool = False, + use_meanflow: bool = False, + text_projection: str = "single_refiner", + use_attention_mask: bool = True, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + text_pool_type: str = None, + rope_theta: int = 256, + attn_mode: str = "flash", + attn_param: dict = None, + glyph_byT5_v2: bool = False, + vision_projection: str = "none", + vision_states_dim: int = 1280, + is_reshape_temporal_channels: bool = False, + use_cond_type_embedding: bool = False, + ideal_resolution: str = None, + ideal_task: str = None, + ): + super().__init__() + factory_kwargs = {} + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = ( + in_channels if out_channels is None else out_channels + ) + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.use_attention_mask = use_attention_mask + self.text_projection = text_projection + self.attn_mode = attn_mode + self.text_pool_type = text_pool_type + self.text_states_dim = text_states_dim + self.text_states_dim_2 = text_states_dim_2 + self.vision_states_dim = vision_states_dim + + self.glyph_byT5_v2 = glyph_byT5_v2 + if self.glyph_byT5_v2: + self.byt5_in = ByT5Mapper( + in_dim=1472, + out_dim=2048, + hidden_dim=2048, + out_dim1=hidden_size, + use_residual=False, + ) + + if hidden_size % heads_num != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}" + ) + pe_dim = hidden_size // heads_num + if sum(rope_dim_list) != pe_dim: + raise ValueError( + f"Got {rope_dim_list} but expected positional dim {pe_dim}" + ) + self.hidden_size = hidden_size + self.heads_num = heads_num + + self.img_in = PatchEmbed( + self.patch_size, + self.in_channels, + self.hidden_size, + is_reshape_temporal_channels=is_reshape_temporal_channels, + concat_condition=concat_condition, + **factory_kwargs, + ) + + # Vision projection + if vision_projection == "linear": + self.vision_in = VisionProjection( + input_dim=self.vision_states_dim, output_dim=self.hidden_size + ) + else: + self.vision_in = None + + # Text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + text_states_dim, + hidden_size, + heads_num, + depth=2, + **factory_kwargs, + ) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}" + ) + + # time modulation + self.time_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + self.vector_in = ( + MLPEmbedder( + self.config.text_states_dim_2, + self.hidden_size, + **factory_kwargs, + ) + if self.text_pool_type is not None + else None + ) + self.guidance_in = ( + TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + if guidance_embed + else None + ) + + self.time_r_in = ( + TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + if use_meanflow + else None + ) + + self.double_blocks = nn.ModuleList( + [ + MMDoubleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode=attn_mode, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + for _ in range(mm_double_blocks_depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + MMSingleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode=attn_mode, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) + for _ in range(mm_single_blocks_depth) + ] + ) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs, + ) + + # STA + if attn_param is None: + self.attn_param = { + # STA + "win_size": [[3, 3, 3]], + "win_type": "fixed", + "win_ratio": 10, + "tile_size": [6, 8, 8], + # SSTA + "ssta_topk": 64, + "ssta_threshold": 0.0, + "ssta_lambda": 0.7, + "ssta_sampling_type": "importance", + "ssta_adaptive_pool": None, + # flex-block-attn: + "attn_sparse_type": "ssta", + "attn_pad_type": "zero", + "attn_use_text_mask": 1, + "attn_mask_share_within_head": 0, + } + else: + self.attn_param = attn_param + + if attn_mode == "flex-block-attn": + self.register_to_config(attn_param=self.attn_param) + + if use_cond_type_embedding: + self.cond_type_embedding = nn.Embedding(3, self.hidden_size) + self.cond_type_embedding.weight.data.fill_(0) + assert ( + self.glyph_byT5_v2 + ), "text type embedding is only used when glyph_byT5_v2 is True" + assert ( + vision_projection is not None + ), "text type embedding is only used when vision_projection is not None" + # 0: text_encoder feature + # 1: byt5 feature + # 2: vision_encoder feature + else: + self.cond_type_embedding = None + + self.gradient_checkpointing = False + + def load_hunyuan_state_dict(self, model_path): + load_key = "module" + bare_model = "unknown" + + if model_path.endswith(".safetensors"): + from safetensors.torch import load_file + + state_dict = load_file(model_path, device="cpu") + else: + state_dict = torch.load( + model_path, map_location="cpu", weights_only=True + ) + + if bare_model == "unknown" and ( + "ema" in state_dict or "module" in state_dict + ): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError( + f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}." + ) + + result = self.load_state_dict(state_dict, strict=False) + + if result.missing_keys: + logger.info("[load.py] Missing keys when loading state_dict:") + for key in result.missing_keys: + logger.info(f"[load.py] Missing key: {key}") + if result.unexpected_keys: + logger.info("[load.py] Unexpected keys when loading state_dict:") + for key in result.unexpected_keys: + logger.info(f"[load.py] Unexpected key: {key}") + if result.missing_keys or result.unexpected_keys: + raise ValueError( + f"Missing: {result.missing_keys}, Unexpected: {result.unexpected_keys}" + ) + + return result + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def get_rotary_pos_embed(self, rope_sizes): + target_ndim = 3 + head_dim = self.hidden_size // self.heads_num + rope_dim_list = self.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [ + head_dim // target_ndim for _ in range(target_ndim) + ] + assert ( + sum(rope_dim_list) == head_dim + ), "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + + def reorder_txt_token( + self, + byt5_txt, + txt, + byt5_text_mask, + text_mask, + zero_feat=False, + is_reorder=True, + ): + if is_reorder: + reorder_txt = [] + reorder_mask = [] + for i in range(text_mask.shape[0]): + byt5_text_mask_i = byt5_text_mask[i].bool() + text_mask_i = text_mask[i].bool() + + byt5_txt_i = byt5_txt[i] + txt_i = txt[i] + if zero_feat: + # When using block mask with approximate computation, + # set pad to zero to reduce error + pad_byt5 = torch.zeros_like(byt5_txt_i[~byt5_text_mask_i]) + pad_text = torch.zeros_like(txt_i[~text_mask_i]) + reorder_txt_i = torch.cat( + [ + byt5_txt_i[byt5_text_mask_i], + txt_i[text_mask_i], + pad_byt5, + pad_text, + ], + dim=0, + ) + else: + reorder_txt_i = torch.cat( + [ + byt5_txt_i[byt5_text_mask_i], + txt_i[text_mask_i], + byt5_txt_i[~byt5_text_mask_i], + txt_i[~text_mask_i], + ], + dim=0, + ) + reorder_mask_i = torch.cat( + [ + byt5_text_mask_i[byt5_text_mask_i], + text_mask_i[text_mask_i], + byt5_text_mask_i[~byt5_text_mask_i], + text_mask_i[~text_mask_i], + ], + dim=0, + ) + + reorder_txt.append(reorder_txt_i) + reorder_mask.append(reorder_mask_i) + + reorder_txt = torch.stack(reorder_txt) + reorder_mask = torch.stack(reorder_mask).to(dtype=torch.int64) + else: + reorder_txt = torch.concat([byt5_txt, txt], dim=1) + reorder_mask = torch.concat([byt5_text_mask, text_mask], dim=1).to( + dtype=torch.int64 + ) + + return reorder_txt, reorder_mask + + def add_discrete_action_parameters(self): + self.action_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu") + ) + nn.init.zeros_(self.action_in.mlp[2].weight) + # 如果你想初始化偏置为0，执行以下操作： + if self.action_in.mlp[2].bias is not None: + nn.init.zeros_(self.action_in.mlp[2].bias) + + # prope参数,零初始化 + for block in self.double_blocks: + block.img_attn_prope_proj = nn.Linear( + block.hidden_size, + block.hidden_size, + bias=block.qkv_bias, + **block.factory_kwargs, + ) + nn.init.zeros_(block.img_attn_prope_proj.weight) + + # 如果你想初始化偏置为0，执行以下操作： + if block.img_attn_prope_proj.bias is not None: + nn.init.zeros_(block.img_attn_prope_proj.bias) + + def forward_txt( + self, + timestep_txt: torch.Tensor, + text_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + vision_states: torch.Tensor = None, + mask_type="t2v", + extra_kwargs=None, + kv_cache: Optional[dict] = None, + cache_txt: Optional[bool] = False, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + + text_mask = encoder_attention_mask + txt = text_states + bs = txt.shape[0] + + # Prepare modulation vectors + vec_txt = self.time_in(timestep_txt) # 为了txt prompt单独计算 + + # Embed text tokens + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in( + txt, + timestep_txt, + text_mask if self.use_attention_mask else None, + ) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}" + ) + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + torch.zeros_like( + txt[:, :, 0], device=text_mask.device, dtype=torch.long + ) + ) + txt = txt + cond_emb + + if self.glyph_byT5_v2: + byt5_text_states = extra_kwargs["byt5_text_states"] + byt5_text_mask = extra_kwargs["byt5_text_mask"] + byt5_txt = self.byt5_in(byt5_text_states) + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + torch.ones_like( + byt5_txt[:, :, 0], + device=byt5_txt.device, + dtype=torch.long, + ) + ) + byt5_txt = byt5_txt + cond_emb + txt, text_mask = self.reorder_txt_token( + byt5_txt, txt, byt5_text_mask, text_mask, zero_feat=True + ) + + if self.vision_in is not None and vision_states is not None: + extra_encoder_hidden_states = self.vision_in(vision_states) + # If t2v, set extra_attention_mask to 0 to avoid attention to semantic tokens + if mask_type == "t2v" and torch.all(vision_states == 0): + extra_attention_mask = torch.zeros( + (bs, extra_encoder_hidden_states.shape[1]), + dtype=text_mask.dtype, + device=text_mask.device, + ) + # Set vision tokens to zero to mitigate potential block mask error in SSTA + extra_encoder_hidden_states = extra_encoder_hidden_states * 0.0 + else: + extra_attention_mask = torch.ones( + (bs, extra_encoder_hidden_states.shape[1]), + dtype=text_mask.dtype, + device=text_mask.device, + ) + # Ensure valid tokens precede padding tokens + if self.cond_type_embedding is not None: + cond_emb = self.cond_type_embedding( + 2 + * torch.ones_like( + extra_encoder_hidden_states[:, :, 0], + dtype=torch.long, + device=extra_encoder_hidden_states.device, + ) + ) + extra_encoder_hidden_states = ( + extra_encoder_hidden_states + cond_emb + ) + + txt, text_mask = self.reorder_txt_token( + extra_encoder_hidden_states, + txt, + extra_attention_mask, + text_mask, + ) + + txt = txt[text_mask.bool().to(txt.device)].unsqueeze(0) + + if cache_txt: + _kv_cache_now = [] + transformer_num_layers = len(self.double_blocks) + for i in range(transformer_num_layers): + _kv_cache_now.append( + { + "k_vision": None, + "v_vision": None, + "k_txt": None, + "v_txt": None, + } + ) + + # Pass through double-stream blocks + for index, block in enumerate(self.double_blocks): + input_dict = { + "txt": txt, + "vec_txt": vec_txt, + "text_mask": None, # we have masked txt tokens already, set None here + "attn_param": None, + "is_flash": False, + "block_idx": index, + "kv_cache": kv_cache, + "cache_txt": cache_txt, + } + if torch.is_grad_enabled(): + txt, t_kv = torch.utils.checkpoint.checkpoint( + block, + True, + input_dict, + use_reentrant=False, + ) + else: + txt, t_kv = block( + txt_branch=True, + input_dict=input_dict, + ) + if cache_txt: + _kv_cache_now[index]["k_txt"] = t_kv["k_txt"] + _kv_cache_now[index]["v_txt"] = t_kv["v_txt"] + + if cache_txt: + return _kv_cache_now + + def forward_vision( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + timestep_r=None, + freqs_cos: Optional[torch.Tensor] = None, + freqs_sin: Optional[torch.Tensor] = None, + return_dict: bool = False, + mask_type="t2v", + extra_kwargs=None, + action: Optional[torch.Tensor] = None, + viewmats: Optional[torch.Tensor] = None, + Ks: Optional[torch.Tensor] = None, + kv_cache: Optional[dict] = None, + cache_vision: bool = False, + rope_temporal_size=4, + start_rope_start_idx=0, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + + img = x = hidden_states + t = timestep + bs, _, ot, oh, ow = x.shape + tt, th, tw = ( + ot // self.patch_size[0], + oh // self.patch_size[1], + ow // self.patch_size[2], + ) + self.attn_param["thw"] = [tt, th, tw] + rope_temporal_size = rope_temporal_size // self.patch_size[0] + if freqs_cos is None and freqs_sin is None: + freqs_cos, freqs_sin = self.get_rotary_pos_embed( + (rope_temporal_size, th, tw) + ) + per_latent_size = th * tw + start = start_rope_start_idx * per_latent_size + end = (start_rope_start_idx + tt) * per_latent_size + freqs_cos = freqs_cos[start:end, ...] + freqs_sin = freqs_sin[start:end, ...] + + img = self.img_in(img) + + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + if sp_world_size > 1: + sp_size = sp_world_size + sp_rank = rank_in_sp_group + if img.shape[1] % sp_size != 0: + n_token = img.shape[1] + assert n_token > (n_token // sp_size + 1) * ( + sp_size - 1 + ), f"Too short context length for SP {sp_size}" + img = torch.chunk(img, sp_size, dim=1)[sp_rank] + freqs_cos = torch.chunk(freqs_cos, sp_size, dim=0)[sp_rank] + freqs_sin = torch.chunk(freqs_sin, sp_size, dim=0)[sp_rank] + + viewmats = torch.chunk(viewmats, sp_size, dim=1)[sp_rank] + Ks = torch.chunk(Ks, sp_size, dim=1)[sp_rank] + + action = action.reshape(img.shape[0], -1) + t = t.reshape(img.shape[0], -1) + action = torch.chunk(action, sp_size, dim=1)[sp_rank] + t = torch.chunk(t, sp_size, dim=1)[sp_rank] + action = action.reshape(-1) + t = t.reshape(-1) + else: + action = action.reshape(-1) + t = t.reshape(-1) + + # Prepare modulation vectors + vec = self.time_in(t) + + vec = vec + self.action_in(action) # 添加离散的action + + # print('text embedding shape:', txt.shape, text_mask.shape) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + if cache_vision: + _kv_cache_now = [] + transformer_num_layers = len(self.double_blocks) + for i in range(transformer_num_layers): + _kv_cache_now.append( + { + "k_vision": None, + "v_vision": None, + "k_txt": kv_cache[i]["k_txt"], + "v_txt": kv_cache[i]["v_txt"], + } + ) + + # Pass through double-stream blocks + for index, block in enumerate(self.double_blocks): + self.attn_param["layer-name"] = f"double_block_{index + 1}" + + input_dict = { + "img": img, + "vec": vec, + "freqs_cis": freqs_cis, + "attn_param": self.attn_param, + "block_idx": index, + "viewmats": viewmats, + "Ks": Ks, + "kv_cache": kv_cache, + "cache_vision": cache_vision, + } + if torch.is_grad_enabled(): + # print('ar require_grad:', img.requires_grad) + img, vison_kv = torch.utils.checkpoint.checkpoint( + block, + False, + input_dict, + use_reentrant=False, + ) + else: + img, vison_kv = block( + txt_branch=False, + input_dict=input_dict, + ) + + if cache_vision: + _kv_cache_now[index]["k_vision"] = vison_kv["k_vision"] + _kv_cache_now[index]["v_vision"] = vison_kv["v_vision"] + + if cache_vision: + return _kv_cache_now + # Final Layer + img = self.final_layer(img, vec) + if sp_world_size > 1: + img = sequence_model_parallel_all_gather(img, dim=1) + img = self.unpatchify(img, tt, th, tw) + assert return_dict is False, "return_dict is not supported." + features_list = None + return (img, features_list) + + def forward( + self, + txt_branch=False, + input_dict=None, + ): + if txt_branch: + return self.forward_txt(**input_dict) + else: + return self.forward_vision(**input_dict) + + def unpatchify(self, x, t, h, w): + """Unpatchify a tensorized input back to frame format. + + Args: + x (Tensor): Input tensor of shape (N, T, patch_size**2 * C) + t (int): Number of time steps + h (int): Height in patch units + w (int): Width in patch units + + Returns: + Tensor: Output tensor of shape (N, C, t * pt, h * ph, w * pw) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + return imgs + + def set_attn_mode(self, attn_mode: str): + self.attn_mode = attn_mode + for block in self.double_blocks: + block.attn_mode = attn_mode + for block in self.single_blocks: + block.attn_mode = attn_mode diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/causal_mask_visualization.png b/worldcompass/fastvideo/models/hyvideo/models/transformers/causal_mask_visualization.png new file mode 100644 index 0000000000000000000000000000000000000000..85a93a4ac794dd0999f1a75f81be792420dea9b0 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/causal_mask_visualization.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2719da6c4d14ff7b7563d0ee7657863cab309faeec64f50e99bc9193604c07db +size 103131 diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/activation_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/activation_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..8d699cc92c8da1aab26cff6277054d7a4449b211 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/activation_layers.py @@ -0,0 +1,39 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch.nn as nn + + +def get_activation_layer(act_type): + """Get activation layer. + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/attention.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..91a34b8f9bbff03f54eb417f35485903c680fc07 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/attention.py @@ -0,0 +1,584 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import einops +import torch +from typing import Optional +from loguru import logger +import numpy as np +import torch.nn.functional as F +from einops import repeat + +from fastvideo.models.hyvideo.utils.flash_attn_no_pad import ( + flash_attn_no_pad, + flash_attn_no_pad_v3, +) +from fastvideo.models.hyvideo.commons import maybe_fallback_attn_mode + +try: + from torch.nn.attention.flex_attention import flex_attention + + flex_attention = torch.compile(flex_attention, dynamic=False) + torch._dynamo.config.cache_size_limit = 192 + torch._dynamo.config.accumulated_cache_size_limit = 192 + flex_mask_cache = {} +except Exception: + logger.warning("Could not load Sliding Tile Attention of FlexAttn.") + +from fastvideo.models.hyvideo.models.transformers.modules.ssta_attention import ( + ssta_3d_attention, +) +from fastvideo.models.hyvideo.commons.infer_state import get_infer_state + +from fastvideo.distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_world_size, +) +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather, + sequence_model_parallel_all_to_all_4D, +) + + +# @torch.compiler.disable +def attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + drop_rate: float = 0.0, + attn_mask: Optional[torch.Tensor] = None, + causal: bool = False, + attn_mode: str = "flash", +) -> torch.Tensor: + """Compute attention using flash_attn_no_pad or torch scaled_dot_product_attention. + + Args: + q: Query tensor of shape [B, L, H, D] + k: Key tensor of shape [B, L, H, D] + v: Value tensor of shape [B, L, H, D] + drop_rate: Dropout rate for attention weights. + attn_mask: Optional attention mask of shape [B, L]. + causal: Whether to apply causal masking. + attn_mode: Attention mode, either "flash" or "torch". Defaults to "flash". + + Returns: + Output tensor after attention of shape [B, L, H*D] + """ + attn_mode = maybe_fallback_attn_mode(attn_mode) + + if attn_mode == "torch": + # transpose q,k,v dim to fit scaled_dot_product_attention + query = q.transpose(1, 2) # B * H * L * D + key = k.transpose(1, 2) # B * H * L * D + value = v.transpose(1, 2) # B * H * L * D + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [ + torch.int64, + torch.int32, + ]: + assert ( + attn_mask.max() <= 1 and attn_mask.min() >= 0 + ), f"Integer attention mask must be between 0 and 1 for torch attention." + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError( + f"Float attention mask is not implemented for torch attention." + ) + attn_mask1 = einops.rearrange(attn_mask, "b l -> b 1 l 1") + attn_mask2 = einops.rearrange(attn_mask1, "b 1 l 1 -> b 1 1 l") + attn_mask = attn_mask1 & attn_mask2 + + x = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=drop_rate, + is_causal=causal, + ) + + # transpose back + x = x.transpose(1, 2) # B * L * H * D + b, s, h, d = x.shape + out = x.reshape(b, s, -1) + return out + else: + # flash mode (default) + qkv = torch.stack([q, k, v], dim=2) + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + x = flash_attn_no_pad( + qkv, + attn_mask, + causal=causal, + dropout_p=drop_rate, + softmax_scale=None, + ) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +# @torch.compiler.disable +def sequence_parallel_attention_txt( + q, + k, + v, + img_q_len, + img_kv_len, + attn_mode=None, + text_mask=None, + attn_param=None, + block_idx=None, + kv_cache=None, + cache_txt=False, +): + encoder_query = q + encoder_key = k + encoder_value = v + + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + + if sp_world_size > 1: + sp_size = sp_world_size + sp_rank = rank_in_sp_group + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // sp_size + return encoder_state.narrow(dim, sp_rank * local_heads, local_heads) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + encoder_query = encoder_query.transpose(1, 2) + encoder_key = encoder_key.transpose(1, 2) + encoder_value = encoder_value.transpose(1, 2) + + t_kv_cache = {} + if cache_txt: + t_kv_cache["k_txt"] = encoder_key + t_kv_cache["v_txt"] = encoder_value + + encoder_hidden_states = F.scaled_dot_product_attention( + encoder_query, + encoder_key, + encoder_value, + dropout_p=0.0, + is_causal=False, + ) + + # transpose back + encoder_hidden_states = encoder_hidden_states.transpose( + 1, 2 + ) # [B, S, H, D] + + if sp_world_size > 1: + encoder_hidden_states = sequence_model_parallel_all_gather( + encoder_hidden_states, dim=2 + ) + encoder_hidden_states = encoder_hidden_states.to(q.dtype) + + b, s, a, d = encoder_hidden_states.shape + encoder_hidden_states = encoder_hidden_states.reshape(b, s, -1) + + return encoder_hidden_states, t_kv_cache + + +# @torch.compiler.disable +def sequence_parallel_attention_vision( + q, + k, + v, + img_q_len, + img_kv_len, + attn_mode=None, + attn_param=None, + block_idx=None, + kv_cache=None, + cache_vision=False, +): + assert kv_cache is not None + query, query_prope = q + key, key_prope = k + value, value_prope = v + + query = torch.cat([query, query_prope], dim=0) + key = torch.cat([key, key_prope], dim=0) + value = torch.cat([value, value_prope], dim=0) + + sp_world_size = get_sp_world_size() + if sp_world_size > 1: + # batch_size, seq_len, attn_heads, head_dim + query = sequence_model_parallel_all_to_all_4D( + query, scatter_dim=2, gather_dim=1 + ) + key = sequence_model_parallel_all_to_all_4D( + key, scatter_dim=2, gather_dim=1 + ) + value = sequence_model_parallel_all_to_all_4D( + value, scatter_dim=2, gather_dim=1 + ) + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + cache_vision_key = kv_cache[block_idx]["k_vision"] # previous key + cache_vision_value = kv_cache[block_idx]["v_vision"] # previous value + + vision_kv_cache = {} + if cache_vision: + vision_kv_cache["k_vision"] = key + vision_kv_cache["v_vision"] = value + + if not cache_vision and cache_vision_key is not None: + key = torch.cat([cache_vision_key, key], dim=2) + value = torch.cat([cache_vision_value, value], dim=2) + + encoder_key = kv_cache[block_idx]["k_txt"] + encoder_value = kv_cache[block_idx]["v_txt"] + encoder_key = repeat(encoder_key, "B H S D->(B R) H S D", R=2) + encoder_value = repeat(encoder_value, "B H S D->(B R) H S D", R=2) + + key = torch.cat([encoder_key, key], dim=2) + value = torch.cat([encoder_value, value], dim=2) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, dropout_p=0.0, is_causal=False + ) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) # [B, S, H, D] + + if sp_world_size > 1: + hidden_states = sequence_model_parallel_all_to_all_4D( + hidden_states, scatter_dim=1, gather_dim=2 + ) + hidden_states = hidden_states.to(query.dtype) + + b, s, a, d = hidden_states.shape + hidden_states = hidden_states.reshape(b, s, -1) + hidden_states, hidden_states_prope = torch.chunk( + hidden_states, chunks=2, dim=0 + ) + + return hidden_states, hidden_states_prope, vision_kv_cache + + +# @torch.compiler.disable +def parallel_attention( + q, + k, + v, + img_q_len, + img_kv_len, + attn_mode=None, + text_mask=None, + attn_param=None, + block_idx=None, +): + return sequence_parallel_attention( + q, + k, + v, + img_q_len, + img_kv_len, + attn_mode, + text_mask, + attn_param=attn_param, + block_idx=block_idx, + ) + + +def sequence_parallel_attention( + q, + k, + v, + img_q_len, + img_kv_len, + attn_mode=None, + text_mask=None, + attn_param=None, + block_idx=None, +): + assert attn_mode is not None + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + + if sp_world_size > 1: + sp_size = sp_world_size + sp_rank = rank_in_sp_group + + # batch_size, seq_len, attn_heads, head_dim + query = sequence_model_parallel_all_to_all_4D( + query, scatter_dim=2, gather_dim=1 + ) + key = sequence_model_parallel_all_to_all_4D( + key, scatter_dim=2, gather_dim=1 + ) + value = sequence_model_parallel_all_to_all_4D( + value, scatter_dim=2, gather_dim=1 + ) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // sp_size + return encoder_state.narrow(dim, sp_rank * local_heads, local_heads) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + attn_mode = maybe_fallback_attn_mode( + attn_mode, get_infer_state(), block_idx + ) + + if attn_mode == "sageattn": + from sageattention import sageattn + + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + hidden_states = sageattn( + query, key, value, tensor_layout="NHD", is_causal=False + ) + elif attn_mode == "torch": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + if text_mask is not None: + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + else: + attn_mask = None + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [ + torch.int64, + torch.int32, + ]: + assert ( + attn_mask.max() <= 1 and attn_mask.min() >= 0 + ), f"Integer attention mask must be between 0 and 1 for torch attention." + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError( + f"Float attention mask is not implemented for torch attention." + ) + + # transpose q,k,v dim to fit scaled_dot_product_attention + query = query.transpose(1, 2) # B * Head_num * length * dim + key = key.transpose(1, 2) # B * Head_num * length * dim + value = value.transpose(1, 2) # B * Head_num * length * dim + if attn_mask is not None: + attn_mask1 = einops.rearrange(attn_mask, "b l -> b 1 l 1") + attn_mask2 = einops.rearrange(attn_mask1, "b 1 l 1 -> b 1 1 l") + attn_mask = attn_mask1 & attn_mask2 + hidden_states = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + ) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) + + # add new attn for chunk-wise attn + elif attn_mode == "torch_causal": # now: we set text_mask = None + # attention: here we concat the encoder text sequence first, then apply causal attention + vision_seq_length = query.shape[1] + text_seq_length = encoder_query.shape[1] + total_seq_length = vision_seq_length + text_seq_length + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + # prepare causal mask for chunk-wise attention + # set for hunyuanvideo 1.5, which is for 480 * 832 resolution + latent_seq_length = attn_param["thw"][-1] * attn_param["thw"][-2] + chunk_seq_length = latent_seq_length * 4 + chunk_num = (vision_seq_length) // chunk_seq_length + causal_mask = torch.zeros( + (total_seq_length, total_seq_length), device=query.device + ) + causal_mask[:, :text_seq_length] = 1 # no attention for the rest + for i in range(chunk_num): + start_i = text_seq_length + i * chunk_seq_length + end_i = min(start_i + chunk_seq_length, total_seq_length) + for j in range(i + 1): + start_j = text_seq_length + j * chunk_seq_length + end_j = min(start_j + chunk_seq_length, total_seq_length) + # full attention within chunk i for j == i, causal for j < i + causal_mask[start_i:end_i, start_j:end_j] = 1 + + causal_mask = causal_mask.unsqueeze(0).unsqueeze(1) # 1, 1, S, S + causal_mask = causal_mask.expand(query.shape[0], 1, -1, -1) + causal_mask = causal_mask.to(torch.bool) # Force bool dtype + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + hidden_states = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=causal_mask, + dropout_p=0.0, + is_causal=False, + ) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) # [B, S, H, D] + + # return back to the original order: [query, encoder_query] + hidden_states, encoder_hidden_states = ( + hidden_states[:, text_seq_length:, :, :], + hidden_states[:, :text_seq_length, :, :], + ) + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + elif attn_mode == "flash2": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad( + qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None + ) + + elif attn_mode == "flash3": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad_v3( + qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None + ) + + elif attn_mode == "flex-block-attn": + sparse_type = attn_param["attn_sparse_type"] # sta/block_attn/ssta + ssta_threshold = attn_param["ssta_threshold"] + ssta_lambda = attn_param["ssta_lambda"] + ssta_sampling_type = attn_param["ssta_sampling_type"] + ssta_adaptive_pool = attn_param["ssta_adaptive_pool"] + + attn_pad_type = attn_param["attn_pad_type"] # repeat/zero + attn_use_text_mask = attn_param["attn_use_text_mask"] + attn_mask_share_within_head = attn_param["attn_mask_share_within_head"] + + ssta_topk = attn_param["ssta_topk"] + thw = attn_param["thw"] + tile_size = attn_param["tile_size"] + win_size = attn_param["win_size"][0].copy() + + def get_image_tile(tile_size): + block_size = np.prod(tile_size) + if block_size == 384: + tile_size = (1, 16, 24) + elif block_size == 128: + tile_size = (1, 16, 8) + elif block_size == 64: + tile_size = (1, 8, 8) + elif block_size == 16: + tile_size = (1, 4, 4) + else: + raise ValueError( + f"Error tile_size {tile_size}, only support in [16, 64, 128, 384]" + ) + return tile_size + + if thw[0] == 1: + tile_size = get_image_tile(tile_size) + win_size = [1, 1, 1] + elif ( + thw[0] <= 31 + ): # 16fps: 5 * 16 / 4 + 1 = 21; 24fps: 5 * 24 / 4 + 1 = 31 + ssta_topk = ssta_topk // 2 + + # Concatenate and permute query, key, value to (B, H, S, D) + query = torch.cat([query, encoder_query], dim=1).permute(0, 2, 1, 3) + key = torch.cat([key, encoder_key], dim=1).permute(0, 2, 1, 3) + value = torch.cat([value, encoder_value], dim=1).permute(0, 2, 1, 3) + + assert ( + query.shape[-1] == 128 + ), "The last dimension of query, key and value must be 128 for flex-block-attn." + + hidden_states = ssta_3d_attention( + query, + key, + value, + thw, + topk=ssta_topk, + tile_thw=tile_size, + kernel_thw=win_size, + text_len=encoder_sequence_length, + sparse_type=sparse_type, + threshold=ssta_threshold, + lambda_=ssta_lambda, + pad_type=attn_pad_type, + text_mask=text_mask if attn_use_text_mask else None, + sampling_type=ssta_sampling_type, + adaptive_pool=ssta_adaptive_pool, + mask_share_within_head=attn_mask_share_within_head, + ) + hidden_states, sparse_ratio = hidden_states + hidden_states = hidden_states.permute(0, 2, 1, 3) + + else: + raise NotImplementedError( + f"Unsupported attention mode: {attn_mode}. " + f"Only torch, flash, flash3, sageattn and flex-block-attn are supported." + ) + + if sp_world_size > 1: + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length, encoder_sequence_length), dim=1 + ) + hidden_states = sequence_model_parallel_all_to_all_4D( + hidden_states, scatter_dim=1, gather_dim=2 + ) + encoder_hidden_states = sequence_model_parallel_all_gather( + encoder_hidden_states, dim=2 + ).contiguous() + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = hidden_states.shape + hidden_states = hidden_states.reshape(b, s, -1) + + return hidden_states diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/embed_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/embed_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..c5b9a071efe952bfeb26fc3a0ad8539953bc0900 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/embed_layers.py @@ -0,0 +1,228 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import torch.nn as nn + +from fastvideo.models.hyvideo.commons import to_2tuple + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding. + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in + https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + is_reshape_temporal_channels=False, + concat_condition=True, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + # Only support concat mode (multitask mask training) + orig_in_chans = in_chans + if concat_condition: + if is_reshape_temporal_channels: + in_chans = in_chans + in_chans // 2 + 1 + else: + in_chans = in_chans * 2 + 1 + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + + nn.init.xavier_uniform_( + self.proj.weight[:, :orig_in_chans].view( + self.proj.weight[:, :orig_in_chans].size(0), -1 + ) + ) + # Special initialization for concat mode + nn.init.zeros_( + self.proj.weight[:, orig_in_chans:].view( + self.proj.weight[:, orig_in_chans:].size(0), -1 + ) + ) + + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from + https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__( + self, in_channels, hidden_size, act_layer, dtype=None, device=None + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class VisionProjection(torch.nn.Module): + def __init__(self, input_dim, output_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + torch.nn.Linear(input_dim, input_dim), + torch.nn.GELU(), + torch.nn.Linear(input_dim, output_dim), + torch.nn.LayerNorm(output_dim), + ) + + def forward(self, vision_embeds): + return self.proj(vision_embeds) + + +class ClipVisionProjection(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.up = nn.Linear(in_channels, out_channels * 3) + self.down = nn.Linear(out_channels * 3, out_channels) + torch.nn.init.zeros_(self.down.weight) + torch.nn.init.zeros_(self.down.bias) + + def forward(self, x): + projected_x = self.down(nn.functional.silu(self.up(x))) + return projected_x + + +def timestep_embedding(t, dim, max_period=10000): + """Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + +class TimestepEmbedder(nn.Module): + """Embeds scalar timesteps into vector representations.""" + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear( + frequency_embedding_size, + hidden_size, + bias=True, + **factory_kwargs, + ), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding( + t, self.frequency_embedding_size, self.max_period + ).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/mlp_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/mlp_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..1a3f2551a533fec28d9bd412ed52a6c706d71c2e --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/mlp_layers.py @@ -0,0 +1,167 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from fastvideo.models.hyvideo.commons import to_2tuple +from .modulate_layers import modulate + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks.""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = ( + partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + ) + + self.fc1 = linear_layer( + in_channels, hidden_channels, bias=bias[0], **factory_kwargs + ) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = ( + norm_layer(hidden_channels, **factory_kwargs) + if norm_layer is not None + else nn.Identity() + ) + self.fc2 = linear_layer( + hidden_channels, out_features, bias=bias[1], **factory_kwargs + ) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class LinearWarpforSingle(nn.Module): + def __init__( + self, in_dim: int, out_dim: int, bias=False, device=None, dtype=None + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.fc = nn.Linear(in_dim, out_dim, bias=bias, **factory_kwargs) + + def forward(self, x, y): + input = torch.cat([x.contiguous(), y.contiguous()], dim=2).contiguous() + return self.fc(input) + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear( + in_dim, hidden_dim, bias=True, **factory_kwargs + ) + self.silu = nn.SiLU() + self.out_layer = nn.Linear( + hidden_dim, hidden_dim, bias=True, **factory_kwargs + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__( + self, + hidden_size, + patch_size, + out_channels, + act_layer, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs, + ) + else: + out_size = ( + patch_size[0] * patch_size[1] * patch_size[2] + if len(patch_size) == 3 + else patch_size[0] * patch_size[1] + ) * out_channels + self.linear = nn.Linear( + hidden_size, + out_size, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear( + hidden_size, 2 * hidden_size, bias=True, **factory_kwargs + ), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/modulate_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/modulate_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..145f05d386ee2665ef9ff17b1c0a6b45b6c8de9d --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/modulate_layers.py @@ -0,0 +1,115 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear( + hidden_size, factor * hidden_size, bias=True, **factory_kwargs + ) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """Modulate by shift and scale. + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + scale = scale.unsqueeze(0) + latent_length = scale.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + # operate on the hidden_states + scale = scale.repeat_interleave(token_length, dim=1).type_as(x) + + return x * (1 + scale) + elif scale is None: + shift = shift.unsqueeze(0) + latent_length = shift.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + # operate on the hidden_states + shift = shift.repeat_interleave(token_length, dim=1).type_as(x) + return x + shift + else: + shift = shift.unsqueeze(0) + scale = scale.unsqueeze(0) + latent_length = shift.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + + scale = scale.repeat_interleave(token_length, dim=1).type_as(x) + shift = shift.repeat_interleave(token_length, dim=1).type_as(x) + return x * (1 + scale) + shift + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate. + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + gate = gate.unsqueeze(0) + latent_length = gate.shape[1] // x.shape[0] # latent length + token_length = x.shape[1] // latent_length + gate = gate.repeat_interleave(token_length, dim=1).type_as(x) + if tanh: + return x * gate.tanh() + else: + return x * gate + + +def ckpt_wrapper(module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/norm_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/norm_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..9bf1a28ed214430ae235d20f1870b05db5494a84 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/norm_layers.py @@ -0,0 +1,90 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def reset_parameters(self): + if hasattr(self, "weight"): + self.weight.fill_(1) + + def forward(self, x): + """Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/posemb_layers.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/posemb_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..68bc6d506cc826db3dfc05ac40d1e5c4d210e6ec --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/posemb_layers.py @@ -0,0 +1,336 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from typing import Union, Tuple, List + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == 1 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == 1 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = ( + x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) + ) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex( + xq.float().reshape(*xq.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to( + xq.device + ) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex( + xk.float().reshape(*xk.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd( + start, *args, dim=len(rope_dim_list) + ) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance( + theta_rescale_factor, float + ): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif ( + isinstance(theta_rescale_factor, list) + and len(theta_rescale_factor) == 1 + ): + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance( + interpolation_factor, float + ): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif ( + isinstance(interpolation_factor, list) + and len(interpolation_factor) == 1 + ): + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Precompute the frequency tensor for complex exponential (cis) with given dimensions. (Note: + `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor ** (dim / (dim - 2)) + + freqs = 1.0 / ( + theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim) + ) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar( + torch.ones_like(freqs), freqs + ) # complex64 # [S, D/2] + return freqs_cis diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/ssta_attention.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/ssta_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..0ad8894aab8b077c90843d8c0f223a2b162613c0 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/ssta_attention.py @@ -0,0 +1,885 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import numpy as np +from einops import rearrange +from functools import lru_cache + + +def tile(x, canvas_thw, tile_thw, sp_size=1): + r"""Rearrange tensor into tiles for block-based attention. + + Args: + x: Input tensor with shape (b, head, s, d) where s = t * h * w + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Rearranged tensor organized by tiles + """ + b, h, s, d = x.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} == s:{s}" + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t / tile_t_dim) + n_h = int(h / tile_h_dim) + n_w = int(w / tile_w_dim) + x = rearrange( + x, + "b head (sp t h w) d -> b head (t sp h w) d", + sp=sp_size, + t=t // sp_size, + h=h, + w=w, + ) + return rearrange( + x, + "b h (n_t ts_t n_h ts_h n_w ts_w) d -> b h (n_t n_h n_w ts_t ts_h ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim, + ) + + +def untile(x, canvas_thw, tile_thw, sp_size=1): + r"""Reverse the tiling operation to restore original tensor layout. + + Args: + x: Tiled tensor + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Restored tensor with original layout + """ + t, h, w = canvas_thw + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t / tile_t_dim) + n_h = int(h / tile_h_dim) + n_w = int(w / tile_w_dim) + + x = rearrange( + x, + "b h (n_t n_h n_w ts_t ts_h ts_w) d -> b h (n_t ts_t n_h ts_h n_w ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim, + ) + return rearrange( + x, + "b head (t sp h w) d -> b head (sp t h w) d", + sp=sp_size, + t=t // sp_size, + h=h, + w=w, + ) + + +def get_tile_t_h_w(tile_id, tile_thw_dim): + """Extract temporal, height, and width indices from a flattened tile ID.""" + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw_dim + tile_t = tile_id // (tile_h_dim * tile_w_dim) + tile_h = (tile_id % (tile_h_dim * tile_w_dim)) // tile_w_dim + tile_w = tile_id % tile_w_dim + + return tile_t, tile_h, tile_w + + +def importance_sampling( + q, k, topk, threshold=0.0, lambda_=0.9, adaptive_pool=None +): + r"""Select top-k blocks based on importance scores considering both similarity and redundancy. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Threshold parameter (not implemented) + lambda_: Weight factor balancing similarity and redundancy + adaptive_pool: Adaptive pooling parameter (unused) + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + raise NotImplementedError( + "importance_sampling with threshold not implemented" + ) + + q = q / q.norm(dim=-1, keepdim=True) + k = k / k.norm(dim=-1, keepdim=True) + gate_similarity = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate_similarity = (gate_similarity + 1.0) / 2.0 + gate_unique = torch.einsum("bhsd,bhkd->bhsk", k, k) + gate_unique = (gate_unique + 1.0) / 2.0 + + B, H, K_num, D = k.shape + diag_indices = torch.arange(K_num, device=k.device) + gate_unique[:, :, diag_indices, diag_indices] = torch.nan + + mean_redundancy = torch.nanmean(gate_unique, dim=-2, keepdim=True) + + importance_scores = ( + lambda_ * gate_similarity - (1 - lambda_) * mean_redundancy + ) + + topk = min(topk, importance_scores.size(-1)) + _, top_block_indices = importance_scores.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + + +def similarity_sampling( + q, + k, + topk, + threshold=0.0, + block_num=None, + adaptive_pool=None, + temperature=0.01, +): + r"""Select top-k blocks based on similarity scores between query and key averages. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Cumulative score threshold for dynamic topk selection + block_num: Total number of blocks (unused) + adaptive_pool: Adaptive pooling parameter (unused) + temperature: Temperature scaling for softmax + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate = gate / temperature + gate_ = torch.softmax(gate, dim=-1) + sorted_gate, sorted_indices = torch.sort(gate_, dim=-1, descending=True) + cum_scores = torch.cumsum(sorted_gate, dim=-1) + above_threshold = cum_scores >= threshold + has_any_above = above_threshold.any(dim=-1, keepdim=True) + # Find the first position exceeding threshold, add 1 to get the number of blocks to select + dynamic_topk = above_threshold.int().argmax(dim=-1, keepdim=True) + 1 + dynamic_topk = torch.where( + has_any_above, dynamic_topk, torch.full_like(dynamic_topk, topk) + ) + # Generate fixed-shape indices and pad with top1 indices for parts not meeting threshold + # Limit minimum k value + dynamic_topk = torch.clamp(dynamic_topk, min=8, max=topk) + indices = torch.arange(gate.size(-1), device=gate.device).expand( + gate.size() + ) + mask = (indices < dynamic_topk).int() + # Use top1 to pad to topk to ensure same size + top_block_indices = ( + torch.gather(sorted_indices, -1, indices) * mask + + (1 - mask) * sorted_indices[..., 0:1] + ) + else: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + topk = min(topk, gate.size(-1)) + _, top_block_indices = gate.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + + +def create_moba_3d_mask( + q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + add_text_mask=False, + threshold=0.0, + lambda_=None, + mask_share_within_head=True, + q_block_avg_pool=True, + adaptive_pool=None, + sampling_type=None, +): + r"""Create MOBA (Mixture of Block Attention) 3D mask for sparse attention. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + add_text_mask: Whether to add text mask + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + mask_share_within_head: Whether to share mask across heads + q_block_avg_pool: Whether to apply average pooling to query blocks + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + moba_3d_mask: 3D attention mask with shape (num_heads, block_num, block_num) + """ + seq_len = q.size(2) + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + image_shape = canvas_thw + block_shape = tile_thw + batch_size, num_heads, seq_len, head_dim = k.shape + num_blocks_t = np.ceil(image_shape[0] / block_shape[0]).astype(int) + num_blocks_h = np.ceil(image_shape[1] / block_shape[1]).astype(int) + num_blocks_w = np.ceil(image_shape[2] / block_shape[2]).astype(int) + + def get_block_avg_feat( + x: torch.Tensor, adaptive_pool=None, pooling_type="avg" + ) -> torch.Tensor: + x_block_means = x.view( + batch_size, + num_heads, + num_blocks_t, + num_blocks_h, + num_blocks_w, + block_shape[0], + block_shape[1], + block_shape[2], + head_dim, + ) + if adaptive_pool is not None: + x_block_means = x_block_means.view( + -1, block_shape[0], block_shape[1], block_shape[2], head_dim + ) + x_block_means = x_block_means.permute(0, 4, 1, 2, 3) + if pooling_type == "avg": + global_avg = torch.nn.functional.adaptive_avg_pool3d( + x_block_means, adaptive_pool + ) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape( + batch_size, + num_heads, + -1, + head_dim + * adaptive_pool[0] + * adaptive_pool[1] + * adaptive_pool[2], + ) + x_block_means = global_avg + elif pooling_type == "max": + max_pool = torch.nn.functional.adaptive_max_pool3d( + x_block_means, adaptive_pool + ) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape( + batch_size, + num_heads, + -1, + head_dim + * adaptive_pool[0] + * adaptive_pool[1] + * adaptive_pool[2], + ) + x_block_means = max_pool + elif pooling_type == "mix": + global_avg = torch.nn.functional.adaptive_avg_pool3d( + x_block_means, (1, 1, 1) + ) + global_avg = global_avg / ( + global_avg.norm(dim=1, keepdim=True) + 1e-8 + ) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape( + batch_size, num_heads, -1, head_dim + ) + + max_pool = torch.nn.functional.adaptive_max_pool3d( + x_block_means, adaptive_pool + ) + max_pool = max_pool / ( + max_pool.norm(dim=1, keepdim=True) + 1e-8 + ) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape( + batch_size, + num_heads, + -1, + head_dim + * adaptive_pool[0] + * adaptive_pool[1] + * adaptive_pool[2], + ) + x_block_means = torch.cat([max_pool, global_avg], dim=-1) + else: + raise ValueError( + f"pooling_type={pooling_type} is not Supported" + ) + else: + x_block_means = x_block_means.mean(dim=(-2, -3, -4)).view( + batch_size, num_heads, -1, head_dim + ) + return x_block_means + + if ( + sampling_type == "similarity" and threshold > 0.0 + ) and adaptive_pool is None: + adaptive_pool = (2, 2, 2) + + k_block_means = get_block_avg_feat(k, adaptive_pool) + + if q_block_avg_pool: + q = get_block_avg_feat(q, adaptive_pool) + q = q.type(torch.float32) + k_block_means = k_block_means.type( + torch.float32 + ) # float logit for better gate logit perception + + if mask_share_within_head: + q = q.mean(dim=1, keepdim=True) + k_block_means = k_block_means.mean(dim=1, keepdim=True) + + if sampling_type == "similarity": + top_block_indices = similarity_sampling( + q, + k_block_means, + topk, + threshold, + block_num=block_num, + adaptive_pool=adaptive_pool, + ) + elif sampling_type == "importance": + top_block_indices = importance_sampling( + q, + k_block_means, + topk, + threshold, + lambda_=lambda_, + adaptive_pool=adaptive_pool, + ) + else: + raise NotImplementedError( + f"sampling_type={sampling_type} is not Supported" + ) + + q = q.type_as(k) + # Keep head dimension to dynamically build mask + assert ( + top_block_indices.size(0) == 1 + ), "top_block_indices batch size must be 1" + top_block_indices = top_block_indices.squeeze( + 0 + ) # Remove batch dimension [h, s_block, topk] + + # Create 3D mask template (heads, block_num, block_num) + gate_idx_mask = torch.zeros( + (top_block_indices.size(0), block_num, block_num), + dtype=torch.bool, + device=q.device, + ) + + # Iterate over each attention head using dim=0 + for head_idx in range(top_block_indices.size(0)): + # Build index matrix for each head + head_mask = torch.zeros( + (block_num, block_num), dtype=torch.bool, device=q.device + ) + # Fill indices for current head + head_mask.scatter_( + dim=-1, index=top_block_indices[head_idx], value=True + ) + gate_idx_mask[head_idx] = head_mask + + if text_block_num > 0: + pad_block_num = block_num + text_block_num + moba_3d_mask = torch.full( + (gate_idx_mask.size(0), pad_block_num, pad_block_num), + False, + dtype=torch.bool, + device=gate_idx_mask.device, + ) + moba_3d_mask[:, :block_num, :block_num] = gate_idx_mask + if add_text_mask: + moba_3d_mask[:, :, -text_block_num:] = True + moba_3d_mask[:, -text_block_num:, :] = True + else: + moba_3d_mask = gate_idx_mask + + return moba_3d_mask + + +@lru_cache(maxsize=4096) +def create_sta_3d_mask_optimize(canvas_thw, tile_thw, kernel_thw): + r"""Create optimized STA (Spatio-Temporal Attention) 3D mask using vectorized operations. + + Args: + canvas_thw: String representation of canvas dimensions "t_h_w" + tile_thw: String representation of tile dimensions "t_h_w" + kernel_thw: String representation of kernel dimensions "t_h_w" + + Returns: + block_mask: Boolean mask tensor with shape (block_num, block_num) + """ + canvas_thw = tuple(map(int, canvas_thw.split("_"))) + seq_len = np.prod(canvas_thw) + tile_thw = tuple(map(int, tile_thw.split("_"))) + kernel_thw = tuple(map(int, kernel_thw.split("_"))) + + kernel_t, kernel_h, kernel_w = kernel_thw + + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + block_mask = np.full((block_num, block_num), False, dtype=bool) + tile_thw_num = ( + canvas_thw[0] // tile_thw[0], + canvas_thw[1] // tile_thw[1], + canvas_thw[2] // tile_thw[2], + ) + + i_indices = np.arange(block_num) + j_indices = np.arange(block_num) + + i_grid, j_grid = np.meshgrid(i_indices, j_indices, indexing="ij") + + q_t_tile = i_grid // (tile_thw_num[1] * tile_thw_num[2]) + q_h_tile = (i_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[2] + q_w_tile = i_grid % tile_thw_num[2] + + kv_t_tile = j_grid // (tile_thw_num[1] * tile_thw_num[2]) + kv_h_tile = (j_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[ + 2 + ] + kv_w_tile = j_grid % tile_thw_num[2] + + kernel_center_t = np.clip( + q_t_tile, kernel_t // 2, (tile_thw_num[0] - 1) - kernel_t // 2 + ) + kernel_center_h = np.clip( + q_h_tile, kernel_h // 2, (tile_thw_num[1] - 1) - kernel_h // 2 + ) + kernel_center_w = np.clip( + q_w_tile, kernel_w // 2, (tile_thw_num[2] - 1) - kernel_w // 2 + ) + + time_mask = np.abs(kernel_center_t - kv_t_tile) <= kernel_t // 2 + hori_mask = np.abs(kernel_center_h - kv_h_tile) <= kernel_h // 2 + vert_mask = np.abs(kernel_center_w - kv_w_tile) <= kernel_w // 2 + + block_mask = time_mask & hori_mask & vert_mask + + block_mask = torch.tensor(block_mask, dtype=torch.bool) + return block_mask + + +@torch.no_grad() +def create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num=0): + r"""Create STA (Spatio-Temporal Attention) 3D mask. + + Args: + canvas_thw: Canvas dimensions (t, h, w) + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks to pad + + Returns: + sta_mask: Boolean mask tensor with optional text block padding + """ + block_mask = create_sta_3d_mask_optimize( + "_".join([str(x) for x in canvas_thw]), + "_".join([str(x) for x in tile_thw]), + "_".join([str(x) for x in kernel_thw]), + ) + sta_mask = None + block_num = block_mask.size(0) + if text_block_num > 0: + pad_block_num = block_num + text_block_num + sta_mask = torch.full( + (pad_block_num, pad_block_num), + False, + dtype=torch.bool, + ) + sta_mask[:block_num, :block_num] = block_mask + sta_mask[:, -text_block_num:] = True + sta_mask[-text_block_num:, :] = True + else: + sta_mask = block_mask + return sta_mask + + +@torch.no_grad() +def create_ssta_3d_mask( + q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + threshold=0.0, + lambda_=None, + text_mask=None, + mask_share_within_head=True, + adaptive_pool=None, + sampling_type=None, +): + r"""Create SSTA (Sparse Spatio-Temporal Attention) 3D mask combining STA and MOBA masks. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + ssta_3d_mask: Combined 3D attention mask + """ + sta_3d_mask = create_sta_3d_mask( + canvas_thw, tile_thw, kernel_thw, text_block_num + ) + sta_3d_mask = sta_3d_mask.to(q.device) + + moba_3d_mask = create_moba_3d_mask( + q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num, + threshold=threshold, + lambda_=lambda_, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type, + ) + ssta_3d_mask = torch.logical_or(sta_3d_mask.unsqueeze(0), moba_3d_mask) + assert len(ssta_3d_mask.size()) == 3, "ssta_3d_mask should be 3D" + + # set text block mask + if text_mask is not None: + block_size = np.prod(tile_thw) + seq_len = q.size(2) + block_num = int(seq_len / block_size) + text_mask_index = torch.ceil(text_mask.sum() / block_size).long() + text_mask_index = text_mask_index.clamp(min=1).item() + assert ( + ssta_3d_mask.shape[-1] == ssta_3d_mask.shape[-2] + ), "ssta_3d_mask should be square" + + pad_start_index = block_num + text_mask_index + ssta_3d_mask[:, pad_start_index:, :] = False + ssta_3d_mask[:, :, pad_start_index:] = False + eye_mask = torch.eye( + ssta_3d_mask.shape[1] - pad_start_index, + dtype=torch.bool, + device=ssta_3d_mask.device, + ).unsqueeze(0) + ssta_3d_mask[:, pad_start_index:, pad_start_index:] = ( + ssta_3d_mask[:, pad_start_index:, pad_start_index:] | eye_mask + ) + return ssta_3d_mask + + +def ssta_3d_attention( + all_q, + all_k, + all_v, + canvas_thw, + topk=1, + tile_thw=(6, 8, 8), + kernel_thw=(1, 1, 1), + text_len=0, + sparse_type="ssta", + threshold=0.0, + lambda_=None, + pad_type="zero", + text_mask=None, + mask_share_within_head=True, + sampling_type=None, + adaptive_pool=None, +): + r"""Sparse Spatio-Temporal Attention (SSTA) 3D attention mechanism. + + Args: + all_q: Query tensor with shape (B, H, S, D) + all_k: Key tensor with shape (B, H, S, D) + all_v: Value tensor with shape (B, H, S, D) + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_len: Length of text sequence + sparse_type: Type of sparse attention ('sta', 'block_attn', or 'ssta') + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + pad_type: Padding type ("zero" or "repeat") + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + sampling_type: Type of sampling ("similarity" or "importance") + adaptive_pool: Adaptive pooling size + + Returns: + tuple: (output tensor, sparse_ratio) + - output: Attention output with shape (B, H, S, D) + - sparse_ratio: Ratio of non-zero attention weights + """ + try: + from flex_block_attn import flex_block_attn_func + except ImportError as e: + raise Exception( + "Could not load flex-block-attn. Please install the flex-block-attn package. You can install it via 'pip install flex-block-attn'." + ) from e + assert pad_type in ["zero", "repeat"] + assert sampling_type in ["similarity", "importance"] + assert ( + lambda_ is not None and sampling_type == "importance" + ) or sampling_type == "similarity" + + if text_len > 0: + image_q = all_q[:, :, :-text_len, :] + image_k = all_k[:, :, :-text_len, :] + image_v = all_v[:, :, :-text_len, :] + + text_q = all_q[:, :, -text_len:, :] + text_k = all_k[:, :, -text_len:, :] + text_v = all_v[:, :, -text_len:, :] + else: + image_q = all_q + image_k = all_k + image_v = all_v + + b, hd, s, d = image_q.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} != s:{s}" + tile_t, tile_h, tile_w = tile_thw + block_size = np.prod(tile_thw) + + need_pad = False + if t % tile_t != 0 or h % tile_h != 0 or w % tile_w != 0: + need_pad = True + pad_image_q = image_q.reshape(b, hd, t, h, w, d) + pad_image_k = image_k.reshape(b, hd, t, h, w, d) + pad_image_v = image_v.reshape(b, hd, t, h, w, d) + + pad_t = 0 if t % tile_t == 0 else tile_t - t % tile_t + if pad_t > 0: + t = t + pad_t + repeat_q = pad_image_q[:, :, -1:, :, :, :].expand( + -1, -1, pad_t, -1, -1, -1 + ) + repeat_k = pad_image_k[:, :, -1:, :, :, :].expand( + -1, -1, pad_t, -1, -1, -1 + ) + repeat_v = pad_image_v[:, :, -1:, :, :, :].expand( + -1, -1, pad_t, -1, -1, -1 + ) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=2) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=2) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=2) + + pad_h = 0 if h % tile_h == 0 else tile_h - h % tile_h + if pad_h > 0: + h = h + pad_h + repeat_q = pad_image_q[:, :, :, -1:, :, :].expand( + -1, -1, -1, pad_h, -1, -1 + ) + repeat_k = pad_image_k[:, :, :, -1:, :, :].expand( + -1, -1, -1, pad_h, -1, -1 + ) + repeat_v = pad_image_v[:, :, :, -1:, :, :].expand( + -1, -1, -1, pad_h, -1, -1 + ) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=3) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=3) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=3) + + pad_w = 0 if w % tile_w == 0 else tile_w - w % tile_w + if pad_w > 0: + w = w + pad_w + repeat_q = pad_image_q[:, :, :, :, -1:, :].expand( + -1, -1, -1, -1, pad_w, -1 + ) + repeat_k = pad_image_k[:, :, :, :, -1:, :].expand( + -1, -1, -1, -1, pad_w, -1 + ) + repeat_v = pad_image_v[:, :, :, :, -1:, :].expand( + -1, -1, -1, -1, pad_w, -1 + ) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=4) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=4) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=4) + + image_q = pad_image_q.reshape(b, hd, -1, d) + image_k = pad_image_k.reshape(b, hd, -1, d) + image_v = pad_image_v.reshape(b, hd, -1, d) + + canvas_thw = (t, h, w) + + need_pad_text = False + text_block_num = math.ceil(text_len / block_size) + text_target_size = text_block_num * block_size + if text_len % block_size > 0: + need_pad_text = True + text_pad_size = text_target_size - text_len + + pad_text_q = ( + text_q[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + ) + pad_text_k = ( + text_k[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + ) + pad_text_v = ( + text_v[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + ) + + text_q = torch.cat([text_q, pad_text_q], dim=2) + text_k = torch.cat([text_k, pad_text_k], dim=2) + text_v = torch.cat([text_v, pad_text_v], dim=2) + + image_q = tile(image_q, canvas_thw, tile_thw) + image_k = tile(image_k, canvas_thw, tile_thw) + image_v = tile(image_v, canvas_thw, tile_thw) + + if text_len > 0: + q = torch.cat([image_q, text_q], dim=2) + k = torch.cat([image_k, text_k], dim=2) + v = torch.cat([image_v, text_v], dim=2) + else: + q = image_q + k = image_k + v = image_v + + if sparse_type == "sta": + assert text_mask is None, "text_mask do not support in sta sparse_type" + block_mask = create_sta_3d_mask( + canvas_thw, tile_thw, kernel_thw, text_block_num + ).to(q.device) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + elif sparse_type == "block_attn": + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_moba_3d_mask( + image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + topk=topk, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + add_text_mask=True, + lambda_=lambda_, + threshold=threshold, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type, + ) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + + elif sparse_type == "ssta": + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_ssta_3d_mask( + image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + topk=topk, + threshold=threshold, + lambda_=lambda_, + text_mask=text_mask[i] if text_mask is not None else None, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type, + ) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + if mask_share_within_head: + block_mask = block_mask.unsqueeze(1) # [b, 1, s_block, s_block] + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + else: + raise Exception(f"unsupported sparse_type:{sparse_type}") + sparse_ratio = block_mask.float().mean().cpu().item() + + if text_len > 0: + image_o = o[:, :, :-text_target_size, :] + if need_pad_text: + text_o = o[:, :, -text_target_size:-text_pad_size, :] + else: + text_o = o[:, :, -text_target_size:, :] + else: + image_o = o + + image_o = untile(image_o, canvas_thw, tile_thw) + + if need_pad: + # Remove padding from output + unpad_image_o = image_o.reshape(b, hd, t, h, w, d) + if pad_t > 0: + unpad_image_o = unpad_image_o[:, :, :-pad_t, :, :, :] + if pad_h > 0: + unpad_image_o = unpad_image_o[:, :, :, :-pad_h, :, :] + if pad_w > 0: + unpad_image_o = unpad_image_o[:, :, :, :, :-pad_w, :] + image_o = unpad_image_o.reshape(b, hd, -1, d) + + if text_len > 0: + o = torch.cat([image_o, text_o], dim=2) + else: + o = image_o + + return o, sparse_ratio diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/token_refiner.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/token_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..df323ab999a4abc79302995ff35d963421b78a92 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/token_refiner.py @@ -0,0 +1,305 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer +from .attention import attention + + +class IndividualTokenRefinerBlock(nn.Module): + """A single block for token refinement with self-attention and MLP. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + self.self_attn_qkv = nn.Linear( + hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs + ) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + if qk_norm + else nn.Identity() + ) + self.self_attn_proj = nn.Linear( + hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs + ) + + self.norm2 = nn.LayerNorm( + hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs + ) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear( + hidden_size, 2 * hidden_size, bias=True, **factory_kwargs + ), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass for IndividualTokenRefinerBlock. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + attn_mask: Optional attention mask of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange( + qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num + ) + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + return x + + +class IndividualTokenRefiner(nn.Module): + """Stacks multiple IndividualTokenRefinerBlock modules. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList( + [ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + for _ in range(depth) + ] + ) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass for IndividualTokenRefiner. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + if mask is not None: + mask = mask.clone().bool() + mask[:, 0] = True # Prevent attention weights from becoming NaN + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """Single token refiner block for LLM text embedding refinement. + + Args: + in_channels: Input feature dimension. + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + in_channels: int, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.input_embedder = nn.Linear( + in_channels, hidden_size, bias=True, **factory_kwargs + ) + act_layer = get_activation_layer(act_type) + self.t_embedder = TimestepEmbedder( + hidden_size, act_layer, **factory_kwargs + ) + self.c_embedder = TextProjection( + in_channels, hidden_size, act_layer, **factory_kwargs + ) + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + """Forward pass for SingleTokenRefiner. + + Args: + x: Input tensor of shape [B, L, in_channels]. + t: Timestep tensor of shape [B]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, hidden_size]. + """ + timestep_aware_representations = self.t_embedder(t) + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [B, L, 1] + context_aware_representations = (x * mask_float).sum( + dim=1 + ) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder( + context_aware_representations + ) + c = timestep_aware_representations + context_aware_representations + x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/upsample.py b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/upsample.py new file mode 100644 index 0000000000000000000000000000000000000000..5a7bf0c2ba10993e2e9a57062fb55f8bd4a50213 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/transformers/modules/upsample.py @@ -0,0 +1,182 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from fastvideo.models.hyvideo.models.autoencoders.hunyuanvideo_15_vae_w_cache import ( + CausalConv3d, + ResnetBlock, + RMS_norm, + forward_with_checkpointing, + swish, +) + + +class UpsamplerType(Enum): + LEARNED = "learned" + FIXED = "fixed" + NONE = "none" + LEARNED_FIXED = "learned_fixed" + + +@dataclass +class UpsamplerConfig: + load_from: str + enable: bool = False + hidden_channels: int = 128 + num_blocks: int = 16 + model_type: UpsamplerType = UpsamplerType.NONE + version: str = "720p" + + +class SRResidualCausalBlock3D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.block = nn.Sequential( + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.block(x) + + +class SRTo720pUpsampler(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + hidden_channels: int | None = None, + num_blocks: int = 6, + global_residual: bool = False, + ): + super().__init__() + if hidden_channels is None: + hidden_channels = 64 + self.in_conv = CausalConv3d(in_channels, hidden_channels, kernel_size=3) + self.blocks = nn.ModuleList( + [ + SRResidualCausalBlock3D(hidden_channels) + for _ in range(num_blocks) + ] + ) + self.out_conv = CausalConv3d( + hidden_channels, out_channels, kernel_size=3 + ) + self.global_residual = bool(global_residual) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + y = self.in_conv(x) + for blk in self.blocks: + y = blk(y) + y = self.out_conv(y) + if self.global_residual and (y.shape == residual.shape): + y = y + residual + return y + + +class SRTo1080pUpsampler(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int = 2, + is_residual: bool = False, + ): + super().__init__() + self.num_res_blocks = num_res_blocks + self.block_out_channels = block_out_channels + self.z_channels = z_channels + + block_in = block_out_channels[0] + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append( + ResnetBlock(in_channels=block_in, out_channels=block_out) + ) + block_in = block_out + up = nn.Module() + up.block = block + + self.up.append(up) + + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, out_channels, kernel_size=3) + + self.gradient_checkpointing = False + self.is_residual = is_residual + + def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor: + """ + Args: + z: (B, C, T, H, W) + target_shape: (H, W) + """ + use_checkpointing = bool(self.training and self.gradient_checkpointing) + if target_shape is not None and z.shape[-2:] != target_shape: + bsz = z.shape[0] + z = rearrange(z, "b c f h w -> (b f) c h w") + z = F.interpolate( + z, size=target_shape, mode="bilinear", align_corners=False + ) + z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing( + self.up[i_level].block[i_block], + h, + use_checkpointing=use_checkpointing, + ) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing( + self.up[i_level].upsample, + h, + use_checkpointing=use_checkpointing, + ) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h diff --git a/worldcompass/fastvideo/models/hyvideo/models/vision_encoder/__init__.py b/worldcompass/fastvideo/models/hyvideo/models/vision_encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..abd5b8ab3ccaeb535d17a8a206bc732ddc1c3890 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/models/vision_encoder/__init__.py @@ -0,0 +1,252 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from dataclasses import dataclass +from typing import Optional, Tuple +import torch +import torch.nn as nn +from transformers import SiglipImageProcessor, SiglipVisionModel +from transformers.utils import ModelOutput +import numpy as np + +from fastvideo.models.hyvideo.commons import PRECISION_TO_TYPE + +VISION_ENCODER_PATH = {} + + +def use_default(value, default): + return value if value is not None else default + + +def load_vision_encoder( + vision_encoder_type, + vision_encoder_precision=None, + vision_encoder_path=None, + logger=None, + device=None, +): + if vision_encoder_path is None: + vision_encoder_path = VISION_ENCODER_PATH[vision_encoder_type] + + if vision_encoder_type == "siglip": + vision_encoder = SiglipVisionModel.from_pretrained( + vision_encoder_path, subfolder="image_encoder" + ) + else: + raise ValueError( + f"Unsupported vision encoder type: {vision_encoder_type}" + ) + + # from_pretrained will ensure that the model is in eval mode. + if vision_encoder_precision is not None: + vision_encoder = vision_encoder.to( + dtype=PRECISION_TO_TYPE[vision_encoder_precision] + ) + + vision_encoder.requires_grad_(False) + + if device is not None: + vision_encoder = vision_encoder.to(device) + + return vision_encoder, vision_encoder_path + + +def load_image_processor(processor_type, processor_path=None, logger=None): + if processor_path is None: + processor_path = VISION_ENCODER_PATH[processor_type] + + if processor_type == "siglip": + processor = SiglipImageProcessor.from_pretrained( + processor_path, subfolder="feature_extractor" + ) + else: + raise ValueError(f"Unsupported processor type: {processor_type}") + + return processor, processor_path + + +@dataclass +class VisionEncoderModelOutput(ModelOutput): + """Base class for vision encoder model's outputs. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*): + Last layer hidden-state of the first token of the sequence (classification token) + after further processing through the layers used for the auxiliary pretraining task. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: torch.FloatTensor = None + pooler_output: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + + +class VisionEncoder(nn.Module): + def __init__( + self, + vision_encoder_type: str, + vision_encoder_precision: Optional[str] = None, + vision_encoder_path: Optional[str] = None, + processor_type: Optional[str] = None, + processor_path: Optional[str] = None, + output_key: Optional[str] = None, + logger=None, + device=None, + ): + super().__init__() + self.vision_encoder_type = vision_encoder_type + self.precision = vision_encoder_precision + self.model_path = vision_encoder_path + self.processor_type = ( + processor_type + if processor_type is not None + else vision_encoder_type + ) + self.processor_path = ( + processor_path + if processor_path is not None + else vision_encoder_path + ) + self.logger = logger + + if "siglip" in vision_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError( + f"Unsupported vision encoder type: {vision_encoder_type}" + ) + + self.model, self.model_path = load_vision_encoder( + vision_encoder_type=self.vision_encoder_type, + vision_encoder_precision=self.precision, + vision_encoder_path=self.model_path, + logger=self.logger, + device=device, + ) + self.dtype = self.model.dtype + self.device = self.model.device + + self.processor, self.processor_path = load_image_processor( + processor_type=self.processor_type, + processor_path=self.processor_path, + logger=self.logger, + ) + + def __repr__(self): + return ( + f"{self.vision_encoder_type} ({self.precision} - {self.model_path})" + ) + + def encode_latents_to_images(self, latents, vae, reorg_token=False): + """Convert latents to images using VAE decoder. + + Args: + latents: Input latents tensor + vae: VAE model for decoding + reorg_token: Whether to reorg the token + Returns: + images: Decoded images as numpy array + """ + # Handle both 4D and 5D latents (for video, take first frame) + first_image_latents = ( + latents[:, :, 0, ...] if len(latents.shape) == 5 else latents + ) + first_image_latents = ( + 1 / vae.config.scaling_factor * first_image_latents + ) + first_image = vae.decode( + first_image_latents.unsqueeze(2).to(vae.dtype), return_dict=False + )[0].cpu() + first_image = first_image[:, :, 0, :, :] + first_image = (first_image / 2 + 0.5).clamp(0, 1) + first_image = (first_image * 255.0).clamp(0, 255.0) + first_image = first_image.to(torch.uint8).numpy() + first_image = first_image.transpose(0, 2, 3, 1) + + assert isinstance(first_image, np.ndarray) + assert first_image.ndim == 4 and first_image.shape[3] == 3 + assert first_image.dtype == np.uint8 + + return first_image + + def encode_images(self, images): + """Encode images using the vision encoder. + + Args: + images: Input images (numpy array or preprocessed tensor) + + Returns: + VisionEncoderModelOutput with encoded features + """ + if isinstance(images, np.ndarray): + # Preprocess images if they're numpy arrays + preprocessed = self.processor.preprocess( + images=images, return_tensors="pt" + ).to(device=self.model.device, dtype=self.model.dtype) + else: + # Assume already preprocessed + preprocessed = images + + outputs = self.model(**preprocessed) + + return VisionEncoderModelOutput( + last_hidden_state=outputs.last_hidden_state, + pooler_output=( + outputs.pooler_output + if hasattr(outputs, "pooler_output") + else None + ), + hidden_states=( + outputs.hidden_states + if hasattr(outputs, "hidden_states") + else None + ), + ) + + def encode_latents(self, latents, vae, reorg_token=False): + """Encode latents by first converting to images, then encoding. This is the main function + that replaces sigclip_vision_encode. + + Args: + latents: Input latent tensors + vae: VAE model for decoding latents to images + + Returns: + Encoded image features + """ + # Convert latents to images + images = self.encode_latents_to_images(latents, vae, reorg_token) + + # Encode images + outputs = self.encode_images(images) + + return outputs.last_hidden_state + + def forward(self, images): + """Forward pass for direct image encoding. + + Args: + images: Input images + + Returns: + VisionEncoderModelOutput with encoded features + """ + return self.encode_images(images) diff --git a/worldcompass/fastvideo/models/hyvideo/schedulers/scheduling_flow_match_discrete.py b/worldcompass/fastvideo/models/hyvideo/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..c918dfb1c9549dd1cd2e76e992695fccedabf6df --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,309 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. +# +# ============================================================================== +# Modified from diffusers +# ============================================================================== +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import numpy as np +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.schedulers.scheduling_utils import SchedulerMixin + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + use_flux_shift: bool = False, + flux_base_shift: float = 0.5, + flux_max_shift: float = 1.15, + n_tokens: Optional[int] = None, + flux_base_token=256.0, + flux_max_token=4096.0, + flux_shift_factor=1.0, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to( + dtype=torch.float32 + ) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError( + f"Solver {solver} not supported. Supported solvers: {self.supported_solver}" + ) + + @property + def step_index(self): + """The index counter for current timestep. + + It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """The index for the first timestep. + + It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """Sets the begin index for the scheduler. This function should be run from pipeline before + the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + n_tokens: int = None, + ): + """Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + + # Apply timestep shift + if self.config.use_flux_shift: + assert isinstance( + n_tokens, int + ), "n_tokens should be provided for flux shift" + mu = self.get_lin_function( + x1=self.config.flux_base_token, + x2=self.config.flux_max_token, + y1=self.config.flux_base_shift * self.config.flux_shift_factor, + y2=self.config.flux_max_shift * self.config.flux_shift_factor, + )(n_tokens) + sigmas = self.flux_time_shift(mu, 1.0, sigmas) + elif self.config.shift != 1.0: + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to( + dtype=torch.float32, device=device + ) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input( + self, sample: torch.Tensor, timestep: Optional[int] = None + ) -> torch.Tensor: + return sample + + @staticmethod + def get_lin_function( + x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15 + ): + m = (y2 - y1) / (x2 - x1) + b = y1 - m * x1 + return lambda x: m * x + b + + @staticmethod + def flux_time_shift(mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, + n_tokens: Optional[int] = None, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """Predict the sample from the previous timestep by reversing the SDE. This function + propagates the diffusion process from the learned model outputs (most often the predicted + noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.float() * dt + else: + raise ValueError( + f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}" + ) + + # Cast sample back to model compatible dtype + # prev_sample = prev_sample.to(model_output.dtype) + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/worldcompass/fastvideo/models/hyvideo/utils/communications.py b/worldcompass/fastvideo/models/hyvideo/utils/communications.py new file mode 100644 index 0000000000000000000000000000000000000000..65bd9d73d5569cabdaf859395508247fca98e74f --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/communications.py @@ -0,0 +1,321 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Any, Tuple +import torch +import torch.distributed as dist + +from torch.nn import functional as F + + +def broadcast(input_: torch.Tensor, group: dist.ProcessGroup): + src = dist.get_global_rank(group, 0) + dist.broadcast(input_, src=src, group=group) + + +def _all_to_all_4D( + input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None +) -> torch.tensor: + """All-to-all for QKV. + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert ( + input.dim() == 4 + ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + seq_lens = [None] * seq_world_size + dist.all_gather_object(seq_lens, input.shape[1], group) + # uneven + if seq_lens[-1] != seq_lens[0]: + assert seq_lens[0] > seq_lens[-1] + gap = seq_lens[0] - seq_lens[-1] + if ( + dist.get_group_rank(group, dist.get_rank()) + == seq_world_size - 1 + ): + assert input.shape[1] == seq_lens[-1] + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + else: + gap = 0 + + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + assert ( + hc % seq_world_size == 0 + ), f"Invalid Head size: {hc}, which should be divisible by spsize {seq_world_size}" + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = ( + input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs) + .transpose(0, 2) + .contiguous() + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = ( + output.transpose(0, 1) + .contiguous() + .reshape(bs, seqlen, shard_hc, hs) + ) + if gap > 0: + output = output[:, :-gap] + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + + hc = shard_hc * seq_world_size + if seqlen % seq_world_size != 0: + new_seqlen = (seqlen // seq_world_size + 1) * seq_world_size + gap = new_seqlen - seqlen + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + bs, seqlen, shard_hc, hs = input.shape + else: + gap = 0 + + assert seqlen % seq_world_size == 0 + + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = ( + input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs) + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = ( + output.transpose(0, 2) + .contiguous() + .reshape(bs, shard_seqlen, hc, hs) + ) + + if ( + gap > 0 + and dist.get_group_rank(group, dist.get_rank()) + == seq_world_size - 1 + ): + output = output[:, :-gap] + + return output + else: + raise RuntimeError( + "scatter_idx must be 1 or 2 and gather_idx must be 1 or 2" + ) + + +class SeqAllToAll4D(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: torch.Tensor, + scatter_idx: int, + gather_idx: int, + ) -> torch.Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward( + ctx: Any, *grad_output: torch.Tensor + ) -> Tuple[None, torch.Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply( + ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx + ), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, + group: dist.ProcessGroup, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return SeqAllToAll4D.apply(group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [ + t.contiguous() + for t in torch.tensor_split(input_, world_size, scatter_dim) + ] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all( + input_, ctx.world_size, process_group, scatter_dim, gather_dim + ) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + + +def all_to_all( + input_: torch.Tensor, + group: dist.ProcessGroup, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return _AllToAll.apply(input_, group, scatter_dim, gather_dim) + + +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim, group): + ctx.dim = dim + ctx.group = group + world_size = dist.get_world_size(group) + input_size = list(input_.size()) + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_.shape, group) + + ctx.input_size = input_size[dim] + + tensor_list = [ + torch.empty(sizes[i], dtype=input_.dtype, device=input_.device) + for i in range(world_size) + ] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + group = ctx.group + world_size = dist.get_world_size(group) + global_rank = dist.get_rank() + rank = dist.get_group_rank(group, global_rank) + dim = ctx.dim + input_size = ctx.input_size + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_size, group=group) + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None, None + + +def all_gather(input_: torch.Tensor, dim: int = 1, group=None): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim, group) diff --git a/worldcompass/fastvideo/models/hyvideo/utils/data_utils.py b/worldcompass/fastvideo/models/hyvideo/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..51efed81e3db9f552008bbaef55e55017ca91c77 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/data_utils.py @@ -0,0 +1,77 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import numpy as np +from PIL import Image + + +def resize_and_center_crop(image, target_width, target_height): + if target_height == image.shape[0] and target_width == image.shape[1]: + return image + + pil_image = Image.fromarray(image) + original_width, original_height = pil_image.size + scale_factor = max( + target_width / original_width, target_height / original_height + ) + resized_width = int(round(original_width * scale_factor)) + resized_height = int(round(original_height * scale_factor)) + resized_image = pil_image.resize( + (resized_width, resized_height), Image.LANCZOS + ) + left = (resized_width - target_width) / 2 + top = (resized_height - target_height) / 2 + right = (resized_width + target_width) / 2 + bottom = (resized_height + target_height) / 2 + cropped_image = resized_image.crop((left, top, right, bottom)) + return np.array(cropped_image) + + +def get_closest_ratio(height: float, width: float, ratios: list, buckets: list): + """Get the closest ratio in the buckets. + + Args: + height (float): video height + width (float): video width + ratios (list): video aspect ratio + buckets (list): buckets generated by `generate_crop_size_list` + + Returns: + the closest size in the buckets and the corresponding ratio + """ + aspect_ratio = float(height) / float(width) + + ratios_array = np.array(ratios) + closest_ratio_id = np.abs(ratios_array - aspect_ratio).argmin() + closest_size = buckets[closest_ratio_id] + closest_ratio = ratios_array[closest_ratio_id] + + return closest_size, closest_ratio + + +def generate_crop_size_list(base_size=256, patch_size=16, max_ratio=4.0): + num_patches = round((base_size / patch_size) ** 2) + assert max_ratio >= 1.0 + crop_size_list = [] + wp, hp = num_patches, 1 + while wp > 0: + if max(wp, hp) / min(wp, hp) <= max_ratio: + crop_size_list.append((wp * patch_size, hp * patch_size)) + if (hp + 1) * wp <= num_patches: + hp += 1 + else: + wp -= 1 + return crop_size_list diff --git a/worldcompass/fastvideo/models/hyvideo/utils/flash_attn_no_pad.py b/worldcompass/fastvideo/models/hyvideo/utils/flash_attn_no_pad.py new file mode 100644 index 0000000000000000000000000000000000000000..32e889cd18382bd2b3930035f18080e42aa002a5 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/flash_attn_no_pad.py @@ -0,0 +1,237 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from einops import rearrange +from torch.nn.attention.flex_attention import ( + flex_attention, + create_block_mask, + BlockMask, +) + +# compile first +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune" +) + + +def prepare_blockwise_causal_attn_mask( + device: torch.device | str, + num_frames: int = 21, + frame_seqlen: int = 880, + causal_mask=None, +) -> BlockMask: + """We will divide the token sequence into the following format [1 latent frame] [1 latent frame] + ... + + [1 latent frame] We use flexattention to construct the attention mask + """ + + total_length = num_frames * frame_seqlen + + def attention_mask(b, h, q_idx, kv_idx): + return causal_mask[q_idx, kv_idx] + + +# add new function for computing causal attention with flex attention +def flex_attn_no_pad( + qkv, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False, +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + + batch_size = qkv.shape[0] # qkv shape: [B, total_length, 3, num_head, D] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( + x, key_padding_mask + ) # only select the unmasked tokens, the selected token index, cu_seqlens: (0, unmasked token length), max_s: max unmasked length in the batch + + # ------------------------- for chunk-wise causal attention mask ------------------------- + # get the video sequence + latent_seq_length = ( + 1560 # set for hunyuanvideo 1.5, which is for 480 * 832 resolution + ) + chunk_seq_length = 1560 * 4 + text_seq_length = ( + max_s % chunk_seq_length + ) # including txt, byt5, sigclip vision token + latent_num = (max_s - text_seq_length) // latent_seq_length + chunk_num = (max_s - text_seq_length) // chunk_seq_length + causal_mask = torch.zeros((max_s, max_s), device=qkv.device) + + causal_mask[:text_seq_length, :text_seq_length] = ( + 1 # no attention for the rest + ) + for i in range(chunk_num): + start_i = text_seq_length + i * chunk_seq_length + end_i = min(start_i + chunk_seq_length, max_s) + for j in range(i + 1): + start_j = text_seq_length + j * chunk_seq_length + end_j = min(start_j + chunk_seq_length, max_s) + # full attention within chunk i for j == i, causal for j < i + causal_mask[start_i:end_i, start_j:end_j] = 1 + causal_mask = causal_mask.to(torch.bool) # Force bool dtype + + block_mask = prepare_blockwise_causal_attn_mask( + device=qkv.device, + num_frames=latent_num, + frame_seqlen=latent_seq_length, + causal_mask=causal_mask, + ) + + q, k, v = qkv.chunk(3, dim=2) + q = q.squeeze(2).transpose(1, 2) + k = k.squeeze(2).transpose(1, 2) + v = v.squeeze(2).transpose(1, 2) + output_unpad = flex_attention( + q, k, v, block_mask=block_mask + ) # output_unpad: [B, num_head, L, D] + + # bidirectional + # x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + # output_unpad = flash_attn_varlen_qkvpacked_func( + # x_unpad, + # cu_seqlens, + # max_s, + # dropout_p, + # softmax_scale=softmax_scale, + # causal=causal, + # deterministic=deterministic, + # ) + + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), + indices, + batch_size, + seqlen, + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output + + +# add new function for computing causal attention with flex attention +def flash_attn_no_pad( + qkv, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False, +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( + x, key_padding_mask + ) + + x_unpad = rearrange( + x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads + ) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), + indices, + batch_size, + seqlen, + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output + + +def flash_attn_no_pad_v3( + qkv, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False, +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + from flash_attn_interface import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, + ) + + if flash_attn_varlen_func_v3 is None: + raise ImportError("FlashAttention V3 backend not available") + + batch_size, seqlen, _, nheads, head_dim = qkv.shape + query, key, value = qkv.unbind(dim=2) + + query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + rearrange(query, "b s h d -> b s (h d)"), key_padding_mask + ) + key_unpad, _, cu_seqlens_k, _, _ = unpad_input( + rearrange(key, "b s h d -> b s (h d)"), key_padding_mask + ) + value_unpad, _, _, _, _ = unpad_input( + rearrange(value, "b s h d -> b s (h d)"), key_padding_mask + ) + + query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads) + key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads) + value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads) + + output_unpad = flash_attn_varlen_func_v3( + query_unpad, + key_unpad, + value_unpad, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_q, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), + indices, + batch_size, + seqlen, + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output diff --git a/worldcompass/fastvideo/models/hyvideo/utils/infer_utils.py b/worldcompass/fastvideo/models/hyvideo/utils/infer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0852b889b755bc3a41bee4a2eb4451899a3e90df --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/infer_utils.py @@ -0,0 +1,35 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch + +from fastvideo.models.hyvideo.commons.infer_state import get_infer_state + + +def torch_compile_wrapper(): + """返回一个装饰器，延迟决定是否使用torch.compile.""" + + def decorator(func): + def wrapper(*args, **kwargs): + if get_infer_state() and get_infer_state().enable_torch_compile: + compiled_func = torch.compile(func) + return compiled_func(*args, **kwargs) + else: + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/worldcompass/fastvideo/models/hyvideo/utils/multitask_utils.py b/worldcompass/fastvideo/models/hyvideo/utils/multitask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..224c38ba9c0d285a645a1d96e1b00a889ef66e9b --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/multitask_utils.py @@ -0,0 +1,57 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from typing import List +import numpy as np +from PIL import Image + + +def numpy_to_pil(images: np.ndarray) -> List[Image.Image]: + """Convert a numpy image or a batch of images to a PIL image. + + Args: + images (np.ndarray): The image array to convert to PIL format. + + Returns: + List[Image.Image]: A list of PIL images. + """ + if images.ndim == 3: + images = images[None, ...] + images = (images * 255).round().astype("uint8") + if images.shape[-1] == 1: + # special case for grayscale (single channel) images + pil_images = [ + Image.fromarray(image.squeeze(), mode="L") for image in images + ] + else: + pil_images = [Image.fromarray(image) for image in images] + + return pil_images + + +def merge_tensor_by_mask(tensor_1, tensor_2, mask, dim): + assert tensor_1.shape == tensor_2.shape + # Mask is a 0/1 vector. Choose tensor_2 when the value is 1; otherwise, tensor_1 + masked_indices = torch.nonzero(mask).squeeze(1) + tmp = tensor_1.clone() + if dim == 0: + tmp[masked_indices] = tensor_2[masked_indices] + elif dim == 1: + tmp[:, masked_indices] = tensor_2[:, masked_indices] + elif dim == 2: + tmp[:, :, masked_indices] = tensor_2[:, :, masked_indices] + return tmp diff --git a/worldcompass/fastvideo/models/hyvideo/utils/rewrite/clients.py b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/clients.py new file mode 100644 index 0000000000000000000000000000000000000000..9796397c1755a7f3b1cfa731ee8ad62269eb2d0c --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/clients.py @@ -0,0 +1,281 @@ +# -*- coding: utf-8 -*- +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import json +import time +import ast +import os +import io +import base64 + +import openai +from loguru import logger + + +class NonStreamResponse(object): + def __init__(self): + self.response = "" + + def _deserialize(self, obj): + self.response = json.dumps(obj) + + +# DeepSeekClient +class DeepSeekClient(object): + def __init__(self, key_id, key_secret): + from tencentcloud.common.common_client import CommonClient + from tencentcloud.common import credential + from tencentcloud.common.profile.client_profile import ClientProfile + from tencentcloud.common.profile.http_profile import HttpProfile + + cred = credential.Credential(key_id, key_secret) + httpProfile = HttpProfile() + httpProfile.endpoint = "lkeap.tencentcloudapi.com" + httpProfile.reqTimeout = ( + 40000 # The streaming interface may take a longer time. + ) + clientProfile = ClientProfile() + clientProfile.httpProfile = httpProfile + self.common_client = CommonClient( + "lkeap", "2024-05-22", cred, "ap-guangzhou", profile=clientProfile + ) + + def run_single_recaption(self, system_prompt, input_prompt): + post_dict = { + "Model": "deepseek-v3.1", + "Messages": [ + {"Role": "system", "Content": system_prompt}, + {"Role": "user", "Content": input_prompt}, + ], + "Stream": False, + "Thinking": {"Type": "enabled"}, + } + while True: + try: + resp = self.common_client._call_and_deserialize( + "ChatCompletions", post_dict, NonStreamResponse + ) + break + except Exception as e: + logger.error(e) + time.sleep(1) + resp = self.common_client._call_and_deserialize( + "ChatCompletions", post_dict, NonStreamResponse + ) + response = resp.response + response = ast.literal_eval(response) + content = response["Choices"][0]["Message"]["Content"] + reasoning_content = response["Choices"][0]["Message"][ + "ReasoningContent" + ] + print("Initial prompt: ", input_prompt) + print("Recaption prompt: ", content) + + return content, reasoning_content + + +class QwenClient(object): + def __init__(self, base_url=None, model_name=None): + # Recommended model: Qwen3-235B-A22B-Thinking-2507 + self.base_url = base_url + self.model_name = model_name + + def qwen_api_call( + self, + system_prompt: str, + user_input: str, + temperature: float, + max_tokens: int, + ): + """Use Qwen Chat API to perform text rewriting, parse ... sections for + reasoning content, and return (thinking, result).""" + client = openai.OpenAI( + base_url=self.base_url, api_key="None", timeout=600 + ) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_input}, + ] + + last_err = None + + for i in range(10): + try: + stream = False + response = client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + stream=stream, + ) + content = response.choices[0].message.content + thinking = "" + result = content or "" + if content and "" in content: + head, tail = content.split("", 1) + thinking = head.replace("", "").strip() + result = tail.strip() + return thinking, result + except Exception as e: + last_err = e + if i < 9: + time.sleep(2**i) + continue + raise last_err + + def run_single_recaption( + self, system_prompt, input_prompt, temperature=0.1, max_tokens=4096 + ): + thinking, result = self.qwen_api_call( + system_prompt, input_prompt, temperature, max_tokens + ) + return result + + +class QwenVLClient(object): + def __init__(self, base_url=None, model_name=None): + self.base_url = base_url + self.model_name = model_name + self.api_key = "None" + self.max_image_size = int( + os.getenv("I2V_REWRITE_MAX_IMAGE_SIZE", "1024") + ) # 控制送入模型的最大图像分辨率 + + def _encode_image_to_base64( + self, image_path: str, max_dimension: int + ) -> str: + """参考 hyvideo/utils/rewrite/qwen_vllm.py 的实现： 加载本地图片，将其按比例缩放到 max_dimension，然后编码为 Base64 + data URL。""" + try: + from PIL import Image + except ImportError as e: + logger.error( + "Pillow (PIL) is required for QwenVLClient image encoding but is not installed." + ) + raise e + + try: + with Image.open(image_path) as img: + if img.width > max_dimension or img.height > max_dimension: + img.thumbnail((max_dimension, max_dimension)) + + buffer = io.BytesIO() + + # 去除透明通道，统一转为 JPEG + if img.mode in ("RGBA", "LA") or ( + img.mode == "P" and "transparency" in img.info + ): + img = img.convert("RGB") + + img.save(buffer, format="JPEG") + mime_type = "image/jpeg" + + encoded_string = base64.b64encode(buffer.getvalue()).decode( + "utf-8" + ) + return f"data:{mime_type};base64,{encoded_string}" + except Exception as e: + logger.error(f"Failed to encode image {image_path} to base64: {e}") + raise + + def qwen_api_call( + self, + system_prompt: str, + user_input: str, + temperature: float, + max_tokens: int, + img_path: str = None, + ): + """Use Qwen3-VL to perform text rewriting.""" + client = openai.OpenAI( + base_url=self.base_url, + api_key=self.api_key, + timeout=3600, + ) + + assert "{}" in system_prompt, "system_prompt must contain {{}}" + prompt_text = system_prompt.format(user_input) + + assert img_path is not None, "img_path is required" + base64_image = self._encode_image_to_base64( + img_path, max_dimension=self.max_image_size + ) + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": base64_image, + }, + }, + { + "type": "text", + "text": prompt_text, + }, + ], + } + ] + + last_err = None + + for i in range(10): + try: + response = client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + stream=False, + ) + content = response.choices[0].message.content + thinking = "" + result = content or "" + # 和 QwenClient 一样，兼容 思维模式输出 + if content and "" in content: + head, tail = content.split("", 1) + thinking = head.replace("", "").strip() + result = tail.strip() + return thinking, result + except Exception as e: + last_err = e + logger.error( + f"QwenVLClient request failed (attempt {i + 1}/10): {e}" + ) + if i < 9: + time.sleep(2) + else: + raise last_err + + def run_single_recaption( + self, + system_prompt, + input_prompt, + temperature=0.1, + max_tokens=4096, + img_path: str = None, + ): + thinking, result = self.qwen_api_call( + system_prompt, + input_prompt, + temperature, + max_tokens, + img_path=img_path, + ) + return result diff --git a/worldcompass/fastvideo/models/hyvideo/utils/rewrite/i2v_prompt.py b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/i2v_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..1b63e29dd55280276595fdbc75e6b3d94d62e6df --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/i2v_prompt.py @@ -0,0 +1,118 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + + +i2v_rewrite_system_prompt = """ +## 角色 +你是一位顶级的图生视频（Image-to-Video）Prompt工程师。你的任务不是生成视频，而是将用户输入的自由形式的自然语言，改写成具有丰富视觉细节、精确动态描述、并采用专业影视语言的中文Prompt。改写后的措辞、句式、表达方式等，必须严格遵循并尽可能接近本指令中定义的语言风格和表达习惯。 + +## 任务 +你的核心任务是进行"文本改写"。接收用户的简短或模糊想法，输出一段符合以下规则的、详细、客观、可执行的中文视频脚本式描述。无论用户输入是中文还是英文，你的输出都必须是中文。 + +## 核心改写规则 + +### 1. 镜头语言标准化 (Camera Language Standardization) +当用户指令中包含镜头运镜描述时，尽可能转换为标准表述。如果不能完全对应到标准表述，保留原意。 +* **运镜标准表述**: `镜头缓缓拉远/后拉`, `镜头向前推进`, `镜头上/下/左/右移动`, `镜头摇动/摇移`, `镜头跟随`, `镜头环绕`, `镜头静止不动`, `手持镜头`。 +* **示例1**: （能对应标准）用户输入"镜头慢慢前推，跟踪小鸟飞行"，应改写为"**镜头向前推进**，跟踪小鸟飞行"。 +* **示例2**: （不能完全对应）用户输入"镜头缓缓顺时针旋转并向前推"，应改写为"镜头缓缓顺时针旋转，并**向前推进**"。 +* 改写时绝对禁止补充用户指令中没有明确提出的镜头运镜描述，除非为了解释或标准化用户已明确提到的运镜所必需；不得擅自联想新增任何镜头运动方式或运镜效果。特别地，在用户未明确说明“镜头静止”或具有同等含义的表述时，严禁擅自添加“镜头静止不动”等描述。 + +### 2. 动态化与时序性 (Dynamic & Sequential) +将用户的静态描述分解成一个微小的时间序列。使用连接词来串联连续发生的、或同时发生的动作，构建出清晰的叙事流。 +* **结构**: 动作A发生，**随后/然后**，动作B发生，**同时**，动作C发生。 +* **常用连接词**: `随后`, `然后`, `接着`, `同时`, `之后`。 +* **示例1**: 用户输入"两个人见面"，应改写为"...画面左侧走进来一个男子，画面右侧走进来一个女子，他们微笑着在中间的心形前停下，**随后**两个人双手握在一起对视，**接着**男子和女子接吻..."。 +* **示例2**: 用户输入"女孩跳舞"，应改写为"**女孩身体开始左右轻轻摇摆,同时双手缓缓举过头顶**。。。"。 + +### 3. 遵循“主体-动作-细节”的客观描述模式 +使用客观、中立的语言，像摄影师一样记录画面中发生的一切。避免使用主观或情感化的词语（如“美丽的”、“悲伤的”），而是通过描述具体的行为来暗示情感（如用“嘴角露出微笑”代替“开心的”）。 +* **句式结构**: `[主体] + [方式状语(如缓缓地)] + [动作(如转动头部)]`。 +* **示例**: "**黄色头的蜥蜴** **转动着头** **向前探出身子**。" 这是一个完美的客观描述链。 + +### 4. 空间与方位的精确化 +明确物体和人物在画面中的位置及其移动方向,如果有新加入物体，需要描述已有物体加以区分。 +* **方位词**: `画面左/右/上/下侧`, `从...上方/下方伸入`, `背景中`, `前景处`, `向...方向移动`。 +* **示例1**: 用户输入"有只手伸进来"，应改写为"**一只手从画面右侧伸出**，摸了摸黑色衣服上的吊牌...**从下方消失在画面中**"。 +* **示例2**: 参考图像中间有一条美人鱼，用户输入"旁边游来两条粉色鱼尾的美人鱼"，应改写为"**画面中间有一条美人鱼，两条粉色鱼尾的美人鱼从画面右侧游入**..."。 + + +### 5. 指代关系清晰 + *模糊指代必须显式化（针对用户输入）**：当用户输入中出现模糊指代（如“他们”“它”“这个/那个”“其”等），在改写中必须替换为明确的实体称谓，并最小合理补充主体类别、数量与性别，但不引入与原意无关的新实体。 + * **示例**： + * 文本指令“他们跳舞。” → 你的输出中写为“一对男女在舞池中跳着华尔兹...” + * 文本指令“把它拿起来。” → 你的输出中写为“一只手从画面左侧伸入，拿起绿色的洗面奶管子...” + * 文本指令“把炸弹递给他。” → 你的输出中写为“黑猫把手中的炸弹递给灰猫。” + +### 6. 过度克制联想 +在不改变用户意图与叙事目标的前提下，仅进行必要的镜头化与轻度细节扩展。 +* 不得新增用户未要求的关键事件或动作（如坐下、拿起、递交、拥抱等）。 +* 可适度补充环境"微动态"，但不得擅自改变主体、方向、数量与时序。 +* **避免补充光影相关描述**：，除非用户指令明确要求,绝对不要添加任何光影相关内容，包括但不限于"光线"、"阴影"、"光影斑驳"、"阳光照射"、"灯光变化"等。 + 同时，绝对严禁输出"无任何光影变化"、"无额外光影渲染"、"整体画面无额外光影变化"等强调不存在光影的表述，否则将视为严重错误！！！ +* 输出长度不必与样例等长，应与用户输入的信息量匹配；避免冗长与过度联想。 + +### 7. 对齐用户意图 +* 结合参考图像和用户指令，正确理解用户的指定的范畴 + - 例1： 参考图像是：三个漂浮在水面上的甜甜圈 + 用户指令是：“美味的甜甜圈从水上慢慢沉下去消失不见，水波浮动” + 不应该改写成："画面中央的巧克力甜甜圈缓缓下沉...旁边的粉色和白色甜甜圈仍在水上轻轻浮动..." + 应该改写成："画面中的三个甜甜圈缓缓下沉..."（因为这里用户指代的是参考图像当中所有的甜甜圈） + - 例2: 参考图像是：狗人、兔子人、猫人三个人形动物头的角色在故宫的背景，面向镜头奔跑 + 用户指令是：兔子人加速向前跑，然后剩下的两个追上去，镜头右摇，拍摄它们跑远宫殿的背影 + 不应该改写成：画面中，一只狗、一只白兔和一只猫正并排在石桥上向前奔跑。突然，中间的白兔加速向前冲刺。随后，它左侧的狗和右侧的猫也加快速度追了上去。镜头向右摇移，跟随拍摄它们跑向远处宫殿的背影。 + 应该改写成：画面中，一只狗、一只白兔和一只猫正并排在石桥上向前奔跑。突然，中间的白兔加速向前冲刺。随后，它左侧的狗和右侧的猫也加快速度追了上去。镜头向右摇移，跟随拍摄它们奔跑远离宫殿的背影。 + + +## 文字处理规则 +如果用户输入包含引用文字（如书名、品牌名、标语等），必须在改写后的Prompt中保持原文和原语言，并一定要用中文双引号包围。不要在没有原文引用的情况下自行添加引用文字。例如，如果用户提到 写着"GROOVY MANGO"，你必须改写为 写着“GROOVY MANGO”。 +注意：绝对不要在指令中让完整的一段文字，逐个字词的出现，例如： + 用户指令：男人在纸上写"iPhone 15" + 不要改写成：...依次写上“i”、“P”、“h”.../依次形成字母“i”、“P”、“h”.../按顺序写上文字“i”、“P”、“h”... （等等任何把文字拆开的描述方式） + 应该改写成：...写上文字“iPhone 15” + +## 改写范例 +以下是一些改写样例，请严格模仿“改写后”的风格。 + +**范例 1** +* **参考图像**：傍晚的海边沙滩，一位长发女子站在近景位置，身体微侧，远处海浪与天际线清晰可见，海风将发丝轻轻吹起。 +* **文本指令**: 在海边的沙滩上站着的一位女人，面朝大海凝望。 +* **你的输出**: 女人的发丝随风轻轻飘动。女人眨了一下眼，向画面右侧转头凝视前方。背景中海浪拍打着沙滩。 + +**范例 2** +* **参考图像**：户外阴天光线下，一朵花的近景特写，花瓣边缘附着细小水珠，背景为柔和的绿色叶片虚化。 +* **文本指令**: 一朵在小雨中被细密雨丝打湿的花朵，花瓣上挂着水珠。 +* **你的输出**: 雨水滴落在花朵上，水珠在花瓣间滚动，花瓣轻微摇动。 + +**范例 3** +* **参考图像**：室内柔和日光下，一个圆形绒面猫窝里蜷着一只橘猫，窗边可见绿色树叶的模糊影子落在地面上。 +* **文本指令**: The cat stretches its body and wags its tail. Camera pulls back. The cat opens and blinks its eyes. +* **你的输出**: 橘猫在柔软的猫窝中舒展身体，把头和一只爪子伸出猫窝外，不停摇晃着尾巴。镜头微微向后拉远，露出窗外更多的树叶景色。猫咪睁开眼睛看了一下前方，随后又眯上了眼睛，尾巴也渐渐停止摇晃。阳光照下的影子也随之动着。 + +**范例 4** +* **参考图像**：灯光温暖的舞池中央，一对男女面对面站立，男士穿深色西装，女士穿深色礼服，周围有零散观众与舞台乐队背景虚化。 +* **文本指令**: A man and a woman are dancing together as the camera follows their movements. +* **你的输出**: 一对男女在舞池中跳着华尔兹，男士穿着黑色西装轻握穿着深色礼服女士的手，扶着她的腰，女子一只手搭在他的肩上。两人眼神交流，面带微笑。女士裙摆随舞步轻轻摆动。镜头跟随拍摄，捕捉他们的动作和表情。后面的人们拍着手，舞台上的乐队在演奏。 + + +## 输出格式要求 +请注意： +1. 无论用户输入是中文还是英文，你的最终输出都必须是中文。 +2. 在你的最终输出中，不要包含任何星号（**）符号。上述示例中的星号仅用于强调重要概念，你的实际输出应该是干净的、不带任何格式标记的纯文本。 + +现在给定参考图像和文本指令: {} +请按照上述规则进行改写，输出改写后的文本: +""" diff --git a/worldcompass/fastvideo/models/hyvideo/utils/rewrite/rewrite_utils.py b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/rewrite_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1667042ad0a19b6c0ff7dc687781360b0462e3a3 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/rewrite_utils.py @@ -0,0 +1,76 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from hyvideo.utils.rewrite.clients import QwenClient, QwenVLClient +from hyvideo.utils.rewrite.t2v_prompt import t2v_rewrite_system_prompt +from hyvideo.utils.rewrite.i2v_prompt import i2v_rewrite_system_prompt + + +def t2v_rewrite(user_prompt, rewrite_client=None): + base_url = os.getenv("T2V_REWRITE_BASE_URL") + model_name = os.getenv("T2V_REWRITE_MODEL_NAME") + if not base_url or not model_name: + raise EnvironmentError( + "T2V_REWRITE_BASE_URL and T2V_REWRITE_MODEL_NAME must be set in the environment variables " + "when prompt rewriting is enabled. Please configure the rewrite service correctly." + ) + if rewrite_client is None: + rewrite_client = QwenClient(base_url, model_name) + try: + rewritten_prompt = rewrite_client.run_single_recaption( + t2v_rewrite_system_prompt, user_prompt + ) + except Exception as e: + raise ValueError( + f"Failed to rewrite prompt using {type(rewrite_client).__name__}: {e}" + ) + return rewritten_prompt + + +def i2v_rewrite(user_input, img_path, rewrite_client=None): + """Use a rewrite client to generate a rewritten prompt for image-to-video.""" + i2v_base_url = os.getenv("I2V_REWRITE_BASE_URL") + i2v_model_name = os.getenv("I2V_REWRITE_MODEL_NAME") + if not i2v_base_url or not i2v_model_name: + raise EnvironmentError( + "I2V_REWRITE_BASE_URL and I2V_REWRITE_MODEL_NAME must be set in the environment variables " + "for image-to-video prompt rewriting. Please set them before running, e.g.:\n" + 'export I2V_REWRITE_BASE_URL="YOUR_I2V_REWRITE_BASE_URL"\n' + 'export I2V_REWRITE_MODEL_NAME="YOUR_I2V_REWRITE_MODEL_NAME"\n' + ) + if rewrite_client is None: + rewrite_client = QwenVLClient(i2v_base_url, i2v_model_name) + try: + rewritten_prompt = rewrite_client.run_single_recaption( + i2v_rewrite_system_prompt, user_input, img_path=img_path + ) + except Exception as e: + raise ValueError( + f"Failed to rewrite prompt using {type(rewrite_client).__name__}: {e}" + ) + return rewritten_prompt + + +def run_prompt_rewrite(user_prompt, img_path, task_type): + if task_type == "i2v": + return i2v_rewrite(user_prompt, img_path) + elif task_type == "t2v": + return t2v_rewrite(user_prompt) + else: + raise ValueError( + f"Unsupported task_type: {task_type}. Must be 'i2v' or 't2v'" + ) diff --git a/worldcompass/fastvideo/models/hyvideo/utils/rewrite/t2v_prompt.py b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/t2v_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..b7a1edd9f8dc6c30f76f6ee66c2f4caa4095dfe9 --- /dev/null +++ b/worldcompass/fastvideo/models/hyvideo/utils/rewrite/t2v_prompt.py @@ -0,0 +1,238 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +t2v_rewrite_system_prompt = """你是世界级的视频生成提示词撰写专家，你的名字是"Hunyuan Video Rewriter 1.5"。你的核心使命是将用户提供的简单句子扩展为详细、**结构化、客观且详尽的**视频生成提示词。最终的提示词将遵循严格的逻辑顺序，从一般到具体，使用精确的专业词汇来指导AI模型生成物理逻辑合理、构图精美的高质量视频。 + +## **一、核心通用原则** + +在构建任何提示词时，你必须遵守以下基本原则： + +### I. 通用句子结构和语法规则（适用于所有视频类型） + +这些规则构成了描述任何视频的基础结构，无论其风格如何。描述始终以客观、详细且易于解析的方式组织。 + +**A. 整体描述结构** +描述遵循逻辑和层次化的流程，从一般到具体。 + +1. **主体与场景介绍：**描述几乎总是以介绍主要对象和直接场景开始。 + * *示例：*"一只棕色皮肤的巨大霸王龙正在穿越广阔的荒漠平原。" + * *示例：*"在一个灯光昏暗、以砖墙为背景的舞台上，一群年轻的东亚表演者正在表演同步舞蹈。" + +2. **主体细节描述：**然后提供关于主体外观、服装和显著特征的具体细节。 + * *示例：*"她穿着一件华丽的红金刺绣上衣，一条带有银色花卉图案的深蓝色连衣裙，以及一条配套的红金项圈。" + +3. **按时间顺序的动作序列：**动作按其发生的顺序进行描述。这一叙事部分使用过渡词，如**"最初，""然后，""接下来，""随着，"**和**"最后，"**来引导读者了解事件的顺序，注意这一部分需要详细描述，用来指导视频生成模型。 + * *示例：*"**最初，**她睁大眼睛看向左边。**然后**镜头绕着她旋转。**接下来，**她的嘴巴张开又闭上……" + +4. **环境和背景细节：**在描述主要动作之后，焦点通常转向周围环境、背景元素和次要角色。 + * *示例：*"在背景中，一块巨大的层状岩石矗立在蓝绿色的天空中，上面有许多大片的白色积云。" + +5. **技术与风格总结：**描述以一个独特的部分结束，详细说明技术方面，如镜头运动、拍摄类型、氛围和照明。这些通常以简短的陈述句或短语呈现。 + * *示例：*"镜头向右平移……低角度。冒险感。" + +**B. 核心语法规则** + +1. **时态：**主要使用的时态是**现在时**（一般现在时或现在进行时）。这使描述感觉即时和主动，就像在播放视频时进行描述一样。 + * *示例：*"一个小女孩……**正在跑**……"，"他**穿着**黑色运动背心……"，"镜头**跟随**她……" + +2. **句子结构：** + * 句子主要是**陈述性的**，陈述关于场景的事实。 + * 结构通常以主语开始，然后是动词：`[主语] + [动词] + [细节]`。 + * **介词短语**被广泛用于添加关于位置（`在土路上`）、外观（`有着苍白的皮肤`）和关系（`在瓢虫后面`）的细节。 + * **分词短语**经常用于简洁地组合动作或描述。 + * *示例：*"一位年轻女性……**戴着棕色帽子**，从马上下来，**微笑着等待**。" + +3. **词汇和语气：** + * **形容词：**语言中充满描述性形容词，用于指定颜色、大小、纹理、情感和外观（例如，"广阔的"，"卷曲的"，"同步的"，"充满活力的"）。 + * **动词：**动作动词精确而动态（例如，"跳跃"，"投掷"，"摇摆"，"攀爬"）。 + * **语气：**语气是客观和事实性的。它描述视觉上呈现的内容，而不做过度主观的解释（除非陈述情绪，例如，"冒险的"，"欢乐的重逢"）。 + +--- + +### II. 不同视频类型的规则和特征 + +虽然上述通用规则适用于所有视频，但某些视频类型在其描述中具有独特的特征。 + +**A. 写实/真人视频** + +1. **关注人物细节：**描述优先考虑人物特征：年龄、种族、肤色、发型/颜色和具体服装项目。情绪状态通过面部表情和肢体语言来描述。 + * *示例：*"这位女性，她的脸上带着**担忧的表情**，抬起头与他对视……她的表情突然转变为**震惊和恐慌**……" + +2. **真实世界的动作和互动：**所描述的动作以现实为基础，关注人与人之间或与物理环境的互动。 + * *示例：*"当跑步者到达她身边时，他们立即**紧紧地拥抱在一起**。" + * *示例：*"他用右手撑起自己，**努力站起来**……" + +3. **电影术语：**描述通常包括暗示物理摄像机的特定电影制作术语。 + * *示例：*"**手持拍摄**向前移动，""**中近景**，""**平视角度**，""镜头**以中景稍微移动拍摄**。" + +4. **照明描述：**照明通常用自然光源（"来自阴天的柔和、漫射自然光"）或刻意的电影布光设置（"高调、漫射的背光"，"高对比度、戏剧性的舞台照明"）来描述。 + +**B. 动画/CGI视频** + +1. **强调风格化：**描述突出动画风格和非写实特征。主体通常是奇幻的（恐龙、蓝精灵）或拟人化的（拿着炸药的北极熊）。 + * *示例：*"四只**动画蓝精灵**骑在一只**毛茸茸的大黄兔子**上。" + * *示例：*"一只有橙色翅膀的瓢虫正在高速飞行……在瓢虫后面，有五只黑色苍蝇，它们有**显眼的大红眼睛**。" + +2. **违反物理和夸张的动作：**所描述的动作通常超越现实世界的限制，反映了动画的创作自由。 + * *示例：*"熊用后腿站立，转身，**将点燃的炸药高高地投入拱形门道**。" + * *示例：*"她突然**被甩出车外**……**安全落地**。" + +3. **明确的风格识别：**技术总结通常明确命名动画风格。 + * *示例：*"一个动态的**风格化昆虫的3D CGI动画**集成在照片级真实的自然环境中……" + * *示例：*"视觉风格是**高质量的3D电脑动画电影**……" + +### **II. 镜头控制指南** +以下是镜头控制系统。你应该使用镜头控制系统来描述提示词中的镜头运动。如果遇到下面的类型，就参考对应的描述。 +* **镜头360度旋转**: ["镜头旋转360度", "镜头进行完整旋转", "镜头绕一圈旋转"] +* **镜头第一人称视角FPV**: ["镜头显示第一人称视角", "场景从第一人称视角拍摄", "镜头采用FPV角度", "镜头处于第一人称视角"] +* **镜头向上移动**: ["镜头向上移动", "镜头上升", "镜头升起"] +* **镜头向下移动**: ["镜头向下移动", "镜头下降", "镜头落下"] +* **镜头低角度/仰拍**: ["镜头从低角度拍摄", "镜头从下方捕捉场景", "镜头位于低视点"] +* **镜头向上倾斜**: ["镜头向上倾斜", "镜头进行向上倾斜运动"] +* **镜头向下倾斜**: ["镜头向下倾斜", "镜头进行向下倾斜运动"] +* **镜头地面拍摄**: ["镜头在地面高度", "镜头从地面拍摄", "镜头从地面视角捕捉场景"] +* **镜头向前推进**: ["镜头向前推进", "镜头向前移动", "镜头向前移"] +* **镜头向右平移**: ["镜头向右移动", "镜头向右移"] +* **镜头向后拉**: ["镜头向后拉", "镜头后退", "镜头向后移"] +* **镜头向左平移**: ["镜头向左移动", "镜头向左移"] +* **镜头延时摄影**: ["镜头捕捉延时摄影", "使用延时拍摄", "场景以延时方式显示"] +* **镜头微距拍摄**: ["镜头进行微距拍摄", "使用微距视角"] +* **镜头慢动作**: ["镜头以慢动作记录", "显示慢动作镜头", "场景以慢动作捕捉"] +* **镜头拉远**: ["镜头拉远", "镜头向后拉", "镜头远离主体"] +* **镜头推近**: ["镜头推近", "镜头向前推", "镜头接近主体"] +* **镜头向右平移**: ["镜头向右平移", "镜头向右摆动", "镜头进行向右平移运动"] +* **镜头向左平移**: ["镜头向左平移", "镜头向左摆动", "镜头进行向左平移运动"] +* **镜头无人机视角**: ["镜头显示无人机视角", "场景从无人机视角拍摄", "镜头采用空中无人机角度"] +* **镜头环绕**: ["镜头环绕主体", "镜头围绕主体旋转"] +* **镜头跟随拍摄**: ["镜头跟随主体", "镜头跟踪运动", "镜头与主体一起移动"] +* **镜头过肩拍摄**: ["镜头使用过肩镜头", "镜头位于主体肩膀后方"] +* **镜头逆时针旋转**: ["镜头逆时针旋转", "镜头以逆时针方向旋转"] +* **镜头静止**: ["镜头保持静止", "镜头静止不动", "镜头保持静态"] +* **镜头顺时针旋转**: ["镜头顺时针旋转", "镜头以顺时针方向旋转"] +* **镜头高角度/俯拍**: ["镜头从高角度拍摄", "镜头从上方捕捉场景", "镜头位于高视点"] +* **镜头鱼眼镜头**: ["镜头使用鱼眼镜头", "场景以鱼眼效果显示", "应用鱼眼透视"] +* **镜头鸟瞰视角**: ["镜头显示鸟瞰视角", "场景从正上方拍摄"] + +** 相机运动也可以是动态的，如"最初，摄像机跟随...然后，焦点平滑地转移，向后拉..." + +## **三、标准生成流程** + +在生成最终提示词之前，你必须在思考和构建时遵循以下步骤： + +0. **语言规则**: + * 整体输出的提示词保持为英文。 + * 文本渲染内容应与用户输入的语言相同。例如，如果用户希望视频显示文本"Hello"，则扩展的提示词应该是英文，但渲染的文本应该是"Hello"。如果是中文，文本渲染内容总是包含在“”内，如果是英文则总是放置在""内。 + * **宝可梦IP的特殊规则**：如果用户输入包含宝可梦的IP角色，始终使用英文的IP名称。（例如，使用Jigglypuff而不是胖丁，因为胖丁来自宝可梦） + +1. **分析核心元素并评估风险**: + * **摘要**：提示词以视频故事的摘要开始。所有主要主体都必须在摘要中描述。摘要应该是一个简洁的句子，并放在提示词的开头。 + * **识别核心元素**：从用户输入中清晰地识别主体（人物、对象）、关键动作/事件、运动、环境和整体叙事弧线。 + * **确定实体数量**：如果用户输入包含多个实体，首先检查用户是否给出了精确的实体数量（例如，"六口之家"，"四个朋友"，"大约五名舞者"），那么你必须严格遵循用户的原始提示词并使用相同的数量描述。否则，如果用户使用模糊的词语（例如，"一群警官"）并且没有精确给出实体数量，那么你应该严格将人物或对象的数量限制在三个或更少，以保持场景清晰度。 + * **识别高风险概念**：特别注意复杂的物理互动（例如，动态体育序列）、随时间展开的抽象概念（例如，"能量"脉动或流动）以及动态排版，需要用简单、可渲染的描述来呈现。 + * **可视化概念**：将非视觉概念转换为视觉可用的序列。例如，"发令枪的声音"暗示比赛的开始，运动员从起跑线冲出并沿着跑道移动。你可以从用户提供的这种非视觉描述中推断事件序列。或出现鸟鸣声，那么场景中就需要出现鸟。 + +2. **确定摄影和构图**: + * **摄影和镜头运动**：使用以下规则决定摄影： + * 如果用户输入包含摄影描述（例如，广角镜头、航拍镜头、跟踪镜头、变焦镜头），你必须严格遵循用户输入，不要更改它。 + * 如果用户没有提供描述，基于场景和核心事件，你应该使用你的摄影知识来选择合适的镜头工作。优先选择能完全捕捉关键动作和所有核心元素的镜头和运动（例如，**缓慢平移穿过场景的略高角度镜头**）。 + * **构图和场景调度**：应用构图技术，如三分法或对称，以确保每一帧都经过专业构图。至关重要的是，考虑主体和元素随时间在画面中的移动方式（调度）——它们的进入、退出和运动路径。 + +3. **选择艺术风格**: + * **风格选择**：如果用户指定了风格（例如，油画动画、动漫、动态图形），严格遵守它。否则，首先根据用户输入推断合适的风格。如果没有明确的风格偏好，默认为**电影感写实风格**。 + +4. **确定镜头运动**: + * **镜头运动**：使用以下规则决定镜头运动： + * 如果用户输入包含镜头运动描述（例如，向前推进、推近、向上倾斜），你必须根据输入描述使用镜头运动关键字。（例如，如果用户输入是"镜头向上移动"，你必须推断它对应关键字"camera upward"，并在"camera upward"的值列表中选择一个词，即["The camera moves upward", "The camera rises", "The camera ascends"]） + * *示例：*输入："镜头向上移动"，使用词语"The camera rises"， + * *示例：*输入："高角度镜头"，使用词语"The camera captures the scene from above"。 + * 如果用户没有提供镜头运动描述。镜头拍摄必须是"The shot is at an eye-level angle with the main subject."。 + 并且镜头运动应该从镜头静止、镜头向前推进、镜头向后拉、镜头向右平移、镜头向左平移、镜头推近、镜头拉远中选择。你需要分析主体的运动和背景的运动，然后选择最合适的镜头运动。首选轻微的运动，以防止画面变化过快产生畸变。 + +4. **填充细节和审查**: + * **时长**：视频时长固定为5秒。填充的细节必须能在5秒内完成。 + * **细节填充**：遵循结构化顺序，描述主体的材料、纹理、动作、手势以及随时间变化的表情，以及环境中的任何次要元素及其变化方式。注意添加恰到好处的细节量；除非汗水是视频叙事的关键，否则不要过分强调汗水等元素。火花、雷电也是容易生成错误的内容，需要注意。**镜头或手机屏幕的正反面**：当场景中出现相机镜头、手机屏幕等物体时，必须明确描述其朝向（正面/背面/侧面），以及观众看到的是屏幕内容还是设备背面，避免产生歧义。例如："手机屏幕朝向镜头，显示......"或"相机的镜头正对着......"。 + * **运动填充**：视频中主体的运动必须在场景中逻辑合理且一致。并且运动必须清楚说明。主体的运动必须能在5秒内完成。一个坏例子是描述大量无法在5秒内完成的动作。主体的动作需要详细描述。 + * **逻辑审查**：检查描述是否完全符合物理定律、因果关系和场景逻辑贯穿整个视频时长。例如，动态场景中的动作流程是否合理？角色移动时的视线和互动是否一致？剪辑开始和结束之间是否有连续性错误？ + * **完整性审查**：确保所有预期的关键元素和事件（例如，裁判发出信号、背景爆炸）都被明确描述。检查镜头运动或主体运动是否导致不自然的裁剪或遮挡，特别是在动作期间的人体肢体或关键对象。 + * **生成畸形审查**：必须高度警惕任何可能引发动作或身体畸形的描述，严密审查并排除一切可能导致角色动作不自然、关节弯曲异常或肢体结构错误的因素，确保所有人体结构、动作轨迹、关节形态均自然、连贯，完全符合物理规律。 + +5. **最终验证**: + * **用户输入遵循检查**：将最终结果与用户输入进行比较，以确保完全描述了用户的核心内容，例如核心实体及其属性、对象或人物的数量、指定的摄影、时长、节奏和事件顺序。最终提示词不得添加用户输入未暗示的任何新前景对象或主要事件。 + * **检查物理和时间逻辑**：检查提示词中是否存在任何物理或时间逻辑错误。例如，对象大小在移动期间是否一致？运动物理是否可信（例如，加速度、重力）？反应的时间是否合理？事件之间的因果关系是否清晰？ + * **检查宝可梦IP**：检查提示词是否包含任何宝可梦IP；你应该使用宝可梦角色的英文名称。（例如，使用Jigglypuff而不是胖丁） + +6. **如果验证失败则重试**: + * **从头开始**：如果验证失败，你应该从头开始重新生成扩展的提示词。 + + + +### **四、风格特定创作指南** + +根据确定的艺术风格激活相应的专业知识库。 + +**1. 摄影和写实风格** +* **总体规则**：想象你是一位摄影大师，用户的输入将被你转换为专业摄影师拍摄的照片。假设你正在查看用户描述的图片，你将使用你的专业摄影知识将用户输入转换为具有专业构图和专业照明的视觉图片。 +* **摄影风格视频的专业照明**：你应该使用你的专业摄影照明技术来增强真实感。你应该根据用户的输入选择合适的照明技术，以下是一些可供选择的示例技术，你可以使用你的世界知识来获得更好的选择。 + * 使用戏剧性照明，强调光影之间的高对比度以创造深度感。使用伦勃朗照明从45度角照亮主体的面部，在脸颊上形成三角形高光，并在面部的另一侧投下深深的阴影。这应该以强烈的维度感突出面部特征。 + * 对于整体场景，应用黄金时段照明，具有温暖的色温，光线在过渡到阴影时逐渐柔化。柔和的侧光应该照亮主体的身体，在光影之间创造柔和的渐变，增强三维形式。背景应该具有渐进的阴影过渡，唤起平静、诱人的氛围。 + * 为了增加真实感，使用背光在主体周围创造剪影效果，强光位于主体后面以突出轮廓。这应该将背景投入深深的阴影中，边缘由附近表面的反射光柔化。确保主体的轮廓清晰，但周围的阴影增加了一种神秘感。 + * 从左侧加入冷光，与温暖的背景光形成对比，创造视觉张力。光强度应该变化，在关键特征上使用强烈的高光，在其他区域使用柔和的阴影。光影之间的过渡应该感觉自然和平衡，焦点区域由锐利的硬光照亮，阴影中的区域更柔和。 + * 为了增加额外的纹理和真实感，在水或玻璃等反射表面上包括反射，光线应该反弹以柔化阴影并在表面上创造光斑。直射光和反射光之间的这种相互作用将为整体构图增加复杂性和趣味性。 + * 确保光线在引导观众的眼睛穿过视频中起着关键作用，无论是通过阴影创造的引导线，还是通过突出场景关键元素的光线定向流动。 + * 调整以获得更多控制的关键参数： + * 照明风格：（例如，伦勃朗、柔和、硬、背光、侧光） + * 光线方向：（例如，45度角、自上而下、侧光） + * 光线质量：（例如，柔和、刺眼、漫射、聚光灯） + * 阴影细节：（例如，深阴影、柔和渐变、高对比度） + * 色温：（例如，温暖的黄金时段、凉爽的日光） + * 反射：（例如，水、玻璃或金属表面上的反射光） + * 剪影和轮廓：（例如，主体背光，创造戏剧性的轮廓） + * 可定制元素： + * 照明情绪：为场景定义整体照明情绪（例如，戏剧性、柔和、高对比度、微妙渐变）。 + * 背景照明：调整光线与背景的相互作用，例如柔和渐变或强烈阴影区域。 + * 柔和对比硬阴影：指定阴影应该有多刺眼或多漫射。 + * 高光细节：关注光线应该突出关键特征的区域（例如，面部、眼睛、纹理）。 + * 氛围：（例如，忧郁、宁静、戏剧性、和平） + * **重要规则**：如果光源不可见，只需描述光效，不要提及视频中未出现的任何光源，一个坏例子是"一个次要的、看不见的温暖光源，可能是画面外的台灯"，像这个例子的描述是被禁止的。 +* **镜头效果**：使用专业术语来描述镜头效果（例如，广角透视、长焦压缩、浅景深），你应该使用你的世界知识为用户的输入选择最佳镜头效果。 +* **构图**：使用专业术语为用户的输入选择最佳构图（例如，引导线、框架构图、三分法），**但不要直接使用构图技术名称，你应该安排人物/对象/环境来反映构图本身**。 +* **极致细节**：深入描述材料纹理（例如，木纹反射、织物纤维）、角色细节（例如，眼睛中的眼神光、皮肤毛孔）和环境氛围（例如，空气中的灰尘颗粒）。 +* **多实体场景**：如果用户的输入表明视频中有多个实体，例如一群人、一个团队或场景中的多个人，如果用户给出了精确的实体数量，则严格遵守用户的输入，不要更改数量。但如果用户没有给出具体的人物或对象数量，严格将人物或对象的数量限制在三个或更少。详细描述每个人/对象，并将它们放在中景到前景，确保他们的面部和肢体清晰、未变形，并且在关节处没有被裁剪。**这一条非常重要** +* **电影摄影写实**：如果用户的输入表明视频是写实风格，风格必须是"cinematic realistic style"（电影摄影写实风格）。 + +**2. 插画和绘画风格（卡通、油画、水彩等）** +* **定义类型**：精确定义风格（例如，"日本赛璐珞动画风格"，"厚涂油画"，"湿画法水彩"，"印象派点彩画法"）。 +* **媒介特定特征**：专注于描述风格的独特视觉语言，例如线条的粗细（"G笔线稿"）、笔触的纹理（"可见的、三维的笔触"）和颜料的特性（"水彩边缘的自然水渍"）。 +* **角色设计**：强调夸张的特征（例如，"Q版身体比例"，"占据脸部三分之一的大眼睛"）和富有表现力的姿势。 + +**3. 字体艺术** +* **主体优先**：描述必须以类似`The words "[Text Content]" rendered as...`的短语开始，以将文本确立为绝对核心主体。 +* **安全透视和构图**：强制使用安全的正面或俯视图（`Front view`，`Top-down`），并将文本主体放在画面中心，使用简单或广阔的背景进行对比，从根本上防止裁剪。 +* **形式而非形成**：描述文本构成的最终"形式"，而不是其"形成过程"。（错误："两个GPU相互倾斜形成一个A"；正确："由GPU制成的金字塔形状的字母A"）。 +* **完整性保险**：在提示词末尾添加强制性指令，例如`The entire [object/phrase] is fully visible`，作为防止裁剪的最后防线。 +* **禁止高风险词汇**：避免使用"巨大的"、"特写"、"复杂的"和"精致的"等词，因为它们可能诱导AI放大主体的一部分，导致裁剪。 + +### **五、最终输出要求** + +1. **仅输出最终提示词**：不要显示任何思考过程、Markdown格式或与文本到视频提示词无关的任何表达（如"摘要显示"）。 +2. **输出语言**：扩展的提示词应该是英文，同时根据语言规则保持文本渲染内容与用户输入的语言相同。 +3. **忠实于输入**：你必须保留用户输入句子中的核心概念、属性、数量和文本渲染内容。 +4. **风格强化**：在提示词中提及核心风格3-5次，并以风格声明句结束（例如，"整个视频是电影感写实风格"）。 +5. **字数控制**：描述主要主体的长度应该在140个单词左右，这里面主体的动作是重要部分。描述背景的长度应该是70个单词。描述其他属性（包括构图、光线、风格、氛围、拍摄角度、拍摄类型）的总长度应该在140个单词以内。 +6. **避免自我引用**：在开头直接描述视频内容，删除冗余短语，如"这个视频"或"这个视频显示了"。 + +接下来，我将提供输入句子，你将提供扩展的提示词。 +""" diff --git a/worldcompass/fastvideo/models/loader/__init__.py b/worldcompass/fastvideo/models/loader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/models/loader/component_loader.py b/worldcompass/fastvideo/models/loader/component_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..39f1cfe4a4dd5d665211419497fd2426a390e2e0 --- /dev/null +++ b/worldcompass/fastvideo/models/loader/component_loader.py @@ -0,0 +1,657 @@ +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import glob +import json +import os +import time +from abc import ABC, abstractmethod +from collections.abc import Generator, Iterable +from copy import deepcopy +from typing import cast + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors.torch import load_file as safetensors_load_file +from torch.distributed import init_device_mesh +from transformers import AutoImageProcessor, AutoTokenizer +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from fastvideo.configs.models import EncoderConfig +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.hf_transformer_utils import get_diffusers_config +from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model, shard_model +from fastvideo.models.loader.utils import set_default_torch_dtype +from fastvideo.models.loader.weight_utils import ( + filter_duplicate_safetensors_files, + filter_files_not_needed_for_inference, + pt_weights_iterator, + safetensors_weights_iterator, +) +from fastvideo.models.registry import ModelRegistry +from fastvideo.platforms import current_platform +from fastvideo.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class ComponentLoader(ABC): + """Base class for loading a specific type of model component.""" + + def __init__(self, device=None) -> None: + self.device = device + + @abstractmethod + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the component based on the model path, architecture, and inference args. + + Args: + model_path: Path to the component model + fastvideo_args: FastVideoArgs + + Returns: + The loaded component + """ + raise NotImplementedError + + @classmethod + def for_module_type( + cls, module_type: str, transformers_or_diffusers: str + ) -> "ComponentLoader": + """Factory method to create a component loader for a specific module type. + + Args: + module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler") + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + A component loader for the specified module type + """ + # Map of module types to their loader classes and expected library + module_loaders = { + "scheduler": (SchedulerLoader, "diffusers"), + "transformer": (TransformerLoader, "diffusers"), + "transformer_2": (TransformerLoader, "diffusers"), + "vae": (VAELoader, "diffusers"), + "text_encoder": (TextEncoderLoader, "transformers"), + "text_encoder_2": (TextEncoderLoader, "transformers"), + "tokenizer": (TokenizerLoader, "transformers"), + "tokenizer_2": (TokenizerLoader, "transformers"), + "image_processor": (ImageProcessorLoader, "transformers"), + "image_encoder": (ImageEncoderLoader, "transformers"), + } + + if module_type in module_loaders: + loader_cls, expected_library = module_loaders[module_type] + # Assert that the library matches what's expected for this module type + assert ( + transformers_or_diffusers == expected_library + ), f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + return loader_cls() + + # For unknown module types, use a generic loader + logger.warning( + "No specific loader found for module type: %s. Using generic loader.", + module_type, + ) + return GenericComponentLoader(transformers_or_diffusers) + + +class TextEncoderLoader(ComponentLoader): + """Loader for text encoders.""" + + @dataclasses.dataclass + class Source: + """A source for weights.""" + + model_or_path: str + """The model ID or path.""" + + prefix: str = "" + """A prefix to prepend to all weights.""" + + fall_back_to_pt: bool = True + """Whether .pt weights can be used.""" + + allow_patterns_overrides: list[str] | None = None + """If defined, weights will load exclusively using these patterns.""" + + counter_before_loading_weights: float = 0.0 + counter_after_loading_weights: float = 0.0 + + def _prepare_weights( + self, + model_name_or_path: str, + fall_back_to_pt: bool, + allow_patterns_overrides: list[str] | None, + ) -> tuple[str, list[str], bool]: + """Prepare weights for the model. + + If the model is not local, it will be downloaded. + """ + # model_name_or_path = (self._maybe_download_from_modelscope( + # model_name_or_path, revision) or model_name_or_path) + + is_local = os.path.isdir(model_name_or_path) + assert is_local, "Model path must be a local directory" + + use_safetensors = False + index_file = SAFE_WEIGHTS_INDEX_NAME + allow_patterns = ["*.safetensors", "*.bin"] + + if fall_back_to_pt: + allow_patterns += ["*.pt"] + + if allow_patterns_overrides is not None: + allow_patterns = allow_patterns_overrides + + hf_folder = model_name_or_path + + hf_weights_files: list[str] = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if len(hf_weights_files) > 0: + if pattern == "*.safetensors": + use_safetensors = True + break + + if use_safetensors: + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file + ) + else: + hf_weights_files = filter_files_not_needed_for_inference( + hf_weights_files + ) + + if len(hf_weights_files) == 0: + raise RuntimeError( + f"Cannot find any model weights with `{model_name_or_path}`" + ) + + return hf_folder, hf_weights_files, use_safetensors + + def _get_weights_iterator( + self, source: "Source", to_cpu: bool + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Get an iterator for the model weights based on the load format.""" + hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + source.model_or_path, + source.fall_back_to_pt, + source.allow_patterns_overrides, + ) + if use_safetensors: + weights_iterator = safetensors_weights_iterator( + hf_weights_files, to_cpu=to_cpu + ) + else: + weights_iterator = pt_weights_iterator( + hf_weights_files, to_cpu=to_cpu + ) + + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # Apply the prefix. + return ( + (source.prefix + name, tensor) + for (name, tensor) in weights_iterator + ) + + def _get_all_weights( + self, + model: nn.Module, + model_path: str, + to_cpu: bool, + ) -> Generator[tuple[str, torch.Tensor], None, None]: + primary_weights = TextEncoderLoader.Source( + model_path, + prefix="", + fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), + allow_patterns_overrides=getattr( + model, "allow_patterns_overrides", None + ), + ) + yield from self._get_weights_iterator(primary_weights, to_cpu) + + secondary_weights = cast( + Iterable[TextEncoderLoader.Source], + getattr(model, "secondary_weights", ()), + ) + for source in secondary_weights: + yield from self._get_weights_iterator(source, to_cpu) + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the text encoders based on the model path, and inference args.""" + # model_config: PretrainedConfig = get_hf_config( + # model=model_path, + # trust_remote_code=fastvideo_args.trust_remote_code, + # revision=fastvideo_args.revision, + # model_override_args=None, + # ) + model_config = get_diffusers_config(model=model_path) + model_config.pop("_name_or_path", None) + model_config.pop("transformers_version", None) + model_config.pop("model_type", None) + model_config.pop("tokenizer_class", None) + model_config.pop("torch_dtype", None) + logger.info("HF Model config: %s", model_config) + + # @TODO(Wei): Better way to handle this? + try: + encoder_config = ( + fastvideo_args.pipeline_config.text_encoder_configs[0] + ) + encoder_config.update_model_arch(model_config) + encoder_precision = ( + fastvideo_args.pipeline_config.text_encoder_precisions[0] + ) + except Exception: + encoder_config = ( + fastvideo_args.pipeline_config.text_encoder_configs[1] + ) + encoder_config.update_model_arch(model_config) + encoder_precision = ( + fastvideo_args.pipeline_config.text_encoder_precisions[1] + ) + + target_device = get_local_torch_device() + # TODO(will): add support for other dtypes + return self.load_model( + model_path, + encoder_config, + target_device, + fastvideo_args, + encoder_precision, + ) + + def load_model( + self, + model_path: str, + model_config: EncoderConfig, + target_device: torch.device, + fastvideo_args: FastVideoArgs, + dtype: str = "fp16", + ): + use_cpu_offload = ( + fastvideo_args.text_encoder_cpu_offload + and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0 + ) + + if fastvideo_args.text_encoder_cpu_offload: + target_device = ( + torch.device("mps") + if current_platform.is_mps() + else torch.device("cpu") + ) + + with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): + with target_device: + architectures = getattr(model_config, "architectures", []) + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + model = model_cls(model_config) + + weights_to_load = {name for name, _ in model.named_parameters()} + loaded_weights = model.load_weights( + self._get_all_weights(model, model_path, to_cpu=use_cpu_offload) + ) + self.counter_after_loading_weights = time.perf_counter() + logger.info( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights + - self.counter_before_loading_weights, + ) + + # Explicitly move model to target device after loading weights + model = model.to(target_device) + + if use_cpu_offload: + # Disable FSDP for MPS as it's not compatible + if current_platform.is_mps(): + logger.info( + "Disabling FSDP sharding for MPS platform as it's not compatible" + ) + else: + mesh = init_device_mesh( + "cuda", + mesh_shape=(1, dist.get_world_size()), + mesh_dim_names=("offload", "replicate"), + ) + shard_model( + model, + cpu_offload=True, + reshard_after_forward=True, + mesh=mesh["offload"], + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=fastvideo_args.pin_cpu_memory, + ) + # We only enable strict check for non-quantized models + # that have loaded weights tracking currently. + # if loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + raise ValueError( + "Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}" + ) + + return model.eval() + + +class ImageEncoderLoader(TextEncoderLoader): + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the text encoders based on the model path, and inference args.""" + # model_config: PretrainedConfig = get_hf_config( + # model=model_path, + # trust_remote_code=fastvideo_args.trust_remote_code, + # revision=fastvideo_args.revision, + # model_override_args=None, + # ) + with open(os.path.join(model_path, "config.json")) as f: + model_config = json.load(f) + model_config.pop("_name_or_path", None) + model_config.pop("transformers_version", None) + model_config.pop("torch_dtype", None) + model_config.pop("model_type", None) + logger.info("HF Model config: %s", model_config) + + encoder_config = fastvideo_args.pipeline_config.image_encoder_config + encoder_config.update_model_arch(model_config) + + if fastvideo_args.image_encoder_cpu_offload: + target_device = ( + torch.device("mps") + if current_platform.is_mps() + else torch.device("cpu") + ) + else: + target_device = get_local_torch_device() + # TODO(will): add support for other dtypes + return self.load_model( + model_path, + encoder_config, + target_device, + fastvideo_args, + fastvideo_args.pipeline_config.image_encoder_precision, + ) + + +class ImageProcessorLoader(ComponentLoader): + """Loader for image processor.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the image processor based on the model path, and inference args.""" + logger.info("Loading image processor from %s", model_path) + + image_processor = AutoImageProcessor.from_pretrained( + model_path, + ) + logger.info( + "Loaded image processor: %s", image_processor.__class__.__name__ + ) + return image_processor + + +class TokenizerLoader(ComponentLoader): + """Loader for tokenizers.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the tokenizer based on the model path, and inference args.""" + logger.info("Loading tokenizer from %s", model_path) + + tokenizer = AutoTokenizer.from_pretrained( + model_path, # "/tokenizer" + # in v0, this was same string as encoder_name "ClipTextModel" + # TODO(will): pass these tokenizer kwargs from inference args? Maybe + # other method of config? + padding_size="right", + ) + logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__) + return tokenizer + + +class VAELoader(ComponentLoader): + """Loader for VAE.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the VAE based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name") + assert ( + class_name is not None + ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + fastvideo_args.model_paths["vae"] = model_path + + vae_config = fastvideo_args.pipeline_config.vae_config + vae_config.update_model_arch(config) + + if fastvideo_args.vae_cpu_offload: + target_device = ( + torch.device("mps") + if current_platform.is_mps() + else torch.device("cpu") + ) + else: + target_device = get_local_torch_device() + + with set_default_torch_dtype( + PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + ): + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = vae_cls(vae_config).to(target_device) + + # Find all safetensors files + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors") + ) + # TODO(PY) + assert ( + len(safetensors_list) == 1 + ), f"Found {len(safetensors_list)} safetensors files in {model_path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict( + loaded, strict=False + ) # We might only load encoder or decoder + + return vae.eval() + + +class TransformerLoader(ComponentLoader): + """Loader for transformer.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the transformer based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + hf_config = deepcopy(config) + cls_name = config.pop("_class_name") + if cls_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported." + ) + + fastvideo_args.model_paths["transformer"] = model_path + + # Config from Diffusers supersedes fastvideo's model config + # dit_config = fastvideo_args.pipeline_config.dit_config + # dit_config.update_model_arch(config) + + # 获取模型类，很重要，如果要用新的模型需要在这里进行注册 + # cls_name = "ARWanTransformer3DModel" + # print(fastvideo_args.module_name) + if ( + fastvideo_args.module_name == "transformer" + ): # 对于ar或者ar_action模型来说加载自己的模型 + cls_name = fastvideo_args.cls_name + load_from_dir = fastvideo_args.load_from_dir + ar_action_load_from_dir = fastvideo_args.ar_action_load_from_dir + else: + cls_name = "WanTransformer3DBidirectModel" + load_from_dir = model_path + ar_load_from_dir = None + model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) + # print(model_cls, cls_name) + # quit() + # Find all safetensors files + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors") + ) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + + logger.info( + "Loading model from %s safetensors files in %s", + len(safetensors_list), + model_path, + ) + + default_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.dit_precision + ] + + # Load the model using FSDP loader + logger.info( + "Loading model from %s, default_dtype: %s", cls_name, default_dtype + ) + assert fastvideo_args.hsdp_shard_dim is not None + # print(model_path, model_cls) + # quit() + model = maybe_load_fsdp_model( + load_from_dir=load_from_dir, + ar_action_load_from_dir=ar_action_load_from_dir, + cls_name=cls_name, + model_cls=model_cls, + # init_params={ + # "config": dit_config, + # "hf_config": hf_config + # }, + init_params=None, + weight_dir_list=safetensors_list, + device=get_local_torch_device(), + hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim, + hsdp_shard_dim=fastvideo_args.hsdp_shard_dim, + cpu_offload=fastvideo_args.dit_cpu_offload, + pin_cpu_memory=fastvideo_args.pin_cpu_memory, + fsdp_inference=fastvideo_args.use_fsdp_inference, + # TODO(will): make these configurable + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + training_mode=fastvideo_args.training_mode, + ) + + total_params = sum(p.numel() for p in model.parameters()) + logger.info("Loaded model with %.2fB parameters", total_params / 1e9) + + dtypes = set(param.dtype for param in model.parameters()) + if len(dtypes) > 1: + model = model.to(default_dtype) + model = model.eval() + return model + + +class SchedulerLoader(ComponentLoader): + """Loader for scheduler.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the scheduler based on the model path, and inference args.""" + config = get_diffusers_config(model=model_path) + + class_name = config.pop("_class_name") + assert ( + class_name is not None + ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + + scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + scheduler = scheduler_cls(**config) + if fastvideo_args.pipeline_config.flow_shift is not None: + scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift) + if fastvideo_args.pipeline_config.timesteps_scale is not None: + scheduler.set_timesteps_scale( + fastvideo_args.pipeline_config.timesteps_scale + ) + return scheduler + + +class GenericComponentLoader(ComponentLoader): + """Generic loader for components that don't have a specific loader.""" + + def __init__(self, library="transformers") -> None: + super().__init__() + self.library = library + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load a generic component based on the model path, and inference args.""" + logger.warning( + "Using generic loader for %s with library %s", + model_path, + self.library, + ) + + if self.library == "transformers": + from transformers import AutoModel + + model = AutoModel.from_pretrained( + model_path, + trust_remote_code=fastvideo_args.trust_remote_code, + revision=fastvideo_args.revision, + ) + logger.info( + "Loaded generic transformers model: %s", + model.__class__.__name__, + ) + return model + elif self.library == "diffusers": + logger.warning( + "Generic loading for diffusers components is not fully implemented" + ) + + model_config = get_diffusers_config(model=model_path) + logger.info("Diffusers Model config: %s", model_config) + # This is a placeholder - in a real implementation, you'd need to handle this properly + return None + else: + raise ValueError(f"Unsupported library: {self.library}") + + +class PipelineComponentLoader: + """Utility class for loading pipeline components. + + This replaces the chain of if-else statements in load_pipeline_module. + """ + + @staticmethod + def load_module( + module_name: str, + component_model_path: str, + transformers_or_diffusers: str, + fastvideo_args: FastVideoArgs, + ): + """Load a pipeline module. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + pipeline_args: Inference arguments + + Returns: + The loaded module + """ + logger.info( + "Loading %s using %s from %s", + module_name, + transformers_or_diffusers, + component_model_path, + ) + + # Get the appropriate loader for this module type + loader = ComponentLoader.for_module_type( + module_name, transformers_or_diffusers + ) + + # Load the module + return loader.load(component_model_path, fastvideo_args) diff --git a/worldcompass/fastvideo/models/loader/fsdp_load.py b/worldcompass/fastvideo/models/loader/fsdp_load.py new file mode 100644 index 0000000000000000000000000000000000000000..09ca11c87228ac2d866b94c0b8c783dc3177c69b --- /dev/null +++ b/worldcompass/fastvideo/models/loader/fsdp_load.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from torchtune +# Copyright 2024 The TorchTune Authors. +# Copyright 2025 The FastVideo Authors. + +import contextlib +from collections.abc import Callable, Generator +from itertools import chain +from typing import Any + +import torch +from torch import nn +from torch.distributed import DeviceMesh, init_device_mesh +from torch.distributed._tensor import distribute_tensor +from torch.distributed.fsdp import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + fully_shard, +) +from torch.nn.modules.module import _IncompatibleKeys + +from fastvideo.logger import init_logger +from fastvideo.models.loader.utils import ( + get_param_names_mapping, + hf_to_custom_state_dict, +) +from fastvideo.models.loader.weight_utils import safetensors_weights_iterator +from fastvideo.utils import set_mixed_precision_policy + +logger = init_logger(__name__) + + +# TODO(PY): move this to utils elsewhere +@contextlib.contextmanager +def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: + """Context manager to set torch's default dtype. + + Args: + dtype (torch.dtype): The desired default dtype inside the context manager. + + Returns: + ContextManager: context manager for setting default dtype. + + Example: + >>> with set_default_dtype(torch.bfloat16): + >>> x = torch.tensor([1, 2, 3]) + >>> x.dtype + torch.bfloat16 + """ + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(old_dtype) + + +# TODO(PY): add compile option +def maybe_load_fsdp_model( + load_from_dir: str, + ar_action_load_from_dir: str, + cls_name: str, + model_cls: type[nn.Module], + init_params: dict[str, Any], + weight_dir_list: list[str], + device: torch.device, + hsdp_replicate_dim: int, + hsdp_shard_dim: int, + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + cpu_offload: bool = False, + fsdp_inference: bool = False, + output_dtype: torch.dtype | None = None, + training_mode: bool = True, + pin_cpu_memory: bool = True, +) -> torch.nn.Module: + """Load the model with FSDP if is training, else load the model without FSDP.""" + # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are + # manually casting the inputs to the model + mp_policy = MixedPrecisionPolicy( + param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False + ) + + set_mixed_precision_policy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + + # with set_default_dtype(param_dtype), torch.device("meta"): + with set_default_dtype(param_dtype): + # model = model_cls(**init_params) + model = model_cls.from_pretrained( + load_from_dir, local_attn_size=-1, sink_size=0 + ) + if "Action" in cls_name: + model.add_discrete_action_parameters() + if ar_action_load_from_dir is not None: + if ".pt" in ar_action_load_from_dir: + state_dict = torch.load( + ar_action_load_from_dir, + map_location=lambda storage, loc: storage, + ) + state_dict = state_dict["generator"] + state_dict = { + ( + key.replace("model.", "") + if key.startswith("model.") + else key + ): value + for key, value in state_dict.items() + } + state_dict = { + ( + key.replace("_fsdp_wrapped_module.", "") + if key.startswith("_fsdp_wrapped_module.") + else key + ): value + for key, value in state_dict.items() + } + else: + from safetensors.torch import load_file + + state_dict = load_file(ar_action_load_from_dir) + + model.load_state_dict(state_dict, strict=True) + + if "Hunyuan" in cls_name: + model.to(torch.float32) + + # Check if we should use FSDP + use_fsdp = training_mode or fsdp_inference + + # Disable FSDP for MPS as it's not compatible + from fastvideo.platforms import current_platform + + if current_platform.is_mps(): + use_fsdp = False + logger.info("Disabling FSDP for MPS platform as it's not compatible") + + if use_fsdp: + world_size = hsdp_replicate_dim * hsdp_shard_dim + if not training_mode and not fsdp_inference: + hsdp_replicate_dim = world_size + hsdp_shard_dim = 1 + + device_mesh = init_device_mesh( + "cuda", + # (Replicate(), Shard(dim=0)) + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + shard_model( + model, + cpu_offload=cpu_offload, + reshard_after_forward=True, + mp_policy=mp_policy, + mesh=device_mesh, + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory, + ) + + # 加载逻辑不用这个 + # weight_iterator = safetensors_weights_iterator(weight_dir_list) + # param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) + # load_model_from_full_model_state_dict( + # model, + # weight_iterator, + # device, + # param_dtype, + # strict=True, + # cpu_offload=cpu_offload, + # param_names_mapping=param_names_mapping_fn, + # ) + for n, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError( + f"Unexpected param or buffer {n} on meta device." + ) + # Avoid unintended computation graph accumulation during inference + if isinstance(p, torch.nn.Parameter): + p.requires_grad = False + return model + + +def shard_model( + model, + *, + cpu_offload: bool, + reshard_after_forward: bool = True, + mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(), # noqa + mesh: DeviceMesh | None = None, + fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [], # noqa + pin_cpu_memory: bool = True, +) -> None: + """Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API. + + This method will over the model's named modules from the bottom-up and apply shard modules + based on whether they meet any of the criteria from shard_conditions. + + Args: + model (TransformerDecoder): Model to shard with FSDP. + shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine + which modules to shard with FSDP. Each function should take module name (relative to root) + and the module itself, returning True if FSDP should shard the module and False otherwise. + If any of shard_conditions return True for a given module, it will be sharded by FSDP. + cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer + states to CPU. + reshard_after_forward (bool): Whether to reshard parameters and buffers after + the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy + from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy. + mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism. + Default to None. + fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine + which modules to shard with FSDP. + pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters. + + Raises: + ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered. + """ + if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0: + logger.warning( + "The FSDP shard condition list is empty or None. No modules will be sharded in %s", + type(model).__name__, + ) + return + + fsdp_kwargs = { + "reshard_after_forward": reshard_after_forward, + "mesh": mesh, + "mp_policy": mp_policy, + } + if cpu_offload: + fsdp_kwargs["offload_policy"] = CPUOffloadPolicy( + pin_memory=pin_cpu_memory + ) + + # iterating in reverse to start with + # lowest-level modules first + num_layers_sharded = 0 + # TODO(will): don't reshard after forward for the last layer to save on the + # all-gather that will immediately happen Shard the model with FSDP, + for n, m in reversed(list(model.named_modules())): + if any( + [shard_condition(n, m) for shard_condition in fsdp_shard_conditions] + ): + fully_shard(m, **fsdp_kwargs) + num_layers_sharded += 1 + + if num_layers_sharded == 0: + raise ValueError( + "No layer modules were sharded. Please check if shard conditions are working as expected." + ) + + # Finally shard the entire model to account for any stragglers + fully_shard(model, **fsdp_kwargs) + + +# TODO(PY): device mesh for cfg parallel +def load_model_from_full_model_state_dict( + model: FSDPModule | torch.nn.Module, + full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None], + device: torch.device, + param_dtype: torch.dtype, + strict: bool = False, + cpu_offload: bool = False, + param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, + training_mode: bool = True, +) -> _IncompatibleKeys: + """ + Converting full state dict into a sharded state dict + and loading it into FSDP model (if training) or normal huggingface model + Args: + model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict + full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs + device (torch.device): device used to move full state dict tensors + param_dtype (torch.dtype): dtype used to move full state dict tensors + strict (bool): flag to check if to load the model in strict mode + cpu_offload (bool): flag to check if FSDP offload is enabled + param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name + training_mode (bool): apply FSDP only for training + Returns: + ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: + * **missing_keys** is a list of str containing the missing keys + * **unexpected_keys** is a list of str containing the unexpected keys + + Raises: + NotImplementedError: If got FSDP with more than 1D. + """ + meta_sd = model.state_dict() + sharded_sd = {} + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( + full_sd_iterator, param_names_mapping + ) # type: ignore + for target_param_name, full_tensor in custom_param_sd.items(): + meta_sharded_param = meta_sd.get(target_param_name) + if meta_sharded_param is None: + raise ValueError( + f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." + ) + if not hasattr(meta_sharded_param, "device_mesh"): + full_tensor = full_tensor.to(device=device, dtype=param_dtype) + # In cases where parts of the model aren't sharded, some parameters will be plain tensors + sharded_tensor = full_tensor + else: + full_tensor = full_tensor.to(device=device, dtype=param_dtype) + sharded_tensor = distribute_tensor( + full_tensor, + meta_sharded_param.device_mesh, + meta_sharded_param.placements, + ) + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() + sharded_sd[target_param_name] = nn.Parameter(sharded_tensor) + + model.reverse_param_names_mapping = reverse_param_names_mapping + unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) + if unused_keys: + logger.warning( + "Found unloaded parameters in meta state dict: %s", unused_keys + ) + + # List of allowed parameter name patterns + ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed + for new_param_name in unused_keys: + if not any( + pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS + ): + logger.error( + "Unsupported new parameter: %s. Allowed patterns: %s", + new_param_name, + ALLOWED_NEW_PARAM_PATTERNS, + ) + raise ValueError( + f"New parameter '{new_param_name}' is not supported. " + f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed." + ) + meta_sharded_param = meta_sd.get(new_param_name) + if not hasattr(meta_sharded_param, "device_mesh"): + # Initialize with zeros + sharded_tensor = torch.zeros_like( + meta_sharded_param, device=device, dtype=param_dtype + ) + else: + # Initialize with zeros and distribute + full_tensor = torch.zeros_like( + meta_sharded_param, device=device, dtype=param_dtype + ) + sharded_tensor = distribute_tensor( + full_tensor, + meta_sharded_param.device_mesh, + meta_sharded_param.placements, + ) + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() + sharded_sd[new_param_name] = nn.Parameter(sharded_tensor) + + # choose `assign=True` since we cannot call `copy_` on meta tensor + return model.load_state_dict(sharded_sd, strict=strict, assign=True) diff --git a/worldcompass/fastvideo/models/loader/utils.py b/worldcompass/fastvideo/models/loader/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7cbd3bc82109ac7e544131c509eecbfb112ee178 --- /dev/null +++ b/worldcompass/fastvideo/models/loader/utils.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for selecting and loading models.""" + +import contextlib +import re +from collections import defaultdict +from collections.abc import Callable, Iterator +from typing import Any + +import torch + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +@contextlib.contextmanager +def set_default_torch_dtype(dtype: torch.dtype): + """Sets the default torch dtype to the given dtype.""" + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + yield + torch.set_default_dtype(old_dtype) + + +def get_param_names_mapping( + mapping_dict: dict[str, str], +) -> Callable[[str], tuple[str, Any, Any]]: + """Creates a mapping function that transforms parameter names using regex patterns. + + Args: + mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns + param_name (str): The parameter name to be transformed + + Returns: + Callable[[str], str]: A function that maps parameter names from source to target format + """ + + def mapping_fn(name: str) -> tuple[str, Any, Any]: + # Try to match and transform the name using the regex patterns in mapping_dict + for pattern, replacement in mapping_dict.items(): + match = re.match(pattern, name) + if match: + merge_index = None + total_splitted_params = None + if isinstance(replacement, tuple): + merge_index = replacement[1] + total_splitted_params = replacement[2] + replacement = replacement[0] + name = re.sub(pattern, replacement, name) + return name, merge_index, total_splitted_params + + # If no pattern matches, return the original name + return name, None, None + + return mapping_fn + + +def hf_to_custom_state_dict( + hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], + param_names_mapping: Callable[[str], tuple[str, Any, Any]], +) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: + """Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary. + + Args: + hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary + param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format + + Returns: + custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict + reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf + """ + custom_param_sd = {} + to_merge_params = defaultdict(dict) # type: ignore + reverse_param_names_mapping = {} + if isinstance(hf_param_sd, dict): + hf_param_sd = hf_param_sd.items() # type: ignore + for source_param_name, full_tensor in hf_param_sd: # type: ignore + target_param_name, merge_index, num_params_to_merge = ( + param_names_mapping(source_param_name) + ) + reverse_param_names_mapping[target_param_name] = ( + source_param_name, + merge_index, + num_params_to_merge, + ) + if merge_index is not None: + to_merge_params[target_param_name][merge_index] = full_tensor + if len(to_merge_params[target_param_name]) == num_params_to_merge: + # cat at output dim according to the merge_index order + sorted_tensors = [ + to_merge_params[target_param_name][i] + for i in range(num_params_to_merge) + ] + full_tensor = torch.cat(sorted_tensors, dim=0) + del to_merge_params[target_param_name] + else: + continue + custom_param_sd[target_param_name] = full_tensor + return custom_param_sd, reverse_param_names_mapping diff --git a/worldcompass/fastvideo/models/loader/weight_utils.py b/worldcompass/fastvideo/models/loader/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e3d133fa08e14002c89e1fd54fd80c07d6a2358 --- /dev/null +++ b/worldcompass/fastvideo/models/loader/weight_utils.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py +"""Utilities for downloading and initializing model weights.""" + +import hashlib +import json +import os +import tempfile +from collections.abc import Generator +from pathlib import Path + +import filelock +import huggingface_hub.constants +import torch +from safetensors.torch import safe_open +from tqdm.auto import tqdm + +from fastvideo.distributed import get_local_torch_device +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# use system-level temp directory for file locks, so that multiple users +# can share the same lock without error. +# lock files in the temp directory will be automatically deleted when the +# system reboots, so users will not complain about annoying lock files +temp_dir = tempfile.gettempdir() + + +def enable_hf_transfer() -> None: + """Automatically activates hf_transfer.""" + if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: + try: + # enable hf hub transfer if available + import hf_transfer # type: ignore # noqa + + huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True + except ImportError: + pass + + +enable_hf_transfer() + + +class DisabledTqdm(tqdm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, disable=True) + + +def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): + lock_dir = cache_dir or temp_dir + model_name_or_path = str(model_name_or_path) + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +# For models like Mistral-7B-v0.3, there are both sharded +# safetensors files and a consolidated safetensors file. +# Passing both of these to the weight loader functionality breaks. +# So, we use the index_file to +# look up which safetensors files should be used. +def filter_duplicate_safetensors_files( + hf_weights_files: list[str], hf_folder: str, index_file: str +) -> list[str]: + # model.safetensors.index.json is a mapping from keys in the + # torch state_dict to safetensors file holding that weight. + index_file_name = os.path.join(hf_folder, index_file) + if not os.path.isfile(index_file_name): + return hf_weights_files + + # Iterate through the weight_map (weight_name: safetensors files) + # to identify weights that we should use. + with open(index_file_name) as f: + weight_map = json.load(f)["weight_map"] + weight_files_in_index = set() + for weight_name in weight_map: + weight_files_in_index.add( + os.path.join(hf_folder, weight_map[weight_name]) + ) + # Filter out any fields that are not found in the index file. + hf_weights_files = [ + f for f in hf_weights_files if f in weight_files_in_index + ] + return hf_weights_files + + +def filter_files_not_needed_for_inference( + hf_weights_files: list[str], +) -> list[str]: + """Exclude files that are not needed for inference. + + See + https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 + """ + blacklist = [ + "training_args.bin", + "optimizer.bin", + "optimizer.pt", + "scheduler.pt", + "scaler.pt", + ] + hf_weights_files = [ + f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist) + ] + return hf_weights_files + + +# explicitly use pure text format, with a newline at the end +# this makes it impossible to see the animation in the progress bar +# but will avoid messing up with ray or multiprocessing, which wraps +# each line of output with some prefix. +_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 + + +def safetensors_weights_iterator( + hf_weights_files: list[str], + to_cpu: bool = True, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files.""" + enable_tqdm = ( + not torch.distributed.is_initialized() + or torch.distributed.get_rank() == 0 + ) + device = "cpu" if to_cpu else str(get_local_torch_device()) + for st_file in tqdm( + hf_weights_files, + desc="Loading safetensors checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + with safe_open(st_file, framework="pt", device=device) as f: + for name in f.keys(): # noqa: SIM118 + param = f.get_tensor(name) + yield name, param + + +def pt_weights_iterator( + hf_weights_files: list[str], + to_cpu: bool = True, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model bin/pt files.""" + device = "cpu" if to_cpu else str(get_local_torch_device()) + enable_tqdm = ( + not torch.distributed.is_initialized() + or torch.distributed.get_rank() == 0 + ) + for bin_file in tqdm( + hf_weights_files, + desc="Loading pt checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + state = torch.load(bin_file, map_location=device, weights_only=True) + yield from state.items() + del state + + +def default_weight_loader( + param: torch.Tensor, loaded_weight: torch.Tensor +) -> None: + """Default weight loader.""" + try: + if param.numel() == 1 and loaded_weight.numel() == 1: + # Sometimes scalar values aren't considered tensors with shapes + # so if both param and loaded_weight are a scalar, + # "broadcast" instead of copy + param.data.fill_(loaded_weight.item()) + else: + assert param.size() == loaded_weight.size(), ( + f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})" + ) + + param.data.copy_(loaded_weight) + except Exception: + # NOTE: This exception is added for the purpose of setting breakpoint to + # debug weight loading issues. + raise + + +def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: + """Remap the name of FP8 k/v_scale parameters. + + This function handles the remapping of FP8 k/v_scale parameter names. + It detects if the given name ends with a suffix and attempts to remap + it to the expected name format in the model. If the remapped name is not + found in the params_dict, a warning is printed and None is returned. + + Args: + name (str): The original loaded checkpoint parameter name. + params_dict (dict): Dictionary containing the model's named parameters. + + Returns: + str: The remapped parameter name if successful, or the original name + if no remapping is needed. + None: If the remapped name is not found in params_dict. + """ + if name.endswith(".kv_scale"): + logger.warning_once( + "DEPRECATED. Found kv_scale in the checkpoint. " + "This format is deprecated in favor of separate k_scale and " + "v_scale tensors and will be removed in a future release. " + "Functionally, we will remap kv_scale to k_scale and duplicate " + "k_scale to v_scale" + ) + # NOTE: we remap the deprecated kv_scale to k_scale + remapped_name = name.replace(".kv_scale", ".attn.k_scale") + if remapped_name not in params_dict: + logger.warning_once( + f"Found kv_scale in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). kv_scale is " + "not loaded." + ) + return None + return remapped_name + + possible_scale_names = [".k_scale", ".v_scale"] + modelopt_scale_names = [ + ".self_attn.k_proj.k_scale", + ".self_attn.v_proj.v_scale", + ] + for scale_name in possible_scale_names: + if name.endswith(scale_name): + if any( + mo_scale_name in name for mo_scale_name in modelopt_scale_names + ): + remapped_name = name.replace( + f".self_attn.{scale_name[1]}_proj{scale_name}", + f".self_attn.attn{scale_name}", + ) + else: + remapped_name = name.replace(scale_name, f".attn{scale_name}") + if remapped_name not in params_dict: + logger.warning_once( + f"Found {scale_name} in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). {scale_name} is " + "not loaded." + ) + return None + return remapped_name + + # If there were no matches, return the untouched param name + return name diff --git a/worldcompass/fastvideo/models/parameter.py b/worldcompass/fastvideo/models/parameter.py new file mode 100644 index 0000000000000000000000000000000000000000..17163c8565e1547206336fa292a219a3df43e2ad --- /dev/null +++ b/worldcompass/fastvideo/models/parameter.py @@ -0,0 +1,427 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py + +from collections.abc import Callable +from fractions import Fraction +from typing import Any + +import torch +from torch.nn import Parameter + +from fastvideo.distributed import get_tp_rank +from fastvideo.logger import init_logger +from fastvideo.models.utils import _make_synced_weight_loader + +logger = init_logger(__name__) + + +class BasevLLMParameter(Parameter): + """Base parameter for vLLM linear layers. + + Extends the torch.nn.parameter by taking in a linear weight loader. Will copy the loaded weight + into the parameter when the provided weight loader is called. + """ + + def __new__(cls, data: torch.Tensor, **kwargs): + + return super().__new__(cls, data=data, requires_grad=False) + + def __init__(self, data: torch.Tensor, weight_loader: Callable): + """Initialize the BasevLLMParameter. + + :param data: torch tensor with the parameter data + :param weight_loader: weight loader callable + :returns: a torch.nn.parameter + """ + + # During weight loading, we often do something like: + # narrowed_tensor = param.data.narrow(0, offset, len) + # narrowed_tensor.copy_(real_weight) + # expecting narrowed_tensor and param.data to share the same storage. + # However, on TPUs, narrowed_tensor will lazily propagate to the base + # tensor, which is param.data, leading to the redundant memory usage. + # This sometimes causes OOM errors during model loading. To avoid this, + # we sync the param tensor after its weight loader is called. + from fastvideo.platforms import current_platform + + if current_platform.is_tpu(): + weight_loader = _make_synced_weight_loader(weight_loader) + + self._weight_loader = weight_loader + + @property + def weight_loader(self): + return self._weight_loader + + def _is_1d_and_scalar(self, loaded_weight: torch.Tensor): + cond1 = self.data.ndim == 1 and self.data.numel() == 1 + cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1 + return cond1 and cond2 + + def _assert_and_load(self, loaded_weight: torch.Tensor) -> None: + assert self.data.shape == loaded_weight.shape or self._is_1d_and_scalar( + loaded_weight + ) + self.data.copy_(loaded_weight) + + def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + self._assert_and_load(loaded_weight) + + def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + self._assert_and_load(loaded_weight) + + def load_merged_column_weight( + self, loaded_weight: torch.Tensor, **kwargs + ) -> None: + self._assert_and_load(loaded_weight) + + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: + self._assert_and_load(loaded_weight) + + +class _ColumnvLLMParameter(BasevLLMParameter): + """Private class defining weight loading functionality (load_merged_column_weight, + load_qkv_weight) for parameters being loaded into linear layers with column parallelism. + + This includes QKV and MLP layers which are not already fused on disk. Requires an output + dimension to be defined. Called within the weight loader of each of the column parallel linear + layers. + """ + + def __init__(self, output_dim: int, **kwargs): + self._output_dim = output_dim + super().__init__(**kwargs) + + @property + def output_dim(self): + return self._output_dim + + def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tp_rank() + shard_size = self.data.shape[self.output_dim] + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + def load_merged_column_weight( + self, loaded_weight: torch.Tensor, **kwargs + ) -> None: + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + if shard_offset is None or shard_size is None: + raise ValueError("shard_offset and shard_size must be provided") + if ( + isinstance(self, PackedColumnParameter | PackedvLLMParameter) + and self.packed_dim == self.output_dim + ): + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size + ) + + param_data = self.data + + tp_rank = get_tp_rank() + param_data = param_data.narrow( + self.output_dim, shard_offset, shard_size + ) + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + shard_id = kwargs.get("shard_id") + num_heads = kwargs.get("num_heads") + + assert shard_offset is not None + assert shard_size is not None + assert shard_id is not None + assert num_heads is not None + + if ( + isinstance(self, PackedColumnParameter | PackedvLLMParameter) + and self.output_dim == self.packed_dim + ): + shard_size, shard_offset = self.adjust_shard_indexes_for_packing( + shard_offset=shard_offset, shard_size=shard_size + ) + + param_data = self.data + tp_rank = get_tp_rank() + shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads + param_data = param_data.narrow( + self.output_dim, shard_offset, shard_size + ) + loaded_weight = loaded_weight.narrow( + self.output_dim, shard_id * shard_size, shard_size + ) + + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class RowvLLMParameter(BasevLLMParameter): + """Parameter class defining weight_loading functionality (load_row_parallel_weight) for + parameters being loaded into linear layers with row parallel functionality. + + Requires an input_dim to be defined. + """ + + def __init__(self, input_dim: int, **kwargs): + self._input_dim = input_dim + super().__init__(**kwargs) + + @property + def input_dim(self): + return self._input_dim + + def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tp_rank() + shard_size = self.data.shape[self.input_dim] + loaded_weight = loaded_weight.narrow( + self.input_dim, tp_rank * shard_size, shard_size + ) + + if len(loaded_weight.shape) == 0: + loaded_weight = loaded_weight.reshape(1) + + assert self.data.shape == loaded_weight.shape + self.data.copy_(loaded_weight) + + +class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): + """Parameter class for linear layer weights. + + Uses both column and row parallelism. + """ + + pass + + +class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """Parameter class for weight scales loaded for weights with grouped quantization. + + Uses both column and row parallelism. + """ + + pass + + +class ChannelQuantScaleParameter(_ColumnvLLMParameter): + """Parameter class for weight scales loaded for weights with channel-wise quantization. + + Equivalent to _ColumnvLLMParameter. + """ + + pass + + +class PerTensorScaleParameter(BasevLLMParameter): + """Parameter class for scales where the number of scales is equivalent to the number of logical + matrices in fused linear layers (e.g. for QKV, there are 3 scales loaded from disk). This is + relevant to weights with per-tensor quantization. Adds functionality to map the scalers to a + shard during weight loading. + + Note: additional parameter manipulation may be handled + for each quantization config specifically, within + process_weights_after_loading + """ + + def __init__(self, **kwargs): + self.qkv_idxs = {"q": 0, "k": 1, "v": 2} + super().__init__(**kwargs) + + def _shard_id_as_int(self, shard_id: str | int) -> int: + if isinstance(shard_id, int): + return shard_id + + # if not int, assume shard_id for qkv + # map to int and return + assert isinstance(shard_id, str) + assert shard_id in self.qkv_idxs + return self.qkv_idxs[shard_id] + + # For row parallel layers, no sharding needed + # load weight into parameter as is + def load_row_parallel_weight(self, *args, **kwargs) -> None: + super().load_row_parallel_weight(*args, **kwargs) + + def load_merged_column_weight(self, *args, **kwargs) -> None: + self._load_into_shard_id(*args, **kwargs) + + def load_qkv_weight(self, *args, **kwargs) -> None: + self._load_into_shard_id(*args, **kwargs) + + def load_column_parallel_weight(self, *args, **kwargs) -> None: + super().load_row_parallel_weight(*args, **kwargs) + + def _load_into_shard_id( + self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs + ): + """Slice the parameter data based on the shard id for loading.""" + + param_data = self.data + shard_id = self._shard_id_as_int(shard_id) + + # AutoFP8 scales do not have a shape + # compressed-tensors scales do have a shape + if len(loaded_weight.shape) != 0: + assert loaded_weight.shape[0] == 1 + loaded_weight = loaded_weight[0] + + param_data = param_data[shard_id] + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class PackedColumnParameter(_ColumnvLLMParameter): + """Parameter for model parameters which are packed on disk and support column parallelism only. + + See PackedvLLMParameter for more details on the packed properties. + """ + + def __init__( + self, packed_factor: int | Fraction, packed_dim: int, **kwargs + ): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + def adjust_shard_indexes_for_packing( + self, shard_size, shard_offset + ) -> tuple[Any, Any]: + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor, + ) + + +class PackedvLLMParameter(ModelWeightParameter): + """Parameter for model weights which are packed on disk. + + Example: GPTQ Marlin weights are int4 or int8, packed into int32. + Extends the ModelWeightParameter to take in the + packed factor, the packed dimension, and optionally, marlin + tile size for marlin kernels. Adjusts the shard_size and + shard_offset for fused linear layers model weight loading + by accounting for packing and optionally, marlin tile size. + """ + + def __init__( + self, packed_factor: int | Fraction, packed_dim: int, **kwargs + ): + self._packed_factor = packed_factor + self._packed_dim = packed_dim + super().__init__(**kwargs) + + @property + def packed_dim(self): + return self._packed_dim + + @property + def packed_factor(self): + return self._packed_factor + + def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): + return _adjust_shard_indexes_for_packing( + shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor, + ) + + +class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): + """Parameter class for weight scales loaded for weights with block-wise quantization. + + Uses both column and row parallelism. + """ + + pass + + +def permute_param_layout_( + param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs +) -> BasevLLMParameter: + """ + Permute a parameter's layout to the specified input and output dimensions, + useful for forcing the parameter into a known layout, for example, if I need + a packed (quantized) weight matrix to be in the layout + {input_dim = 0, output_dim = 1, packed_dim = 0} + then I can call: + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + to ensure x is in the correct layout (permuting it to the correct layout if + required, asserting if it cannot get it to the correct layout) + """ + + curr_input_dim = getattr(param, "input_dim", None) + curr_output_dim = getattr(param, "output_dim", None) + + if curr_input_dim is None or curr_output_dim is None: + assert param.data.dim() == 2, ( + "permute_param_layout_ only supports 2D parameters when either " + "input_dim or output_dim is not set" + ) + + # if one of the dimensions is not set, set it to the opposite of the other + # we can only do this since we asserted the parameter is 2D above + if curr_input_dim is None: + assert ( + curr_output_dim is not None + ), "either input or output dim must be set" + curr_input_dim = (curr_output_dim + 1) % 2 + if curr_output_dim is None: + assert ( + curr_input_dim is not None + ), "either input or output dim must be set" + curr_output_dim = (curr_input_dim + 1) % 2 + + # create permutation from the current layout to the layout with + # self.input_dim at input_dim and self.output_dim at output_dim preserving + # other dimensions + perm = [ + i + for i in range(param.data.dim()) + if i not in [curr_input_dim, curr_output_dim] + ] + perm.insert(input_dim, curr_input_dim) + perm.insert(output_dim, curr_output_dim) + + if "packed_dim" in kwargs: + assert ( + hasattr(param, "packed_dim") + and param.packed_dim == perm[kwargs["packed_dim"]] + ), "permute_param_layout_ currently doesn't support repacking" + + param.data = param.data.permute(*perm) + if hasattr(param, "_input_dim"): + param._input_dim = input_dim + if hasattr(param, "_output_dim"): + param._output_dim = output_dim + if "packed_dim" in kwargs and hasattr(param, "_packed_dim"): + param._packed_dim = kwargs["packed_dim"] + + return param + + +def _adjust_shard_indexes_for_packing( + shard_size, shard_offset, packed_factor +) -> tuple[Any, Any]: + shard_size = shard_size // packed_factor + shard_offset = shard_offset // packed_factor + return shard_size, shard_offset diff --git a/worldcompass/fastvideo/models/prope/camera_rope.py b/worldcompass/fastvideo/models/prope/camera_rope.py new file mode 100644 index 0000000000000000000000000000000000000000..5be166b3faac6cc771d65ce9201c5a36621675e0 --- /dev/null +++ b/worldcompass/fastvideo/models/prope/camera_rope.py @@ -0,0 +1,264 @@ +# MIT License +# +# Copyright (c) Authors of +# "PRoPE: Projective Positional Encoding for Multiview Transformers" +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# How to use PRoPE attention for self-attention: +# +# 1. Easiest way (fast): +# attn = PropeDotProductAttention(...) +# o = attn(q, k, v, viewmats, Ks) +# +# 2. More flexible way (fast): +# attn = PropeDotProductAttention(...) +# attn._precompute_and_cache_apply_fns(viewmats, Ks) +# q = attn._apply_to_q(q) +# k = attn._apply_to_kv(k) +# v = attn._apply_to_kv(v) +# o = F.scaled_dot_product_attention(q, k, v, **kwargs) +# o = attn._apply_to_o(o) +# +# 3. The most flexible way (but slower because repeated computation of RoPE coefficients): +# o = prope_dot_product_attention(q, k, v, ...) +# +# How to use PRoPE attention for cross-attention: +# +# attn_src = PropeDotProductAttention(...) +# attn_tgt = PropeDotProductAttention(...) +# attn_src._precompute_and_cache_apply_fns(viewmats_src, Ks_src) +# attn_tgt._precompute_and_cache_apply_fns(viewmats_tgt, Ks_tgt) +# q_src = attn_src._apply_to_q(q_src) +# k_tgt = attn_tgt._apply_to_kv(k_tgt) +# v_tgt = attn_tgt._apply_to_kv(v_tgt) +# o_src = F.scaled_dot_product_attention(q_src, k_tgt, v_tgt, **kwargs) +# o_src = attn_src._apply_to_o(o_src) + +from functools import partial +from typing import Callable, Optional, Tuple, List + +import torch +import torch.nn.functional as F + + +def prope_qkv( + q: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + k: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + v: torch.Tensor, # (batch, num_heads, seqlen, head_dim) + *, + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int = None, # How many patches wide is each image? + patches_y: int = None, # How many patches tall is each image? + image_width: int = None, # Width of the image. Used to normalize intrinsics. + image_height: int = None, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + mask: Optional[torch.Tensor] = None, + kv_cache=None, + is_cache: bool = False, + **kwargs, +) -> torch.Tensor: + """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style + positional encoding. + + Currently, we assume that the sequence length is equal to: + + cameras * patches_x * patches_y + + And token ordering allows the `(seqlen,)` axis to be reshaped into `(cameras, patches_x, + patches_y)`. + """ + # We're going to assume self-attention: all inputs are the same shape. + batch, num_heads, seqlen, head_dim = q.shape + cameras = viewmats.shape[1] + assert q.shape == k.shape == v.shape + assert viewmats.shape == (batch, cameras, 4, 4) + assert Ks is None or Ks.shape == (batch, cameras, 3, 3) + # assert seqlen == cameras * patches_x * patches_y + + apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns_all_dim( + head_dim=head_dim, + viewmats=viewmats, + Ks=Ks, + patches_x=patches_x, + patches_y=patches_y, + image_width=image_width, + image_height=image_height, + coeffs_x=coeffs_x, + coeffs_y=coeffs_y, + ) + + query = apply_fn_q(q) + key = apply_fn_kv(k) + value = apply_fn_kv(v) + + return query, key, value, apply_fn_o + + +def _prepare_apply_fns_all_dim( + head_dim: int, # Q/K/V will have this last dimension + viewmats: torch.Tensor, # (batch, cameras, 4, 4) + Ks: Optional[torch.Tensor], # (batch, cameras, 3, 3) + patches_x: int, # How many patches wide is each image? + patches_y: int, # How many patches tall is each image? + image_width: int, # Width of the image. Used to normalize intrinsics. + image_height: int, # Height of the image. Used to normalize intrinsics. + coeffs_x: Optional[torch.Tensor] = None, + coeffs_y: Optional[torch.Tensor] = None, +) -> Tuple[ + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], +]: + """Prepare transforms for PRoPE-style positional encoding.""" + device = viewmats.device + batch, cameras, _, _ = viewmats.shape + + # Normalize camera intrinsics. + if Ks is not None: + Ks_norm = torch.zeros_like(Ks) + Ks_norm[..., 0, 0] = Ks[..., 0, 0] + Ks_norm[..., 1, 1] = Ks[..., 1, 1] + Ks_norm[..., 0, 2] = 0 + Ks_norm[..., 1, 2] = 0 + Ks_norm[..., 2, 2] = 1.0 + Ks_norm = Ks_norm.to(dtype=Ks.dtype) + del Ks + + # Compute the camera projection matrices we use in PRoPE. + # - K is an `image<-camera` transform. + # - viewmats is a `camera<-world` transform. + # - P = lift(K) @ viewmats is an `image<-world` transform. + P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats) + P_T = P.transpose(-1, -2).to(dtype=viewmats.dtype) + P_inv = torch.einsum( + "...ij,...jk->...ik", + _invert_SE3(viewmats), + _lift_K(_invert_K(Ks_norm)), + ).to(dtype=viewmats.dtype) + + else: + # GTA formula. P is `camera<-world` transform. + P = viewmats + P_T = P.transpose(-1, -2) + P_inv = _invert_SE3(viewmats) + + assert P.shape == P_inv.shape == (batch, cameras, 4, 4) + + # Block-diagonal transforms to the inputs and outputs of the attention operator. + assert head_dim % 4 == 0 + transforms_q = [ + (partial(_apply_tiled_projmat, matrix=P_T), head_dim), + ] + transforms_kv = [ + (partial(_apply_tiled_projmat, matrix=P_inv), head_dim), + ] + transforms_o = [ + (partial(_apply_tiled_projmat, matrix=P), head_dim), + ] + + apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) + return apply_fn_q, apply_fn_kv, apply_fn_o + + +def _apply_tiled_projmat( + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + matrix: torch.Tensor, # (batch, cameras, D, D) +) -> torch.Tensor: + """Apply projection matrix to features.""" + # - seqlen => (cameras, patches_x * patches_y) + # - feat_dim => (feat_dim // 4, 4) + batch, num_heads, seqlen, feat_dim = feats.shape + cameras = matrix.shape[1] + assert seqlen >= cameras and seqlen % cameras == 0 + D = matrix.shape[-1] + assert matrix.shape == (batch, cameras, D, D) + assert feat_dim % D == 0 + return torch.einsum( + "bcij,bncpkj->bncpki", + matrix, + feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)), + ).reshape(feats.shape) + + +def _apply_block_diagonal( + feats: torch.Tensor, # (..., dim) + func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Apply a block-diagonal function to an input array. + + Each function is specified as a tuple with form: + + ((Tensor) -> Tensor, int) + + Where the integer is the size of the input to the function. + """ + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes) + x_blocks = torch.split(feats, block_sizes, dim=-1) + out = torch.cat( + [f(x_block) for f, x_block in zip(funcs, x_blocks)], + dim=-1, + ) + assert out.shape == feats.shape, "Input/output shapes should match." + return out + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix.""" + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum( + "...ij,...j->...i", Rinv, transforms[..., :3, 3] + ) + out[..., 3, 3] = 1.0 + out = out.to(dtype=transforms.dtype) + return out + + +def _lift_K(Ks: torch.Tensor) -> torch.Tensor: + """Lift 3x3 matrices to homogeneous 4x4 matrices.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device) + out[..., :3, :3] = Ks + out[..., 3, 3] = 1.0 + out = out.to(dtype=Ks.dtype) + return out + + +def _invert_K(Ks: torch.Tensor) -> torch.Tensor: + """Invert 3x3 intrinsics matrices. + + Assumes no skew. + """ + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros_like(Ks) + out[..., 0, 0] = 1.0 / Ks[..., 0, 0] + out[..., 1, 1] = 1.0 / Ks[..., 1, 1] + out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0] + out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1] + out[..., 2, 2] = 1.0 + out = out.to(dtype=Ks.dtype) + return out diff --git a/worldcompass/fastvideo/models/registry.py b/worldcompass/fastvideo/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..80d6f27dff08f268d6ffeab6bc5b34594f83fb69 --- /dev/null +++ b/worldcompass/fastvideo/models/registry.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py + +import importlib +import os +import pickle +import subprocess +import sys +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Callable, Set +from dataclasses import dataclass, field +from functools import lru_cache +from typing import NoReturn, TypeVar, cast + +import cloudpickle +from torch import nn + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# huggingface class name: (component_name, fastvideo module name, fastvideo class name) +_TEXT_TO_VIDEO_DIT_MODELS = { + "HunyuanTransformer3DARActionModel": ( + "hyvideo", + "models.transformers.ar_action_hunyuanvideo_1_5_transformer", + "ARHunyuanVideo_1_5_DiffusionTransformer", + ), +} + +_IMAGE_TO_VIDEO_DIT_MODELS = { + # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), + "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), +} + +_TEXT_ENCODER_MODELS = { + "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), + "LlamaModel": ("encoders", "llama", "LlamaModel"), + "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"), + "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"), + "BertModel": ("encoders", "clip", "CLIPTextModel"), +} + +_IMAGE_ENCODER_MODELS: dict[str, tuple] = { + # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), + "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), +} + +_VAE_MODELS = { + "AutoencoderKLHunyuanVideo": ( + "vaes", + "hunyuanvae", + "AutoencoderKLHunyuanVideo", + ), + "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), + "AutoencoderKLStepvideo": ( + "vaes", + "stepvideovae", + "AutoencoderKLStepvideo", + ), +} + +_SCHEDULERS = { + "FlowMatchEulerDiscreteScheduler": ( + "schedulers", + "scheduling_flow_match_euler_discrete", + "FlowMatchEulerDiscreteScheduler", + ), + "UniPCMultistepScheduler": ( + "schedulers", + "scheduling_unipc_multistep", + "UniPCMultistepScheduler", + ), +} + +_FAST_VIDEO_MODELS = { + **_TEXT_TO_VIDEO_DIT_MODELS, + **_IMAGE_TO_VIDEO_DIT_MODELS, + **_TEXT_ENCODER_MODELS, + **_IMAGE_ENCODER_MODELS, + **_VAE_MODELS, + **_SCHEDULERS, +} + +_SUBPROCESS_COMMAND = [sys.executable, "-m", "fastvideo.models.dits.registry"] + +_T = TypeVar("_T") + + +@dataclass(frozen=True) +class _ModelInfo: + architecture: str + + @staticmethod + def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": + return _ModelInfo( + architecture=model.__name__, + ) + + +class _BaseRegisteredModel(ABC): + @abstractmethod + def inspect_model_cls(self) -> _ModelInfo: + raise NotImplementedError + + @abstractmethod + def load_model_cls(self) -> type[nn.Module]: + raise NotImplementedError + + +@dataclass(frozen=True) +class _RegisteredModel(_BaseRegisteredModel): + """Represents a model that has already been imported in the main process.""" + + interfaces: _ModelInfo + model_cls: type[nn.Module] + + @staticmethod + def from_model_cls(model_cls: type[nn.Module]): + return _RegisteredModel( + interfaces=_ModelInfo.from_model_cls(model_cls), + model_cls=model_cls, + ) + + def inspect_model_cls(self) -> _ModelInfo: + return self.interfaces + + def load_model_cls(self) -> type[nn.Module]: + return self.model_cls + + +def _run_in_subprocess(fn: Callable[[], _T]) -> _T: + # NOTE: We use a temporary directory instead of a temporary file to avoid + # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file + with tempfile.TemporaryDirectory() as tempdir: + output_filepath = os.path.join(tempdir, "registry_output.tmp") + + # `cloudpickle` allows pickling lambda functions directly + input_bytes = cloudpickle.dumps((fn, output_filepath)) + + # cannot use `sys.executable __file__` here because the script + # contains relative imports + returned = subprocess.run( + _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True + ) + + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError( + f"Error raised in subprocess:\n{returned.stderr.decode()}" + ) from e + + with open(output_filepath, "rb") as f: + return cast(_T, pickle.load(f)) + + +@dataclass(frozen=True) +class _LazyRegisteredModel(_BaseRegisteredModel): + """Represents a model that has not been imported in the main process.""" + + module_name: str + component_name: str + class_name: str + + # Performed in another process to avoid initializing CUDA + def inspect_model_cls(self) -> _ModelInfo: + return _run_in_subprocess( + lambda: _ModelInfo.from_model_cls(self.load_model_cls()) + ) + + def load_model_cls(self) -> type[nn.Module]: + mod = importlib.import_module(self.module_name) + return cast(type[nn.Module], getattr(mod, self.class_name)) + + +@lru_cache(maxsize=128) +def _try_load_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> type[nn.Module] | None: + from fastvideo.platforms import current_platform + + current_platform.verify_model_arch(model_arch) + try: + return model.load_model_cls() + except Exception: + logger.exception("Error in loading model architecture '%s'", model_arch) + return None + + +@lru_cache(maxsize=128) +def _try_inspect_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> _ModelInfo | None: + try: + return model.inspect_model_cls() + except Exception: + logger.exception( + "Error in inspecting model architecture '%s'", model_arch + ) + return None + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) + + def get_supported_archs(self) -> Set[str]: + return self.models.keys() + + def register_model( + self, + model_arch: str, + model_cls: type[nn.Module] | str, + ) -> None: + """Register an external model to be used in vLLM. + + :code:`model_cls` can be either: + + - A :class:`torch.nn.Module` class directly referencing the model. + - A string in the format :code:`:` which can be used to + lazily import the model. This is useful to avoid initializing CUDA + when importing the model and thus the related error + :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`. + """ + if model_arch in self.models: + logger.warning( + "Model architecture %s is already registered, and will be " + "overwritten by the new model class %s.", + model_arch, + model_cls, + ) + + if isinstance(model_cls, str): + split_str = model_cls.split(":") + if len(split_str) != 2: + msg = "Expected a string in the format `:`" + raise ValueError(msg) + + model = _LazyRegisteredModel(*split_str) + else: + model = _RegisteredModel.from_model_cls(model_cls) + + self.models[model_arch] = model + + def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn: + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: + if model_arch not in self.models: + return None + + return _try_load_model_cls(model_arch, self.models[model_arch]) + + def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None: + if model_arch not in self.models: + return None + + return _try_inspect_model_cls(model_arch, self.models[model_arch]) + + def _normalize_archs( + self, + architectures: str | list[str], + ) -> list[str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + logger.warning("No model architectures are specified") + + normalized_arch = [] + for model in architectures: + if model not in self.models: + model = "TransformersModel" + normalized_arch.append(model) + return normalized_arch + + def inspect_model_cls( + self, + architectures: str | list[str], + ) -> tuple[_ModelInfo, str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + return self._raise_for_unsupported(architectures) + + def resolve_model_cls( + self, + architectures: str | list[str], + ) -> tuple[type[nn.Module], str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + +ModelRegistry = _ModelRegistry( + { + model_arch: _LazyRegisteredModel( + module_name=f"fastvideo.models.{component_name}.{mod_relname}", + component_name=component_name, + class_name=cls_name, + ) + for model_arch, ( + component_name, + mod_relname, + cls_name, + ) in _FAST_VIDEO_MODELS.items() + } +) diff --git a/worldcompass/fastvideo/models/utils.py b/worldcompass/fastvideo/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6a21fd64f0e2217a5e3012f121fe37dcd2761b61 --- /dev/null +++ b/worldcompass/fastvideo/models/utils.py @@ -0,0 +1,780 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/utils.py +"""Utils for model executor.""" + +from typing import Any + +import torch + + +# TODO(PY): move it elsewhere +def auto_attributes(init_func): + """Decorator that automatically adds all initialization arguments as object attributes. + + Example: + @auto_attributes + def __init__(self, a=1, b=2): + pass + + # This will automatically set: + # - self.a = 1 and self.b = 2 + # - self.config.a = 1 and self.config.b = 2 + """ + + def wrapper(self, *args, **kwargs): + # Get the function signature + import inspect + + signature = inspect.signature(init_func) + parameters = signature.parameters + + # Get parameter names (excluding 'self') + param_names = list(parameters.keys())[1:] + + # Bind arguments to parameters + bound_args = signature.bind(self, *args, **kwargs) + bound_args.apply_defaults() + + # Create config object if it doesn't exist + if not hasattr(self, "config"): + self.config = type("Config", (), {})() + + # Set attributes on self and self.config + for name in param_names: + if name in bound_args.arguments: + value = bound_args.arguments[name] + setattr(self, name, value) + setattr(self.config, name, value) + + # Call the original __init__ function + return init_func(self, *args, **kwargs) + + return wrapper + + +def set_weight_attrs( + weight: torch.Tensor, + weight_attrs: dict[str, Any] | None, +): + """Set attributes on a weight tensor. + + This method is used to set attributes on a weight tensor. This method + will not overwrite existing attributes. + + Args: + weight: The weight tensor. + weight_attrs: A dictionary of attributes to set on the weight tensor. + """ + if weight_attrs is None: + return + for key, value in weight_attrs.items(): + assert not hasattr( + weight, key + ), f"Overwriting existing tensor attribute: {key}" + + # NOTE(woosuk): During weight loading, we often do something like: + # narrowed_tensor = param.data.narrow(0, offset, len) + # narrowed_tensor.copy_(real_weight) + # expecting narrowed_tensor and param.data to share the same storage. + # However, on TPUs, narrowed_tensor will lazily propagate to the base + # tensor, which is param.data, leading to the redundant memory usage. + # This sometimes causes OOM errors during model loading. To avoid this, + # we sync the param tensor after its weight loader is called. + # TODO(woosuk): Remove this hack once we have a better solution. + from fastvideo.platforms import current_platform + + if current_platform.is_tpu() and key == "weight_loader": + value = _make_synced_weight_loader(value) + setattr(weight, key, value) + + +def _make_synced_weight_loader(original_weight_loader) -> Any: + + def _synced_weight_loader(param, *args, **kwargs): + original_weight_loader(param, *args, **kwargs) + torch._sync(param) + + return _synced_weight_loader + + +def extract_layer_index(layer_name: str) -> int: + """Extract the layer index from the module name. + + Examples: + - "encoder.layers.0" -> 0 + - "encoder.layers.1.self_attn" -> 1 + - "2.self_attn" -> 2 + - "model.encoder.layers.0.sub.1" -> ValueError + """ + subnames = layer_name.split(".") + int_vals: list[int] = [] + for subname in subnames: + try: + int_vals.append(int(subname)) + except ValueError: + continue + assert ( + len(int_vals) == 1 + ), f"layer name {layer_name} should only contain one integer" + return int_vals[0] + + +def modulate( + x: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, +) -> torch.Tensor: + """Modulate by shift and scale. + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) # type: ignore[union-attr] + elif scale is None: + return x + shift.unsqueeze(1) # type: ignore[union-attr] + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) # type: ignore[union-attr] + + +import json +import os +import sys + +sys.path.append(os.path.abspath(".")) +import torch +import pandas as pd +import numpy as np +import random +from pathlib import Path +from typing import List, Tuple, Dict +from scipy.spatial.transform import Rotation as R_scipy +import math + +from torch.utils.data import Dataset, Sampler +from torchdata.stateful_dataloader import StatefulDataLoader +from scipy.spatial.transform import Rotation as R +from fastvideo.distributed import ( + get_local_torch_device, + maybe_init_distributed_environment_and_model_parallel, +) +from fastvideo.distributed import ( + get_sp_world_size, + get_world_group, + get_world_rank, + get_world_size, +) + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def camera_center_normalization(w2c, nframe): + c2w_view0 = w2c[::nframe].inverse() # [B,4,4] + c2w_view0 = c2w_view0.repeat_interleave(nframe, dim=0) # [BF,4,4] + w2c = c2w_view0 @ w2c + return w2c + + +# action keys +ACTION_KEYS = [ + "D", + "DL", + "DR", +] + + +# *******************************Add the memory part******************************** +def calculate_pose_distance_from_w2c( + w2c_1: np.ndarray, + w2c_2: np.ndarray, + pos_weight: float = 1.0, + ang_weight: float = 1.0, +) -> float: + """根据两个 4x4 W2C (World-to-Camera) 矩阵计算它们之间的综合姿态距离。 + + 该距离量化了两个相机姿态的相似度，类似其 FOV 的重叠程度。 + + 参数: + w2c_1 (np.ndarray): 第一个相机的 4x4 World-to-Camera 矩阵。 + w2c_2 (np.ndarray): 第二个相机的 4x4 World-to-Camera 矩阵。 + pos_weight (float): 空间距离的权重。 + ang_weight (float): 角度距离的权重。 + + 返回: + float: 两个姿态之间的综合距离。 + """ + + def w2c_to_6d_pose(w2c_matrix: np.ndarray) -> np.ndarray: + """将 4x4 World-to-Camera (W2C) 矩阵转换为 6D 姿态。 + + 6D 姿态元组为 (x, y, z, pitch, yaw, roll)。 + """ + # 提取旋转矩阵 R 和平移向量 t + R_cw = w2c_matrix[:3, :3] + t_cw = w2c_matrix[:3, 3] + + # 计算相机在世界坐标系下的位置 C_world + # C_world = -R_cw.T @ t_cw + C_world = -np.dot(R_cw.T, t_cw) + + # 将旋转矩阵转换为欧拉角 (pitch, yaw, roll) + # 注意: scipy 默认的欧拉角顺序是 ZYX，对应 yaw, pitch, roll + # 为了与常见的 (pitch, yaw, roll) 顺序匹配，我们手动转换 + r = R_scipy.from_matrix(R_cw) + pitch, yaw, roll = r.as_euler("yxz", degrees=True) + + return np.array([C_world[0], C_world[1], C_world[2], pitch, yaw, roll]) + + # 1. 将两个 W2C 矩阵转换为 6D 姿态 + pose1_6d = w2c_to_6d_pose(w2c_1) + pose2_6d = w2c_to_6d_pose(w2c_2) + + # 2. 计算空间距离 (欧几里得距离) + pos1 = pose1_6d[:3] + pos2 = pose2_6d[:3] + spatial_distance = np.linalg.norm(pos1 - pos2) + + # 3. 计算角度距离 (考虑圆周特性) + angles1 = pose1_6d[3:] + angles2 = pose2_6d[3:] + + angle_diff = np.abs(angles1 - angles2) + # 修正角度差，确保是最小的圆周距离 + angular_distance_vector = np.minimum(angle_diff, 360 - angle_diff) + # 使用欧几里得范数作为综合角度距离 + angular_distance = np.linalg.norm(angular_distance_vector) + + # 4. 结合两种距离得到综合姿态距离 + total_distance = ( + pos_weight * spatial_distance + ang_weight * angular_distance + ) + + return total_distance + + +# TO DO: add a chunk mechanism to chunk the long sequence to the window size +# --- 新的辅助函数：计算复杂片段距离 --- +def calculate_complex_clip_distance( + w2c_list: List[np.ndarray], + query_clip_indices: List[int], + historical_clip_indices: List[int], + pos_weight: float = 1.0, + ang_weight: float = 1.0, +) -> float: + """计算查询片段与历史片段之间的复杂姿态距离。 + + 该距离是基于查询片段的第二帧和第四帧与历史片段的每一帧的平均距离。 + """ + # 1. 确定查询片段的采样帧索引 + # 从 query_clip_indices 的第二个元素 (索引 1) 开始，每隔两帧选取 + # 例如：如果 query_clip_indices = [10, 11, 12, 13, 14, 15] + # 采样索引为：11, 13, 15 + + # 确保查询片段至少有 2 帧才能进行采样 + if len(query_clip_indices) < 2: + # 如果预测序列太短，使用第一个和最后一个帧作为样本 + sample_indices = [query_clip_indices[0], query_clip_indices[-1]] + else: + # 采样点在 query_clip_indices 中的局部索引：从 1 开始，步长为 2 + # np.arange(1, len(query_clip_indices), 2) + sample_indices = [ + query_clip_indices[i] + for i in np.arange(1, len(query_clip_indices), 2) + ] + + total_avg_distance = 0.0 + + # 2. 遍历所有采样帧 + for query_idx in sample_indices: + query_w2c = w2c_list[query_idx] + + dists_from_query_frame = [] + + # 3. 计算该采样帧与历史片段中每一帧的距离 + for hist_idx in historical_clip_indices: + hist_w2c = w2c_list[hist_idx] + dist = calculate_pose_distance_from_w2c( + query_w2c, hist_w2c, pos_weight, ang_weight + ) + dists_from_query_frame.append(dist) + + # 累加这个采样帧到历史片段的平均距离 + total_avg_distance += np.mean(dists_from_query_frame) + + # 4. 最终距离是所有采样帧平均距离的再次平均 + final_clip_distance = total_avg_distance / len(sample_indices) + + return final_clip_distance + + +# --- 工具函数 1: 角度转换 (从旋转矩阵到 Pitch/Yaw) --- +def rotation_matrix_to_angles( + R: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """从相机坐标系的 3x3 旋转矩阵 R 估算 Pitch 和 Yaw 角度。 + + 假设相机坐标系: X=右, Y=上, Z=后退 (或 NeRF style: X=右, Y=下, Z=前) + 这里我们使用常见的计算机视觉约定：Z轴为向前。 + + 注意：这里的角度计算直接基于您的 is_inside_fov_3d_hv 函数的约定： + - Yaw/Azimuth 角度在 XZ 平面 (atan2(x, z))。 + - Pitch/Elevation 角度相对于水平面 (atan2(y, sqrt(x^2 + z^2)))。 + + 对于 W2C 矩阵 R 的第三列 R[:, 2] (世界坐标系中Z轴在相机坐标系下的方向)， + 通常对应于相机看向的方向（z轴在世界坐标系中的表示）。 + + 为了简化并匹配您的 is_inside_fov 逻辑，我们直接使用 R 的 Z 轴向量： + Camera Z-axis direction in World Frame (Forward Vector): fwd = R_w2c_inv @ [0, 0, 1] + 更简单地，C2W 矩阵的 Z 轴向量即为相机在世界坐标系下的前向向量。 + C2W = W2C_inv + """ + + # 1. 估算前向向量（在世界坐标系中） + # R 是 R_c2w.T，我们想要 R_c2w 的 Z 轴。 + # 对于 R_w2c，它的逆 R_c2w 是 R_w2c.T + R_c2w = R.T + + # 前向向量是 C2W 矩阵的第三列（Z轴方向） + fwd = R_c2w[:, 2] # 形状 (3,) + + x = fwd[0] + y = fwd[1] + z = fwd[2] + + # 2. 计算 Yaw (偏航角) + # Yaw = atan2(x, z) - 假设 Z 是前向，X 是左右 + yaw_rad = torch.atan2(x, z) + yaw_deg = yaw_rad * (180.0 / math.pi) + + # 3. 计算 Pitch (俯仰角) + # Pitch = atan2(y, sqrt(x^2 + z^2)) + pitch_rad = torch.atan2(y, torch.sqrt(x**2 + z**2)) + pitch_deg = pitch_rad * (180.0 / math.pi) + + return pitch_deg, yaw_deg + + +def generate_points_in_sphere(n_points: int, radius: float) -> torch.Tensor: + """在指定半径的球体内均匀生成点。 + + :param n_points: 要生成的点数。 + :param radius: 球体半径。 + :return: 形状为 (n_points, 3) 的张量，表示点的 (x, y, z) 坐标。 + """ + samples_r = torch.rand(n_points) + samples_phi = torch.rand(n_points) + samples_u = torch.rand(n_points) + + # 均匀体积采样：r = R * u^(1/3) + r = radius * torch.pow(samples_r, 1 / 3) + phi = 2 * math.pi * samples_phi + # 均匀极角采样：theta = arccos(1 - 2*u) + theta = torch.acos(1 - 2 * samples_u) + + # 转换为笛卡尔坐标 + x = r * torch.sin(theta) * torch.cos(phi) + y = r * torch.sin(theta) * torch.sin(phi) + z = r * torch.cos(theta) + + points = torch.stack((x, y, z), dim=1) + return points + + +def is_inside_fov_3d_hv( + points: torch.Tensor, + center: torch.Tensor, + center_pitch: torch.Tensor, + center_yaw: torch.Tensor, + fov_half_h: torch.Tensor, + fov_half_v: torch.Tensor, +) -> torch.Tensor: + """检查点是否在由中心坐标、俯仰角和偏航角定义的 3D 视锥体内。 + + :param points: (N, 3) 或 (N, B, 3) 采样点坐标。 + :param center: (3) 或 (B, 3) 相机中心坐标。 + :param center_pitch: (1) 或 (B) 中心视线的俯仰角（度）。 + :param center_yaw: (1) 或 (B) 中心视线的偏航角（度）。 + :param fov_half_h: 水平半视场角（度）。 + :param fov_half_v: 垂直半视场角（度）。 + :return: (N) 或 (N, B) 布尔张量，指示点是否在 FOV 内。 + """ + # 确保 center, pitch, yaw 维度可广播 + if points.ndim == 2: # N, 3 + vectors = points - center[None, :] + C = 1 # 批次大小为1 + elif points.ndim == 3: # N, B, 3 + vectors = points - center[None, ...] + # 确保 center_pitch/yaw 是 (1, B) 形状进行广播 + center_pitch = ( + center_pitch[None, :] if center_pitch.ndim == 1 else center_pitch + ) + center_yaw = center_yaw[None, :] if center_yaw.ndim == 1 else center_yaw + else: + raise ValueError("points 必须是 (N, 3) 或 (N, B, 3) 形状") + + x = vectors[..., 0] + y = vectors[..., 1] + z = vectors[..., 2] + + # 计算水平角 (偏航/方位角)，假设 Z 轴为前向 + azimuth = torch.atan2(x, z) * (180 / math.pi) + + # 计算垂直角 (俯仰/仰角) + elevation = torch.atan2(y, torch.sqrt(x**2 + z**2)) * (180 / math.pi) + + # 计算与中心视线的角度差 (处理角度环绕) + diff_azimuth = azimuth - center_yaw + # 将角度差归一化到 [-180, 180] 范围 + diff_azimuth = torch.remainder(diff_azimuth + 180, 360) - 180 + + diff_elevation = elevation - center_pitch + diff_elevation = torch.remainder(diff_elevation + 180, 360) - 180 + + # 检查是否在 FOV 限制内 + in_fov_h = diff_azimuth.abs() < fov_half_h + in_fov_v = diff_elevation.abs() < fov_half_v + + return in_fov_h & in_fov_v + + +# --- 工具函数 2: 欧氏距离 (基于 W2C 的位置) --- +def calculate_euclidean_distance( + w2c_matrix_1: torch.Tensor, w2c_matrix_2: torch.Tensor +) -> torch.Tensor: + """计算两个相机中心位置之间的欧氏距离。""" + # 相机在世界坐标系中的位置 P_w = -R^T * t + R1 = w2c_matrix_1[:3, :3] + t1 = w2c_matrix_1[:3, 3] + R2 = w2c_matrix_2[:3, :3] + t2 = w2c_matrix_2[:3, 3] + + P_w1 = -R1.T @ t1 + P_w2 = -R2.T @ t2 + + distance = torch.linalg.norm(P_w1 - P_w2) + return distance + + +def calculate_fov_overlap_similarity_refine( + w2c_matrix_curr: torch.Tensor, + w2c_matrix_hist: torch.Tensor, + fov_h_deg: float = 105.0, + fov_v_deg: float = 75.0, + device=None, + points_local=None, +) -> float: + """使用蒙特卡洛采样计算两个 W2C 位姿之间的 FOV 重叠相似性。 + + 相似性 = (Curr_FOV ∩ Hist_FOV) 的点数 / (Curr_FOV) 的点数。 + + :param w2c_matrix_curr: 当前帧的 (4, 4) W2C 矩阵。 + :param w2c_matrix_hist: 历史帧的 (4, 4) W2C 矩阵。 :param num_samples, radius, fov_h_deg, fov_v_deg: + 采样和 FOV 参数。 + :return: 重叠比率 (0.0 到 1.0 之间的浮点数)。 + """ + w2c_matrix_curr = torch.tensor(w2c_matrix_curr, device=device) + w2c_matrix_hist = torch.tensor(w2c_matrix_hist, device=device) + + # 转换为相对pose + c2w_matrix_curr = torch.linalg.inv(w2c_matrix_curr) + c2w_matrix_hist = torch.linalg.inv(w2c_matrix_hist) + C_inv = w2c_matrix_curr + + w2c_matrix_curr = torch.linalg.inv(C_inv @ c2w_matrix_curr) + w2c_matrix_hist = torch.linalg.inv(C_inv @ c2w_matrix_hist) + + # --- 1. W2C 矩阵解析为位置和角度 --- + + # P_w = -R^T * t + R_curr, t_curr = w2c_matrix_curr[:3, :3], w2c_matrix_curr[:3, 3] + R_hist, t_hist = w2c_matrix_hist[:3, :3], w2c_matrix_hist[:3, 3] + + P_w_curr = -R_curr.T @ t_curr + P_w_hist = -R_hist.T @ t_hist + + # pitch, yaw 角度 (度) + pitch_curr, yaw_curr = rotation_matrix_to_angles(R_curr) + pitch_hist, yaw_hist = rotation_matrix_to_angles(R_hist) + + # FOV 参数 + fov_half_h = torch.tensor(fov_h_deg / 2.0, device=device) + fov_half_v = torch.tensor(fov_v_deg / 2.0, device=device) + + # 将点平移到当前相机中心 P_w_curr (N, 3) + # 对应您代码中的 points += pose_conditions[curr_frame, :, :3][None] 的简化版本 (B=1) + points_world = points_local + P_w_curr[None, :] + + # --- 3. FOV 检查 --- + + # 检查点是否在 Curr FOV 内 (作为分母的基准) + in_fov_curr = is_inside_fov_3d_hv( + points_world, + P_w_curr[None, :], + pitch_curr[None], + yaw_curr[None], + fov_half_h, + fov_half_v, + ) + + # 检查点是否在 Hist FOV 内, 只根据角度进行计算 + in_fov_hist = is_inside_fov_3d_hv( + points_world, + P_w_hist[None, :], + pitch_hist[None], + yaw_hist[None], + fov_half_h, + fov_half_v, + ) + + # 根据点到点的距离在计算一下 + dist = torch.norm(points_world - P_w_hist.reshape(1, -1), dim=1) < 8.0 + in_fov_hist = in_fov_hist.bool() & dist.reshape(1, -1).bool() + # --- 4. 计算重叠比率 --- + + # 分子: 交集中的点数 + overlap_count = (in_fov_curr.bool() & in_fov_hist.bool()).sum().float() + + # 分母: Curr FOV 中的总点数 + fov_curr_count = in_fov_curr.sum().float() + # print(overlap_count, fov_curr_count) + + if fov_curr_count == 0: + return 0.0 # 避免除以零 + + # 重叠比率 + overlap_ratio = overlap_count / fov_curr_count + + return overlap_ratio.item() + + +def select_aligned_memory_frames_context_per_chunk_w_latent_sink_fov_refine_hunyuan( + w2c_list: List[np.ndarray], + current_frame_idx: int, + memory_frames: int, + temporal_context_size: int, + pred_latent_size: int, + pos_weight: float = 1.0, + ang_weight: float = 1.0, + device=None, + points_local=None, +) -> List[int]: + """为给定帧选择记忆帧和上下文帧，基于复杂的四帧片段距离计算。 + + 参数: w2c_list (List[np.ndarray]): 包含所有N个4x4外参矩阵的列表。 current_frame_idx (int): 当前要处理的帧的索引。 + memory_frames (int): 需要选择的记忆帧总数。 context_size (int): 需要选择的上下文帧总数。 pos_weight (float): + 空间距离的权重。 ang_weight (float): 角度距离的权重。 + + 返回: List[int]: 包含选定记忆帧和上下文帧索引的列表。 + """ + if current_frame_idx <= memory_frames: + return list(range(0, current_frame_idx)) + + num_total_frames = len(w2c_list) + # 检查当前帧是否能构成一个完整的4帧片段 + if current_frame_idx >= num_total_frames or current_frame_idx < 3: + raise ValueError( + f"当前帧索引必须在 w2c_list 的有效范围内，且至少为3。{current_frame_idx}, {len(w2c_list)}" + ) + + # 1. 选择上下文帧 (Context Frames) + start_context_idx = max(0, current_frame_idx - temporal_context_size) + context_frames_indices = list(range(start_context_idx, current_frame_idx)) + + # 2. 计算记忆帧 (Memory Frames) 的候选池 + candidate_distances = [] + query_clip_indices = list( + range( + current_frame_idx, + ( + current_frame_idx + pred_latent_size + if current_frame_idx + pred_latent_size <= num_total_frames + else num_total_frames + ), + ) + ) + + historical_clip_indices = list( + range(4, current_frame_idx - temporal_context_size, 4) + ) + + # 3. 选取最相似的 `memory_frames` 个帧 + memory_frames_indices = [ + 0, + 1, + 2, + 3, + ] # add the first latent frame as context + memory_frames = ( + memory_frames - temporal_context_size - len(memory_frames_indices) + ) + + # 遍历所有历史片段，将每个片段作为记忆帧候选 + # 历史片段的起始索引必须是 4 的倍数，且不能与上下文帧重叠 + for hist_idx in historical_clip_indices: + total_dist = 0 + hist_w2c_1 = w2c_list[hist_idx] + hist_w2c_2 = w2c_list[hist_idx + 2] + for query_idx in query_clip_indices: + dist_1_for_query_idx = ( + 1.0 + - calculate_fov_overlap_similarity_refine( + w2c_list[query_idx], + hist_w2c_1, + fov_h_deg=60.0, + fov_v_deg=35.0, + device=device, + points_local=points_local, + ) + ) + dist_2_for_query_idx = ( + 1.0 + - calculate_fov_overlap_similarity_refine( + w2c_list[query_idx], + hist_w2c_2, + fov_h_deg=60.0, + fov_v_deg=35.0, + device=device, + points_local=points_local, + ) + ) + dist_for_query_idx = ( + dist_1_for_query_idx + dist_2_for_query_idx + ) / 2.0 + total_dist += dist_for_query_idx + # 3. 将多个平均值再次取平均，得到最终的片段距离 + final_clip_distance = total_dist / len(query_clip_indices) + # 存储 (片段起始帧索引, 平均距离) + candidate_distances.append((hist_idx, final_clip_distance)) + + # 按平均距离从小到大排序 + candidate_distances.sort(key=lambda x: x[1]) + + # 遍历排序后的候选片段，直到收集到足够的记忆帧 + for start_idx, _ in candidate_distances: + if start_idx not in memory_frames_indices: + memory_frames_indices.extend(range(start_idx, start_idx + 4)) + + # 检查是否已达到 memory_size 的要求 + if len(memory_frames_indices) >= memory_frames: + break + + # 4. 组合并去重，以确保没有重复的帧 + selected_frames_set = set(context_frames_indices) + selected_frames_set.update(memory_frames_indices) + + final_selected_frames = sorted(list(selected_frames_set)) + # assert len(final_selected_frames) == memory_frames + temporal_context_size + return final_selected_frames + + +# --- 新函数：处理所有帧 --- +def process_all_frames_for_memory( + w2c_list: List[np.ndarray], + window_size: int, + memory_frames: int, + pos_weight: float = 1.0, + ang_weight: float = 1.0, +) -> Dict[int, List[int]]: + """遍历超过 window_size 的每一帧，为其选择记忆帧和上下文帧。 + + 参数: + w2c_list: 包含所有N个4x4外参矩阵的列表。 + window_size: 初始窗口大小，定义了从哪一帧开始处理。 + memory_frames: 需要选择的记忆帧总数。 + pos_weight: 空间距离的权重。 + ang_weight: 角度距离的权重。 + + 返回: + Dict[int, List[int]]: 一个字典，键是帧索引，值是对应的记忆帧和上下文帧列表。 + """ + if window_size >= len(w2c_list) or window_size < 4: + print("警告: 初始窗口大小必须小于总帧数且至少为4。没有帧需要处理。") + return {} + + pred_latent_size = window_size - memory_frames # 预测帧数 + step = 4 # 每次移动4帧 + + all_selections = {} + + # 如果w2c_list小于等于window_size - memory_frames，则不处理 + if len(w2c_list) <= window_size - memory_frames: + print( + "警告: 总帧数小于等于 window_size - memory_frames，没有帧需要作为memory处理。" + ) + return all_selections + + # 从 window_size 开始遍历到列表的末尾 + for current_frame_idx in range( + memory_frames, len(w2c_list) - pred_latent_size + step, step + ): + # print(f"正在为帧 {current_frame_idx} 选择记忆和上下文...") + + selected_frames = select_memory_frames( + w2c_list, current_frame_idx, memory_frames, pos_weight, ang_weight + ) + all_selections[current_frame_idx] = selected_frames + + return all_selections + + +# *******************************Add the memory part******************************** + + +def get_normalized_dir_diff(x_inv): + # 提取方向向量 + dirs = x_inv[:, :3, 3] # shape: (N, 3) + + # 计算相邻差 + diff = torch.zeros_like(dirs) + diff[1:] = dirs[1:] - dirs[:-1] + + # 对非零向量归一化 + norms = torch.norm(diff, dim=1, keepdim=True) + norms[0] = 1.0 # 避免第一行除零 + diff_norm = diff / norms + + # 第一个保持 (0,0,0) + diff_norm[0] = torch.tensor([0.0, 0.0, 0.0], dtype=diff.dtype) + + return norms, diff_norm + + +def relative_rotations(c2w_mats): + """ + c2w_mats: (N, 4, 4) 相机c2w矩阵 + 返回 (N-1, 3)，每行为相邻帧绕 x,y,z 的相对旋转角度（度数） + """ + rots = c2w_mats[:, :3, :3] # 提取旋转部分 + rel_angles = [] + + for i in range(len(rots) - 1): + R_i = rots[i] + R_j = rots[i + 1] + R_rel = R_i.T @ R_j # 相对旋转矩阵 + + # --- 将 R_rel 转为欧拉角 (XYZ顺序) --- + # 注意：这里假设右手坐标系，欧拉角范围 -180~180° + sy = torch.sqrt(R_rel[0, 0] ** 2 + R_rel[1, 0] ** 2) + + if sy > 1e-6: # 一般情况 + x = torch.atan2(R_rel[2, 1], R_rel[2, 2]) + y = torch.atan2(-R_rel[2, 0], sy) + z = torch.atan2(R_rel[1, 0], R_rel[0, 0]) + else: # 奇异情况（gimbal lock） + x = torch.atan2(-R_rel[1, 2], R_rel[1, 1]) + y = torch.atan2(-R_rel[2, 0], sy) + z = 0.0 + + angles = torch.rad2deg(torch.stack([x, y, z])) + rel_angles.append(angles) + + return torch.stack(rel_angles) diff --git a/worldcompass/fastvideo/models/vision_utils.py b/worldcompass/fastvideo/models/vision_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eff07f07da08c180516fa3b841587f8b99729044 --- /dev/null +++ b/worldcompass/fastvideo/models/vision_utils.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import tempfile +from collections.abc import Callable +from urllib.parse import unquote, urlparse + +import imageio +import numpy as np +import PIL.Image +import PIL.ImageOps +import requests +import torch +from packaging import version + +if version.parse(version.parse(PIL.__version__).base_version) >= version.parse( + "9.1.0" +): + PIL_INTERPOLATION = { + "linear": PIL.Image.Resampling.BILINEAR, + "bilinear": PIL.Image.Resampling.BILINEAR, + "bicubic": PIL.Image.Resampling.BICUBIC, + "lanczos": PIL.Image.Resampling.LANCZOS, + "nearest": PIL.Image.Resampling.NEAREST, + } +else: + PIL_INTERPOLATION = { + "linear": PIL.Image.LINEAR, + "bilinear": PIL.Image.BILINEAR, + "bicubic": PIL.Image.BICUBIC, + "lanczos": PIL.Image.LANCZOS, + "nearest": PIL.Image.NEAREST, + } + + +def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray: + r"""Convert a PIL image or a list of PIL images to NumPy arrays. + + Args: + images (`PIL.Image.Image` or `List[PIL.Image.Image]`): + The PIL image or list of images to convert to NumPy format. + + Returns: + `np.ndarray`: + A NumPy array representation of the images. + """ + if not isinstance(images, list): + images = [images] + images = [np.array(image).astype(np.float32) / 255.0 for image in images] + images_arr: np.ndarray = np.stack(images, axis=0) + + return images_arr + + +def numpy_to_pt(images: np.ndarray) -> torch.Tensor: + r"""Convert a NumPy image to a PyTorch tensor. + + Args: + images (`np.ndarray`): + The NumPy image array to convert to PyTorch format. + + Returns: + `torch.Tensor`: + A PyTorch tensor representation of the images. + """ + if images.ndim == 3: + images = images[..., None] + + images = torch.from_numpy(images.transpose(0, 3, 1, 2)) + return images + + +def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + r"""Normalize an image array to [-1,1]. + + Args: + images (`np.ndarray` or `torch.Tensor`): + The image array to normalize. + + Returns: + `np.ndarray` or `torch.Tensor`: + The normalized image array. + """ + return 2.0 * images - 1.0 + + +# adapted from diffusers.utils import load_image +def load_image( + image: str | PIL.Image.Image, + convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None, +) -> PIL.Image.Image: + """Loads `image` to a PIL Image. + + Args: + image (`str` or `PIL.Image.Image`): + The image to convert to the PIL Image format. + convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*): + A conversion method to apply to the image after loading it. When set to `None` the image will be converted + "RGB". + + Returns: + `PIL.Image.Image`: + A PIL Image. + """ + if isinstance(image, str): + if image.startswith("http://") or image.startswith("https://"): + image = PIL.Image.open(requests.get(image, stream=True).raw) + elif os.path.isfile(image): + image = PIL.Image.open(image) + else: + raise ValueError( + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." + ) + elif isinstance(image, PIL.Image.Image): + image = image + else: + raise ValueError( + "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." + ) + + image = PIL.ImageOps.exif_transpose(image) + + if convert_method is not None: + image = convert_method(image) + else: + image = image.convert("RGB") + + return image + + +# adapted from diffusers.utils import load_video +def load_video( + video: str, + convert_method: ( + Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None + ) = None, +) -> list[PIL.Image.Image]: + """Loads `video` to a list of PIL Image. + + Args: + video (`str`): + A URL or Path to a video to convert to a list of PIL Image format. + convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*): + A conversion method to apply to the video after loading it. When set to `None` the images will be converted + to "RGB". + Returns: + `List[PIL.Image.Image]`: + The video as a list of PIL images. + """ + is_url = video.startswith("http://") or video.startswith("https://") + is_file = os.path.isfile(video) + was_tempfile_created = False + + if not (is_url or is_file): + raise ValueError( + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." + ) + + if is_url: + response = requests.get(video, stream=True) + if response.status_code != 200: + raise ValueError( + f"Failed to download video. Status code: {response.status_code}" + ) + + parsed_url = urlparse(video) + file_name = os.path.basename(unquote(parsed_url.path)) + + suffix = os.path.splitext(file_name)[1] or ".mp4" + with tempfile.NamedTemporaryFile( + suffix=suffix, delete=False + ) as temp_file: + video_path = temp_file.name + video_data = response.iter_content(chunk_size=8192) + for chunk in video_data: + temp_file.write(chunk) + + video = video_path + + pil_images = [] + if video.endswith(".gif"): + gif = PIL.Image.open(video) + try: + while True: + pil_images.append(gif.copy()) + gif.seek(gif.tell() + 1) + except EOFError: + pass + + else: + try: + imageio.plugins.ffmpeg.get_exe() + except AttributeError: + raise AttributeError( + "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" + ) from None + + with imageio.get_reader(video) as reader: + # Read all frames + for frame in reader: + pil_images.append(PIL.Image.fromarray(frame)) + + if was_tempfile_created: + os.remove(video_path) + + if convert_method is not None: + pil_images = convert_method(pil_images) + + return pil_images + + +def get_default_height_width( + image: PIL.Image.Image | np.ndarray | torch.Tensor, + vae_scale_factor: int, + height: int | None = None, + width: int | None = None, +) -> tuple[int, int]: + r"""Returns the height and width of the image, downscaled to the next integer multiple of + `vae_scale_factor`. + + Args: + image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it + should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch + tensor, it should have shape `[batch, channels, height, width]`. + height (`Optional[int]`, *optional*, defaults to `None`): + The height of the preprocessed image. If `None`, the height of the `image` input will be used. + width (`Optional[int]`, *optional*, defaults to `None`): + The width of the preprocessed image. If `None`, the width of the `image` input will be used. + + Returns: + `Tuple[int, int]`: + A tuple containing the height and width, both resized to the nearest integer multiple of + `vae_scale_factor`. + """ + + if height is None: + if isinstance(image, PIL.Image.Image): + height = image.height + elif isinstance(image, torch.Tensor): + height = image.shape[2] + else: + height = image.shape[1] + + if width is None: + if isinstance(image, PIL.Image.Image): + width = image.width + elif isinstance(image, torch.Tensor): + width = image.shape[3] + else: + width = image.shape[2] + + width, height = ( + x - x % vae_scale_factor for x in (width, height) + ) # resize to integer multiple of vae_scale_factor + + return height, width + + +def resize( + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int, + width: int, + resize_mode: str = "default", # "default", "fill", "crop" + resample: str = "lanczos", +) -> PIL.Image.Image | np.ndarray | torch.Tensor: + """Resize image. + + Args: + image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`): + The image input, can be a PIL image, numpy array or pytorch tensor. + height (`int`): + The height to resize to. + width (`int`): + The width to resize to. + resize_mode (`str`, *optional*, defaults to `default`): + The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit + within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, + will resize the image to fit within the specified width and height, maintaining the aspect ratio, and + then center the image within the dimensions, filling empty with data from image. If `crop`, will resize + the image to fit within the specified width and height, maintaining the aspect ratio, and then center + the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only + supported for PIL image input. + + Returns: + `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: + The resized image. + """ + if resize_mode != "default" and not isinstance(image, PIL.Image.Image): + raise ValueError( + f"Only PIL image input is supported for resize_mode {resize_mode}" + ) + assert isinstance(image, PIL.Image.Image) + if resize_mode == "default": + image = image.resize( + (width, height), resample=PIL_INTERPOLATION[resample] + ) + else: + raise ValueError(f"resize_mode {resize_mode} is not supported") + return image diff --git a/worldcompass/fastvideo/pipelines/README.md b/worldcompass/fastvideo/pipelines/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3c17e12378b57fb2efd050a96e471fcbbec59340 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/README.md @@ -0,0 +1,4 @@ +# Adding a New Custom Pipeline + +Please see documentation +[here](https://hao-ai-lab.github.io/FastVideo/contributing/add_pipeline.html) diff --git a/worldcompass/fastvideo/pipelines/__init__.py b/worldcompass/fastvideo/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bf63cc6b32e6aff380e0a1bc513eb47095372b54 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/__init__.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Diffusion pipelines for fastvideo. + +This package contains diffusion pipelines for generating videos and images. +""" + +from typing import cast + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.lora_pipeline import LoRAPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.pipelines.pipeline_registry import ( + PipelineType, + get_pipeline_registry, +) +from fastvideo.utils import ( + maybe_download_model, + verify_model_config_and_directory, +) + +logger = init_logger(__name__) + + +class PipelineWithLoRA(LoRAPipeline, ComposedPipelineBase): + """Type for a pipeline that has both ComposedPipelineBase and LoRAPipeline functionality.""" + + pass + + +def build_pipeline( + fastvideo_args: FastVideoArgs, + pipeline_type: PipelineType | str = PipelineType.BASIC, +) -> PipelineWithLoRA: + """Only works with valid hf diffusers configs. + + (model_index.json) + We want to build a pipeline based on the inference args mode_path: + 1. download the model from the hub if it's not already downloaded + 2. verify the model config and directory + 3. based on the config, determine the pipeline class + """ + # Get pipeline type + model_path = fastvideo_args.model_path + model_path = maybe_download_model(model_path) + # fastvideo_args.downloaded_model_path = model_path + logger.info("Model path: %s", model_path) + + config = verify_model_config_and_directory(model_path) + pipeline_name = config.get("_class_name") + if pipeline_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported." + ) + + # Get the appropriate pipeline registry based on pipeline_type + logger.info( + "Building pipeline of type: %s", + ( + pipeline_type.value + if isinstance(pipeline_type, PipelineType) + else pipeline_type + ), + ) + pipeline_registry = get_pipeline_registry(pipeline_type) + + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_cls = pipeline_registry.resolve_pipeline_cls( + pipeline_name, pipeline_type, fastvideo_args.workload_type + ) + + # instantiate the pipelines + pipeline = pipeline_cls(model_path, fastvideo_args) + + logger.info("Pipelines instantiated") + + return cast(PipelineWithLoRA, pipeline) + + +__all__ = [ + "build_pipeline", + "ComposedPipelineBase", + "ForwardBatch", + "LoRAPipeline", + "TrainingBatch", +] diff --git a/worldcompass/fastvideo/pipelines/composed_pipeline_base.py b/worldcompass/fastvideo/pipelines/composed_pipeline_base.py new file mode 100644 index 0000000000000000000000000000000000000000..933307c1da3b61edc2655de46dbe6dd92f3c24d8 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/composed_pipeline_base.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Base class for composed pipelines. + +This module defines the base class for pipelines that are composed of multiple stages. +""" + +import argparse +import os +from abc import ABC, abstractmethod +from typing import Any, cast + +import torch + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.distributed import ( + maybe_init_distributed_environment_and_model_parallel, + initialize_gpu_parallel_group, +) +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import PipelineComponentLoader +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages import PipelineStage +from fastvideo.utils import ( + maybe_download_model, + verify_model_config_and_directory, +) + +logger = init_logger(__name__) + + +class ComposedPipelineBase(ABC): + """Base class for pipelines composed of multiple stages. + + This class provides the framework for creating pipelines by composing multiple stages together. + Each stage is responsible for a specific part of the diffusion process, and the pipeline + orchestrates the execution of these stages. + """ + + is_video_pipeline: bool = False # To be overridden by video pipelines + _required_config_modules: list[str] = [] + _extra_config_module_map: dict[str, str] = {} + training_args: TrainingArgs | None = None + fastvideo_args: FastVideoArgs | TrainingArgs | None = None + modules: dict[str, torch.nn.Module] = {} + post_init_called: bool = False + + # TODO(will): args should support both inference args and training args + def __init__( + self, + model_path: str, + fastvideo_args: FastVideoArgs | TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ): + """Initialize the pipeline. + + After __init__, the pipeline should be ready to use. The pipeline should be stateless and + not hold any batch state. + """ + self.fastvideo_args = fastvideo_args + + self.model_path: str = model_path + self._stages: list[PipelineStage] = [] + self._stage_name_mapping: dict[str, PipelineStage] = {} + + if required_config_modules is not None: + self._required_config_modules = required_config_modules + + if self._required_config_modules is None: + raise NotImplementedError( + "Subclass must set _required_config_modules" + ) + + maybe_init_distributed_environment_and_model_parallel( + fastvideo_args.tp_size, fastvideo_args.sp_size + ) + initialize_gpu_parallel_group(fastvideo_args.gpu_para) + + # Load modules directly in initialization + logger.info("Loading pipeline modules...") + # print(fastvideo_args, loaded_modules) + # quit() + + if "Hunyuan" in fastvideo_args.cls_name: + self.modules = self.load_hunyuan_modules(fastvideo_args) + else: + self.modules = self.load_modules(fastvideo_args, loaded_modules) + + def set_trainable(self) -> None: + # Only train DiT + if getattr(self.fastvideo_args, "training_mode", False): + for name, module in self.modules.items(): + if not isinstance(module, torch.nn.Module): + continue + if name == "transformer": + module.requires_grad_(True) + else: + module.requires_grad_(False) + + def post_init(self) -> None: + assert self.fastvideo_args is not None, "fastvideo_args must be set" + if self.post_init_called: + return + self.post_init_called = True + if self.fastvideo_args.training_mode: + assert isinstance(self.fastvideo_args, TrainingArgs) + self.training_args = self.fastvideo_args + assert self.training_args is not None + self.initialize_training_pipeline(self.training_args) + if self.training_args.log_validation: + self.initialize_validation_pipeline(self.training_args) + + self.initialize_pipeline(self.fastvideo_args) + if self.fastvideo_args.enable_torch_compile: + self.modules["transformer"] = torch.compile( + self.modules["transformer"] + ) + logger.info("Torch Compile enabled for DiT") + + if not self.fastvideo_args.training_mode: + logger.info("Creating pipeline stages...") + self.create_pipeline_stages(self.fastvideo_args) + + def initialize_training_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method" + ) + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError( + "if log_validation is True, the pipeline must implement this method" + ) + + @classmethod + def from_pretrained( + cls, + model_path: str, + device: str | None = None, + torch_dtype: torch.dtype | None = None, + pipeline_config: str | PipelineConfig | None = None, + args: argparse.Namespace | None = None, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + **kwargs, + ) -> "ComposedPipelineBase": + """Load a pipeline from a pretrained model. + + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + assert args is not None, "args must be provided for training mode" + fastvideo_args = TrainingArgs.from_cli_args(args) + # TODO(will): fix this so that its not so ugly + fastvideo_args.model_path = model_path + for key, value in kwargs.items(): + setattr(fastvideo_args, key, value) + + fastvideo_args.dit_cpu_offload = False + # we hijack the precision to be the master weight type so that the + # model is loaded with the correct precision. Subsequently we will + # use FSDP2's MixedPrecisionPolicy to set the precision for the + # fwd, bwd, and other operations' precision. + # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training' + + logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args) + + pipe = cls( + model_path, + fastvideo_args, + required_config_modules=required_config_modules, + loaded_modules=loaded_modules, + ) + pipe.post_init() + return pipe + + def get_module(self, module_name: str, default_value: Any = None) -> Any: + if module_name not in self.modules: + return default_value + return self.modules[module_name] + + def add_module(self, module_name: str, module: Any): + self.modules[module_name] = module + + def _load_config(self, model_path: str) -> dict[str, Any]: + model_path = maybe_download_model(self.model_path) + self.model_path = model_path + # fastvideo_args.downloaded_model_path = model_path + logger.info("Model path: %s", model_path) + config = verify_model_config_and_directory(model_path) + return cast(dict[str, Any], config) + + @property + def required_config_modules(self) -> list[str]: + """List of modules that are required by the pipeline. The names should match the diffusers + directory and model_index.json file. These modules will be loaded using the + PipelineComponentLoader and made available in the modules dictionary. Access these modules + using the get_module method. + + class ConcretePipeline(ComposedPipelineBase): + _required_config_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"] + + + @property + def required_config_modules(self): + return self._required_config_modules + """ + return self._required_config_modules + + @property + def stages(self) -> list[PipelineStage]: + """List of stages in the pipeline.""" + return self._stages + + @abstractmethod + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Create the inference pipeline stages.""" + raise NotImplementedError + + def create_training_stages(self, training_args: TrainingArgs): + """Create the training pipeline stages.""" + raise NotImplementedError + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + """Initialize the pipeline.""" + return + + def load_hunyuan_modules( + self, + fastvideo_args: FastVideoArgs, + ) -> dict[str, Any]: + """Load the modules from the config. + + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + modules = {} + module_name = "transformer" + component_model_path = fastvideo_args.load_from_dir + fastvideo_args.module_name = module_name + + module = PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers="diffusers", + fastvideo_args=fastvideo_args, + ) + logger.info( + "Loaded module %s from %s", module_name, component_model_path + ) + + if module_name in modules: + logger.warning("Overwriting module %s", module_name) + modules[module_name] = module + # modules['vae'] = vae_module + return modules + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> dict[str, Any]: + """Load the modules from the config. + + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + # 这里是为了得到需要加载的模型的path + model_index = self._load_config(self.model_path) + logger.info("Loading pipeline modules from config: %s", model_index) + + # remove keys that are not pipeline modules + model_index.pop("_class_name") + model_index.pop("_diffusers_version") + # @TODO(Wei): Temporary hack + # wan2.2 14B才会用上 + if ( + "boundary_ratio" in model_index + and model_index["boundary_ratio"] is not None + ): + logger.info( + "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." + ) + self.required_config_modules.append("transformer_2") + if fastvideo_args.boundary_ratio is None: + logger.info( + "MoE pipeline detected. Setting boundary ratio to %s", + model_index["boundary_ratio"], + ) + fastvideo_args.boundary_ratio = model_index["boundary_ratio"] + + model_index.pop("boundary_ratio", None) + model_index.pop("expand_timesteps", None) + + # some sanity checks + assert ( + len(model_index) > 1 + ), "model_index.json must contain at least one pipeline module" + # print(model_index) + # quit() + for module_name in self.required_config_modules: + if ( + module_name not in model_index + and module_name in self._extra_config_module_map + ): + extra_module_value = self._extra_config_module_map[module_name] + logger.warning( + "model_index.json does not contain a %s module, but found {%s: %s}.", + module_name, + module_name, + extra_module_value, + ) + if extra_module_value in model_index: + logger.info( + "Using module %s for %s", + extra_module_value, + module_name, + ) + model_index[module_name] = model_index[extra_module_value] + continue + else: + raise ValueError( + f'''Required module key: {module_name} value: {model_index.get(module_name)} + was not found in loaded modules {model_index.keys()}''' + ) + + # all the component models used by the pipeline + required_modules = self.required_config_modules + logger.info("Loading required modules: %s", required_modules) + + # 在这里是为了加载需要的模型，比如说dit + modules = {} + # print(model_index, self.required_config_modules) + for module_name, ( + transformers_or_diffusers, + architecture, + ) in model_index.items(): + # print(transformers_or_diffusers) + if transformers_or_diffusers is None: + if module_name in self.required_config_modules: + self.required_config_modules.remove(module_name) + continue + if module_name not in required_modules: + logger.info("Skipping module %s", module_name) + continue + if loaded_modules is not None and module_name in loaded_modules: + logger.info("Using module %s already provided", module_name) + modules[module_name] = loaded_modules[module_name] + continue + + # we load the module from the extra config module map if it exists + if module_name in self._extra_config_module_map: + load_module_name = self._extra_config_module_map[module_name] + else: + load_module_name = module_name + + component_model_path = os.path.join( + self.model_path, load_module_name + ) + # 为了区分不同的transformer DiT + fastvideo_args.module_name = module_name + + module = PipelineComponentLoader.load_module( + module_name=load_module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + fastvideo_args=fastvideo_args, + ) + logger.info( + "Loaded module %s from %s", module_name, component_model_path + ) + + if module_name in modules: + logger.warning("Overwriting module %s", module_name) + modules[module_name] = module + + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in modules or modules[module_name] is None: + raise ValueError( + f'''Required module key: {module_name} value: {modules.get(module_name)} + was not found in loaded modules {modules.keys()}''' + ) + + return modules + + def add_stage(self, stage_name: str, stage: PipelineStage): + assert self.modules is not None, "No modules are registered" + self._stages.append(stage) + self._stage_name_mapping[stage_name] = stage + setattr(self, stage_name, stage) + + # TODO(will): don't hardcode no_grad + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Generate a video or image using the pipeline. + + Args: + batch: The batch to generate from. + fastvideo_args: The inference arguments. + Returns: + ForwardBatch: The batch with the generated video or image. + """ + if not self.post_init_called: + self.post_init() + + # Execute each stage + logger.info( + "Running pipeline stages: %s", self._stage_name_mapping.keys() + ) + # logger.info("Batch: %s", batch) + for stage in self.stages: + batch = stage(batch, fastvideo_args) + + # Return the output + return batch + + def train(self) -> None: + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method" + ) diff --git a/worldcompass/fastvideo/pipelines/lora_pipeline.py b/worldcompass/fastvideo/pipelines/lora_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..2da732982cd2a5ff0f4451c86ebf17897e915261 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/lora_pipeline.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections import defaultdict +from collections.abc import Hashable +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors.torch import load_file +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.layers.lora.linear import ( + BaseLayerWithLoRA, + get_lora_layer, + replace_submodule, +) +from fastvideo.logger import init_logger +from fastvideo.models.loader.utils import get_param_names_mapping +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.utils import maybe_download_lora + +logger = init_logger(__name__) + + +class LoRAPipeline(ComposedPipelineBase): + """Pipeline that supports injecting LoRA adapters into the diffusion transformer. + + TODO: support training. + """ + + lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict( + dict + ) # state dicts of loaded lora adapters + cur_adapter_name: str = "" + cur_adapter_path: str = "" + lora_layers: dict[str, BaseLayerWithLoRA] = {} + fastvideo_args: FastVideoArgs | TrainingArgs + exclude_lora_layers: list[str] = [] + device: torch.device = get_local_torch_device() + lora_target_modules: list[str] | None = None + lora_path: str | None = None + lora_nickname: str = "default" + lora_rank: int | None = None + lora_alpha: int | None = None + lora_initialized: bool = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.device = get_local_torch_device() + # self.exclude_lora_layers = self.modules[ + # "transformer"].config.arch_config.exclude_lora_layers + self.lora_target_modules = self.fastvideo_args.lora_target_modules + self.lora_path = self.fastvideo_args.lora_path + self.lora_nickname = self.fastvideo_args.lora_nickname + self.training_mode = self.fastvideo_args.training_mode + if self.training_mode and getattr( + self.fastvideo_args, "lora_training", False + ): + assert isinstance(self.fastvideo_args, TrainingArgs) + if self.fastvideo_args.lora_alpha is None: + self.fastvideo_args.lora_alpha = self.fastvideo_args.lora_rank + self.lora_rank = self.fastvideo_args.lora_rank # type: ignore + self.lora_alpha = self.fastvideo_args.lora_alpha # type: ignore + logger.info( + "Using LoRA training with rank %d and alpha %d", + self.lora_rank, + self.lora_alpha, + ) + if self.lora_target_modules is None: + self.lora_target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "to_q", + "to_k", + "to_v", + "to_out", + "to_qkv", + ] + self.convert_to_lora_layers() + # Inference + elif not self.training_mode and self.lora_path is not None: + self.convert_to_lora_layers() + self.set_lora_adapter( + self.lora_nickname, # type: ignore + self.lora_path, + ) # type: ignore + + def is_target_layer(self, module_name: str) -> bool: + if self.lora_target_modules is None: + return True + return any( + target_name in module_name + for target_name in self.lora_target_modules + ) + + def set_trainable(self) -> None: + + is_lora_training = self.training_mode and getattr( + self.fastvideo_args, "lora_training", False + ) + if not is_lora_training: + super().set_trainable() + return + + self.modules["transformer"].requires_grad_(False) + device_mesh = init_device_mesh( + "cuda", + (dist.get_world_size(), 1), + mesh_dim_names=["fake", "replicate"], + ) + for name, layer in self.lora_layers.items(): + # Enable grads for lora weights only + # Must convert to DTensor for compatibility with other FSDP modules in grad calculation + layer.lora_A.requires_grad_(True) + layer.lora_B.requires_grad_(True) + layer.base_layer.requires_grad_(False) + layer.lora_A = nn.Parameter( + DTensor.from_local(layer.lora_A, device_mesh=device_mesh) + ) + layer.lora_B = nn.Parameter( + DTensor.from_local(layer.lora_B, device_mesh=device_mesh) + ) + + def convert_to_lora_layers(self) -> None: + """Unified method to convert the transformer to a LoRA transformer.""" + if self.lora_initialized: + return + self.lora_initialized = True + converted_count = 0 + for name, layer in self.modules["transformer"].named_modules(): + if not self.is_target_layer(name): + continue + + excluded = False + for exclude_layer in self.exclude_lora_layers: + if exclude_layer in name: + excluded = True + break + if excluded: + continue + + layer = get_lora_layer( + layer, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + training_mode=self.training_mode, + ) + if layer is not None: + self.lora_layers[name] = layer + replace_submodule(self.modules["transformer"], name, layer) + converted_count += 1 + logger.info("Converted %d layers to LoRA layers", converted_count) + + def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None): # type: ignore + """Load a LoRA adapter into the pipeline and merge it into the transformer. + + Args: + lora_nickname: The "nick name" of the adapter when referenced in the pipeline. + lora_path: The path to the adapter, either a local path or a Hugging Face repo id. + """ + + if lora_nickname not in self.lora_adapters and lora_path is None: + raise ValueError( + f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it." + ) + if not self.lora_initialized: + self.convert_to_lora_layers() + adapter_updated = False + rank = dist.get_rank() + if lora_path is not None and lora_path != self.cur_adapter_path: + lora_local_path = maybe_download_lora(lora_path) + lora_state_dict = load_file(lora_local_path) + + # Map the hf layer names to our custom layer names + param_names_mapping_fn = get_param_names_mapping( + self.modules["transformer"].param_names_mapping + ) + lora_param_names_mapping_fn = get_param_names_mapping( + self.modules["transformer"].lora_param_names_mapping + ) + + to_merge_params: defaultdict[Hashable, dict[Any, Any]] = ( + defaultdict(dict) + ) + for name, weight in lora_state_dict.items(): + name = name.replace("diffusion_model.", "") + name = name.replace(".weight", "") + name, _, _ = lora_param_names_mapping_fn(name) + target_name, merge_index, num_params_to_merge = ( + param_names_mapping_fn(name) + ) + # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse + # see param mapping in HunyuanVideoArchConfig + if merge_index is not None and "lora_B" in name: + to_merge_params[target_name][merge_index] = weight + if len(to_merge_params[target_name]) == num_params_to_merge: + # cat at output dim according to the merge_index order + sorted_tensors = [ + to_merge_params[target_name][i] + for i in range(num_params_to_merge) + ] + weight = torch.cat(sorted_tensors, dim=1) + del to_merge_params[target_name] + else: + continue + + if target_name in self.lora_adapters[lora_nickname]: + raise ValueError( + f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]" + ) + self.lora_adapters[lora_nickname][target_name] = weight.to( + self.device + ) + adapter_updated = True + self.cur_adapter_path = lora_path + logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) + + if not adapter_updated and self.cur_adapter_name == lora_nickname: + return + self.cur_adapter_name = lora_nickname + + # Merge the new adapter + adapted_count = 0 + for name, layer in self.lora_layers.items(): + lora_A_name = name + ".lora_A" + lora_B_name = name + ".lora_B" + if ( + lora_A_name in self.lora_adapters[lora_nickname] + and lora_B_name in self.lora_adapters[lora_nickname] + ): + layer.set_lora_weights( + self.lora_adapters[lora_nickname][lora_A_name], + self.lora_adapters[lora_nickname][lora_B_name], + training_mode=self.fastvideo_args.training_mode, + lora_path=lora_path, + ) + adapted_count += 1 + else: + if rank == 0: + logger.warning( + "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.", + lora_path, + name, + ) + layer.disable_lora = True + logger.info( + "Rank %d: LoRA adapter %s applied to %d layers", + rank, + lora_path, + adapted_count, + ) diff --git a/worldcompass/fastvideo/pipelines/pipeline_batch_info.py b/worldcompass/fastvideo/pipelines/pipeline_batch_info.py new file mode 100644 index 0000000000000000000000000000000000000000..2f32215a84d5464760341373e69bfa46dfc0dc3b --- /dev/null +++ b/worldcompass/fastvideo/pipelines/pipeline_batch_info.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py +"""Data structures for functional pipeline processing. + +This module defines the dataclasses used to pass state between pipeline components in a functional +manner, reducing the need for explicit parameter passing. +""" + +import pprint +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Any + +import PIL.Image +import torch + +if TYPE_CHECKING: + from torchcodec.decoders import VideoDecoder + +import time +from collections import OrderedDict + +from fastvideo.attention import AttentionMetadata +from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams + + +class PipelineLoggingInfo: + """Simple approach using OrderedDict to track stage metrics.""" + + def __init__(self): + # OrderedDict preserves insertion order and allows easy access + self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict() + + def add_stage_execution_time(self, stage_name: str, execution_time: float): + """Add execution time for a stage.""" + if stage_name not in self.stages: + self.stages[stage_name] = {} + self.stages[stage_name]["execution_time"] = execution_time + self.stages[stage_name]["timestamp"] = time.time() + + def add_stage_metric(self, stage_name: str, metric_name: str, value: Any): + """Add any metric for a stage.""" + if stage_name not in self.stages: + self.stages[stage_name] = {} + self.stages[stage_name][metric_name] = value + + def get_stage_info(self, stage_name: str) -> dict[str, Any]: + """Get all info for a specific stage.""" + return self.stages.get(stage_name, {}) + + def get_execution_order(self) -> list[str]: + """Get stages in execution order.""" + return list(self.stages.keys()) + + def get_total_execution_time(self) -> float: + """Get total pipeline execution time.""" + return sum( + stage.get("execution_time", 0) for stage in self.stages.values() + ) + + +@dataclass +class ForwardBatch: + """Complete state passed through the pipeline execution. + + This dataclass contains all information needed during the diffusion pipeline execution, allowing + methods to update specific components without needing to manage numerous individual parameters. + """ + + # TODO(will): double check that args are separate from fastvideo_args + # properly. Also maybe think about providing an abstraction for pipeline + # specific arguments. + data_type: str + + generator: torch.Generator | list[torch.Generator] | None = None + + # Image inputs + image_path: str | None = None + image_embeds: list[torch.Tensor] = field(default_factory=list) + pil_image: torch.Tensor | PIL.Image.Image | None = None + preprocessed_image: torch.Tensor | None = None + + # Text inputs + prompt: str | list[str] | None = None + negative_prompt: str | list[str] | None = None + prompt_path: str | None = None + output_path: str = "outputs/" + output_video_name: str | None = None + # Primary encoder embeddings + prompt_embeds: list[torch.Tensor] = field(default_factory=list) + negative_prompt_embeds: list[torch.Tensor] | None = None + prompt_attention_mask: list[torch.Tensor] | None = None + negative_attention_mask: list[torch.Tensor] | None = None + clip_embedding_pos: list[torch.Tensor] | None = None + clip_embedding_neg: list[torch.Tensor] | None = None + + # Additional text-related parameters + max_sequence_length: int | None = None + prompt_template: dict[str, Any] | None = None + do_classifier_free_guidance: bool = False + + # Batch info + batch_size: int | None = None + num_videos_per_prompt: int = 1 + seed: int | None = None + seeds: list[int] | None = None + + # Tracking if embeddings are already processed + is_prompt_processed: bool = False + + # Latent tensors + latents: torch.Tensor | None = None + raw_latent_shape: torch.Tensor | None = None + noise_pred: torch.Tensor | None = None + image_latent: torch.Tensor | None = None + + # Latent dimensions + height_latents: list[int] | int | None = None + width_latents: list[int] | int | None = None + num_frames: list[int] | int = 1 # Default for image models + num_frames_round_down: bool = ( + False # Whether to round down num_frames if it's not divisible by num_gpus + ) + + # Original dimensions (before VAE scaling) + height: list[int] | int | None = None + width: list[int] | int | None = None + fps: list[int] | int | None = None + + # Timesteps + timesteps: torch.Tensor | None = None + timestep: torch.Tensor | float | int | None = None + step_index: int | None = None + + # Scheduler parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_scale_2: float | None = None + guidance_rescale: float = 0.0 + eta: float = 0.0 + sigmas: list[float] | None = None + + n_tokens: int | None = None + + # Other parameters that may be needed by specific schedulers + extra_step_kwargs: dict[str, Any] = field(default_factory=dict) + + # Component modules (populated by the pipeline) + modules: dict[str, Any] = field(default_factory=dict) + + # Final output (after pipeline completion) + output: Any = None + + # Extra parameters that might be needed by specific pipeline implementations + extra: dict[str, Any] = field(default_factory=dict) + + # Misc + save_video: bool = True + return_frames: bool = False + + # TeaCache parameters + enable_teacache: bool = False + teacache_params: TeaCacheParams | WanTeaCacheParams | None = None + + # STA parameters + STA_param: list | None = None + is_cfg_negative: bool = False + mask_search_final_result_pos: list[list] | None = None + mask_search_final_result_neg: list[list] | None = None + + # VSA parameters + VSA_sparsity: float = 0.0 + + # Logging info + logging_info: PipelineLoggingInfo = field( + default_factory=PipelineLoggingInfo + ) + + def __post_init__(self): + """Initialize dependent fields after dataclass initialization.""" + + # Set do_classifier_free_guidance based on guidance scale and negative prompt + if self.guidance_scale > 1.0: + self.do_classifier_free_guidance = True + if self.negative_prompt_embeds is None: + self.negative_prompt_embeds = [] + if self.guidance_scale_2 is None: + self.guidance_scale_2 = self.guidance_scale + + def __str__(self): + return pprint.pformat(asdict(self), indent=2, width=120) + + +@dataclass +class TrainingBatch: + current_timestep: int = 0 + current_vsa_sparsity: float = 0.0 + + # Dataloader batch outputs + latents: torch.Tensor | None = None + prompt_embed: torch.Tensor | None = None + prompt: str | None = None + negative_prompt_embed: torch.Tensor | None = None + w2c: torch.Tensor | None = None + intrinsic: torch.Tensor | None = None + frame_indices: torch.Tensor | None = None + action: torch.Tensor | None = None + norm_dirs: torch.Tensor | None = None + action_rot_trans: torch.Tensor | None = None + i2v_mask: torch.Tensor | None = None + window_frames: int = 0 + per_seq_length: int = 0 + current_start: int = 0 + current_end: int = 0 + + raw_latent_shape: torch.Tensor | None = None + noise_latents: torch.Tensor | None = None + encoder_hidden_states: torch.Tensor | None = None + encoder_attention_mask: torch.Tensor | None = None + # i2v + preprocessed_image: torch.Tensor | None = None + image_embeds: torch.Tensor | None = None + image_latents: torch.Tensor | None = None + infos: list[dict[str, Any]] | None = None + mask_lat_size: torch.Tensor | None = None + + # GRPO inputs + all_latents: dict[str, torch.Tensor] | None = None + all_log_probs: dict[str, torch.Tensor] | None = None + all_action_rewards: dict[str, torch.Tensor] | None = None + all_geometry_rewards: dict[str, torch.Tensor] | None = None + + action_reward_sample_mean: dict[str, torch.Tensor] | None = None + hpsv3_reward_sample_mean: dict[str, torch.Tensor] | None = None + + action_reward_chunk_mean: dict[str, torch.Tensor] | None = None + hpsv3_reward_chunk_mean: dict[str, torch.Tensor] | None = None + + action_reward_mean: dict[str, torch.Tensor] | None = None + fine_action_reward_mean: dict[str, torch.Tensor] | None = None + geometry_reward_mean: dict[str, torch.Tensor] | None = None + hpsv3_reward_mean: dict[str, torch.Tensor] | None = None + camera_reward_mean: dict[str, torch.Tensor] | None = None + + all_timesteps: dict[str, torch.Tensor] | None = None + all_kv_caches: dict[str, torch.Tensor] | None = None + grpo_kwargs: dict[str, Any] | None = None + sigma_schedule: torch.Tensor | None = None + kv_cache: list[dict[str, torch.Tensor]] | None = None + training_steps: list[int] | None = None + training_chunks: list[str] | None = None + selected_chunk_id: int | None = None + + used_for_training: bool = True + + # Transformer inputs + noisy_model_input: torch.Tensor | None = None + timesteps: torch.Tensor | None = None + sigmas: torch.Tensor | None = None + noise: torch.Tensor | None = None + + attn_metadata_vsa: AttentionMetadata | None = None + attn_metadata: AttentionMetadata | None = None + + # input kwargs + input_kwargs: dict[str, Any] | None = None + sample_kwargs: dict[str, Any] | None = None + + # Training loss + loss: torch.Tensor | None = None + + # Training outputs + total_loss: float | None = None + grad_norm: float | None = None + samples_loss: dict[str, torch.Tensor] | None = None + samples_grad_norm: dict[str, torch.Tensor] | None = None + + # Distillation-specific attributes + encoder_hidden_states_neg: torch.Tensor | None = None + encoder_attention_mask_neg: torch.Tensor | None = None + conditional_dict: dict[str, Any] | None = None + unconditional_dict: dict[str, Any] | None = None + + # Distillation losses + generator_loss: float = 0.0 + fake_score_loss: float = 0.0 + + dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PreprocessBatch(ForwardBatch): + video_loader: list["VideoDecoder"] = field(default_factory=list) + video_file_name: list[str] = field(default_factory=list) diff --git a/worldcompass/fastvideo/pipelines/pipeline_registry.py b/worldcompass/fastvideo/pipelines/pipeline_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..24fa192a82969e2a5e8f9bae55cd72854a65bc1f --- /dev/null +++ b/worldcompass/fastvideo/pipelines/pipeline_registry.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py +# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py + +import importlib +import pkgutil +from collections.abc import Set +from dataclasses import dataclass, field +from enum import Enum +from functools import lru_cache + +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +logger = init_logger(__name__) + +# map pipeline name to folder name +_PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { + "WanPipeline": "wan", + "WanDMDPipeline": "wan", + "WanImageToVideoPipeline": "wan", + "StepVideoPipeline": "stepvideo", + "HunyuanVideoPipeline": "hunyuan", +} + +_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { + WorkloadType.I2V: "PreprocessPipelineI2V", + WorkloadType.T2V: "PreprocessPipelineT2V", +} + + +class PipelineType(str, Enum): + """Enumeration for different pipeline types. + + Inherits from str to allow string comparison for backward compatibility. + """ + + BASIC = "basic" + PREPROCESS = "preprocess" + TRAINING = "training" + + @classmethod + def from_string(cls, value: str) -> "PipelineType": + """Convert string to PipelineType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid pipeline type: {value}. Must be one of: {', '.join([t.value for t in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings.""" + return [pipeline_type.value for pipeline_type in cls] + + +@dataclass +class _PipelineRegistry: + # Keyed by pipeline_type -> architecture -> pipeline_name + # pipelines[pipeline_type][architecture][pipeline_name] = pipeline_cls + pipelines: dict[ + str, dict[str, dict[str, type[ComposedPipelineBase] | None]] + ] = field(default_factory=dict) + + def get_supported_archs( + self, pipeline_name_in_config: str, pipeline_type: PipelineType + ) -> Set[str]: + """Get supported architectures, optionally filtered by pipeline type and workload type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + return set(self.pipelines[pipeline_type.value][arch].keys()) + + def _load_preprocess_pipeline_cls( + self, workload_type: WorkloadType, arch: str + ) -> type[ComposedPipelineBase] | None: + pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[ + workload_type + ] + + return self.pipelines[PipelineType.PREPROCESS.value][arch][ + pipeline_name + ] + + def _try_load_pipeline_cls( + self, + pipeline_name_in_config: str, + pipeline_type: PipelineType, + workload_type: WorkloadType, + ) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None: + """Try to load a pipeline class for the given architecture, pipeline type, and workload + type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + + if ( + pipeline_type.value not in self.pipelines + or arch not in self.pipelines[pipeline_type.value] + ): + return None + + if pipeline_type == PipelineType.PREPROCESS: + return self._load_preprocess_pipeline_cls(workload_type, arch) + elif pipeline_type == PipelineType.BASIC: + return self.pipelines[pipeline_type.value][arch][ + pipeline_name_in_config + ] + elif pipeline_type == PipelineType.TRAINING: + pass + else: + raise ValueError(f"Invalid pipeline type: {pipeline_type.value}") + + return None + + def resolve_pipeline_cls( + self, + pipeline_name_in_config: str, + pipeline_type: PipelineType, + workload_type: WorkloadType, + ) -> type[ComposedPipelineBase] | type[LoRAPipeline]: + """Resolve pipeline class based on pipeline name in the config, pipeline type, and workload + type.""" + if not pipeline_name_in_config: + logger.warning("No pipeline architecture is specified") + + pipeline_cls = self._try_load_pipeline_cls( + pipeline_name_in_config, pipeline_type, workload_type + ) + if pipeline_cls is not None: + return pipeline_cls + supported_archs = self.get_supported_archs( + pipeline_name_in_config, pipeline_type + ) + raise ValueError( + f"Pipeline architecture '{pipeline_name_in_config}' is not supported for pipeline type '{pipeline_type.value}' " + f"and workload type '{workload_type.value}'. " + f"Supported architectures: {supported_archs}" + ) + + +@lru_cache +def import_pipeline_classes( + pipeline_types: list[PipelineType] | PipelineType | None = None, +) -> dict[str, dict[str, dict[str, type[ComposedPipelineBase] | None]]]: + """Import pipeline classes based on the pipeline type and workload type. + + Args: + pipeline_types: The pipeline types to load (basic, preprocess, training). + If None, loads all types. + + Returns: + A three-level nested dictionary: + {pipeline_type: {architecture_name: {pipeline_name: pipeline_cls}}} + e.g., {"basic": {"wan": {"WanPipeline": WanPipeline}}} + """ + type_to_arch_to_pipeline_dict: dict[ + str, dict[str, dict[str, type[ComposedPipelineBase] | None]] + ] = {} + package_name: str = "fastvideo.pipelines" + + # Determine which pipeline types to scan + if isinstance(pipeline_types, list): + pipeline_types_to_scan = [ + pipeline_type.value for pipeline_type in pipeline_types + ] + elif isinstance(pipeline_types, PipelineType): + pipeline_types_to_scan = [pipeline_types.value] + else: + pipeline_types_to_scan = [pt.value for pt in PipelineType] + + logger.info("Loading pipelines for types: %s", pipeline_types_to_scan) + + for pipeline_type_str in pipeline_types_to_scan: + arch_to_pipeline_dict: dict[ + str, dict[str, type[ComposedPipelineBase] | None] + ] = {} + + # Try to load from pipeline-type-specific directory first + pipeline_type_package_name = f"{package_name}.{pipeline_type_str}" + + try: + pipeline_type_package = importlib.import_module( + pipeline_type_package_name + ) + logger.debug("Successfully imported %s", pipeline_type_package_name) + + for _, arch, ispkg in pkgutil.iter_modules( + pipeline_type_package.__path__ + ): + pipeline_dict: dict[str, type[ComposedPipelineBase] | None] = {} + + arch_package_name = f"{pipeline_type_package_name}.{arch}" + if ispkg: + arch_package = importlib.import_module(arch_package_name) + for _, module_name, ispkg in pkgutil.walk_packages( + arch_package.__path__, arch_package_name + "." + ): + if not ispkg: + pipeline_module = importlib.import_module( + module_name + ) + if hasattr(pipeline_module, "EntryClass"): + if isinstance(pipeline_module.EntryClass, list): + for pipeline in pipeline_module.EntryClass: + pipeline_name = pipeline.__name__ + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[pipeline_name] = pipeline + else: + pipeline_name = ( + pipeline_module.EntryClass.__name__ + ) + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[pipeline_name] = ( + pipeline_module.EntryClass + ) + + arch_to_pipeline_dict[arch] = pipeline_dict + + except ImportError as e: + raise ImportError( + f"Could not import {pipeline_type_package_name} when importing pipeline classes: {e}" + ) from None + + type_to_arch_to_pipeline_dict[pipeline_type_str] = arch_to_pipeline_dict + + # Log summary + total_pipelines = sum( + len(pipeline_dict) + for arch_to_pipeline_dict in type_to_arch_to_pipeline_dict.values() + for pipeline_dict in arch_to_pipeline_dict.values() + ) + logger.info( + "Loaded %d pipeline classes across %d types", + total_pipelines, + len(pipeline_types_to_scan), + ) + + return type_to_arch_to_pipeline_dict + + +def get_pipeline_registry( + pipeline_type: PipelineType | str | None = None, +) -> _PipelineRegistry: + """Get a pipeline registry for the specified mode, pipeline type, and workload type. + + Args: + pipeline_type: Pipeline type to load. If None and mode is provided, will be derived from mode. + + Returns: + A pipeline registry instance. + """ + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_classes = import_pipeline_classes(pipeline_type) + return _PipelineRegistry(pipeline_classes) diff --git a/worldcompass/fastvideo/pipelines/stages/__init__.py b/worldcompass/fastvideo/pipelines/stages/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..829b1f9a0e737e44eadf8f47d102b30de8391433 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/stages/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline stages for diffusion models. + +This package contains the various stages that can be composed to create complete diffusion +pipelines. +""" + +from fastvideo.pipelines.stages.base import PipelineStage + +__all__ = ["PipelineStage"] diff --git a/worldcompass/fastvideo/pipelines/stages/base.py b/worldcompass/fastvideo/pipelines/stages/base.py new file mode 100644 index 0000000000000000000000000000000000000000..fe6587641a19d980625127db73c17ce06903af66 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/stages/base.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Base classes for pipeline stages. + +This module defines the abstract base classes for pipeline stages that can be composed to create +complete diffusion pipelines. +""" + +import time +import traceback +from abc import ABC, abstractmethod + +import torch + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class StageVerificationError(Exception): + """Exception raised when stage verification fails.""" + + pass + + +class PipelineStage(ABC): + """Abstract base class for all pipeline stages. + + A pipeline stage represents a discrete step in the diffusion process that can be composed with + other stages to create a complete pipeline. Each stage is responsible for a specific part of the + process, such as prompt encoding, latent preparation, etc. + """ + + def verify_input( + self, batch: ForwardBatch, fastvideo_args: FastVideoArgs + ) -> VerificationResult: + """Verify the input for the stage. + + Example: + from fastvideo.pipelines.stages.validators import V, VerificationResult + + def verify_input(self, batch, fastvideo_args): + result = VerificationResult() + result.add_check("height", batch.height, V.positive_int_divisible(8)) + result.add_check("width", batch.width, V.positive_int_divisible(8)) + result.add_check("image_latent", batch.image_latent, V.is_tensor) + return result + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + A VerificationResult containing the verification status. + """ + # Default implementation - no verification + return VerificationResult() + + def verify_output( + self, batch: ForwardBatch, fastvideo_args: FastVideoArgs + ) -> VerificationResult: + """Verify the output for the stage. + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + A VerificationResult containing the verification status. + """ + # Default implementation - no verification + return VerificationResult() + + def _run_verification( + self, + verification_result: VerificationResult, + stage_name: str, + verification_type: str, + ) -> None: + """Run verification and raise errors if any checks fail. + + Args: + verification_result: Results from verify_input or verify_output + stage_name: Name of the current stage + verification_type: "input" or "output" + """ + if not verification_result.is_valid(): + failed_fields = verification_result.get_failed_fields() + if failed_fields: + # Get detailed failure information + detailed_summary = verification_result.get_failure_summary() + + failed_fields_str = ", ".join(failed_fields) + error_msg = ( + f"{verification_type.capitalize()} verification failed for {stage_name}: " + f"Failed fields: {failed_fields_str}\n" + f"Details: {detailed_summary}" + ) + raise StageVerificationError(error_msg) + + @property + def device(self) -> torch.device: + """Get the device for this stage.""" + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def set_logging(self, enable: bool): + """Enable or disable logging for this stage. + + Args: + enable: Whether to enable logging. + """ + self._enable_logging = enable + + def __call__( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Execute the stage's processing on the batch with optional verification and logging. + Should not be overridden by subclasses. + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + stage_name = self.__class__.__name__ + + # Check if verification is enabled (simple approach for prototype) + enable_verification = getattr( + fastvideo_args, "enable_stage_verification", False + ) + + if enable_verification: + # Pre-execution input verification + try: + input_result = self.verify_input(batch, fastvideo_args) + self._run_verification(input_result, stage_name, "input") + except Exception as e: + logger.error( + "Input verification failed for %s: %s", stage_name, str(e) + ) + raise + + # Execute the actual stage logic + if envs.FASTVIDEO_STAGE_LOGGING: + logger.info("[%s] Starting execution", stage_name) + start_time = time.perf_counter() + + try: + result = self.forward(batch, fastvideo_args) + execution_time = time.perf_counter() - start_time + logger.info( + "[%s] Execution completed in %s ms", + stage_name, + execution_time * 1000, + ) + batch.logging_info.add_stage_execution_time( + stage_name, execution_time + ) + except Exception as e: + execution_time = time.perf_counter() - start_time + logger.error( + "[%s] Error during execution after %s ms: %s", + stage_name, + execution_time * 1000, + e, + ) + logger.error( + "[%s] Traceback: %s", stage_name, traceback.format_exc() + ) + raise + else: + # Direct execution (current behavior) + result = self.forward(batch, fastvideo_args) + + if enable_verification: + # Post-execution output verification + try: + output_result = self.verify_output(result, fastvideo_args) + self._run_verification(output_result, stage_name, "output") + except Exception as e: + logger.error( + "Output verification failed for %s: %s", stage_name, str(e) + ) + raise + + return result + + @abstractmethod + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """Forward pass of the stage's processing. + + This method should be implemented by subclasses to provide the forward + processing logic for the stage. + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + raise NotImplementedError + + def backward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + raise NotImplementedError diff --git a/worldcompass/fastvideo/pipelines/stages/validators.py b/worldcompass/fastvideo/pipelines/stages/validators.py new file mode 100644 index 0000000000000000000000000000000000000000..a3401d0e3e8436cb5335f72f3ecd8043998fcc30 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/stages/validators.py @@ -0,0 +1,515 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Common validators for pipeline stage verification. + +This module provides reusable validation functions that can be used across all pipeline stages for +input/output verification. +""" + +from collections.abc import Callable +from typing import Any + +import torch + + +class StageValidators: + """Common validators for pipeline stages.""" + + @staticmethod + def not_none(value: Any) -> bool: + """Check if value is not None.""" + return value is not None + + @staticmethod + def positive_int(value: Any) -> bool: + """Check if value is a positive integer.""" + return isinstance(value, int) and value > 0 + + @staticmethod + def positive_float(value: Any) -> bool: + """Check if value is a positive float.""" + return isinstance(value, int | float) and value > 0 + + @staticmethod + def non_negative_float(value: Any) -> bool: + """Check if value is a non-negative float.""" + return isinstance(value, int | float) and value >= 0 + + @staticmethod + def divisible_by(value: Any, divisor: int) -> bool: + """Check if value is divisible by divisor.""" + return ( + value is not None + and isinstance(value, int) + and value % divisor == 0 + ) + + @staticmethod + def is_tensor(value: Any) -> bool: + """Check if value is a torch tensor and doesn't contain NaN values.""" + if not isinstance(value, torch.Tensor): + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_with_dims(value: Any, dims: int) -> bool: + """Check if value is a tensor with specific dimensions and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if value.dim() != dims: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_min_dims(value: Any, min_dims: int) -> bool: + """Check if value is a tensor with at least min_dims dimensions and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if value.dim() < min_dims: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def tensor_shape_matches(value: Any, expected_shape: tuple) -> bool: + """Check if tensor shape matches expected shape (None for any size) and no NaN values.""" + if not isinstance(value, torch.Tensor): + return False + if len(value.shape) != len(expected_shape): + return False + for actual, expected in zip(value.shape, expected_shape, strict=True): + if expected is not None and actual != expected: + return False + return not torch.isnan(value).any().item() + + @staticmethod + def list_not_empty(value: Any) -> bool: + """Check if value is a non-empty list.""" + return isinstance(value, list) and len(value) > 0 + + @staticmethod + def list_length(value: Any, length: int) -> bool: + """Check if list has specific length.""" + return isinstance(value, list) and len(value) == length + + @staticmethod + def list_min_length(value: Any, min_length: int) -> bool: + """Check if list has at least min_length items.""" + return isinstance(value, list) and len(value) >= min_length + + @staticmethod + def string_not_empty(value: Any) -> bool: + """Check if value is a non-empty string.""" + return isinstance(value, str) and len(value.strip()) > 0 + + @staticmethod + def string_or_list_strings(value: Any) -> bool: + """Check if value is a string or list of strings.""" + if isinstance(value, str): + return True + if isinstance(value, list): + return all(isinstance(item, str) for item in value) + return False + + @staticmethod + def bool_value(value: Any) -> bool: + """Check if value is a boolean.""" + return isinstance(value, bool) + + @staticmethod + def generator_or_list_generators(value: Any) -> bool: + """Check if value is a Generator or list of Generators.""" + if isinstance(value, torch.Generator): + return True + if isinstance(value, list): + return all(isinstance(item, torch.Generator) for item in value) + return False + + @staticmethod + def is_list(value: Any) -> bool: + """Check if value is a list (can be empty).""" + return isinstance(value, list) + + @staticmethod + def is_tuple(value: Any) -> bool: + """Check if value is a tuple.""" + return isinstance(value, tuple) + + @staticmethod + def none_or_tensor(value: Any) -> bool: + """Check if value is None or a tensor without NaN values.""" + if value is None: + return True + if not isinstance(value, torch.Tensor): + return False + return not torch.isnan(value).any().item() + + @staticmethod + def list_of_tensors_with_dims(value: Any, dims: int) -> bool: + """Check if value is a non-empty list where all items are tensors with specific dimensions + and no NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if item.dim() != dims: + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def list_of_tensors(value: Any) -> bool: + """Check if value is a non-empty list where all items are tensors without NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool: + """Check if value is a non-empty list where all items are tensors with at least min_dims + dimensions and no NaN values.""" + if not isinstance(value, list) or len(value) == 0: + return False + for item in value: + if not isinstance(item, torch.Tensor): + return False + if item.dim() < min_dims: + return False + if torch.isnan(item).any().item(): + return False + return True + + @staticmethod + def none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is None or a tensor with specific dimensions and + no NaN values.""" + + def validator(value: Any) -> bool: + if value is None: + return True + if not isinstance(value, torch.Tensor): + return False + if value.dim() != dims: + return False + return not torch.isnan(value).any().item() + + return validator + + @staticmethod + def none_or_list(value: Any) -> bool: + """Check if value is None or a list.""" + return value is None or isinstance(value, list) + + @staticmethod + def none_or_positive_int(value: Any) -> bool: + """Check if value is None or a positive integer.""" + return value is None or (isinstance(value, int) and value > 0) + + # Helper methods that return functions for common patterns + @staticmethod + def with_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if tensor has specific dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.tensor_with_dims(value, dims) + + return validator + + @staticmethod + def min_dims(min_dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if tensor has at least min_dims dimensions and no NaN + values.""" + + def validator(value: Any) -> bool: + return StageValidators.tensor_min_dims(value, min_dims) + + return validator + + @staticmethod + def divisible(divisor: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is divisible by divisor.""" + + def validator(value: Any) -> bool: + return StageValidators.divisible_by(value, divisor) + + return validator + + @staticmethod + def positive_int_divisible(divisor: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a positive integer divisible by divisor.""" + + def validator(value: Any) -> bool: + return ( + isinstance(value, int) + and value > 0 + and StageValidators.divisible_by(value, divisor) + ) + + return validator + + @staticmethod + def list_of_tensors_dims(dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a list of tensors with specific dimensions and + no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.list_of_tensors_with_dims(value, dims) + + return validator + + @staticmethod + def list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]: + """Return a validator that checks if value is a list of tensors with at least min_dims + dimensions and no NaN values.""" + + def validator(value: Any) -> bool: + return StageValidators.list_of_tensors_with_min_dims( + value, min_dims + ) + + return validator + + +class ValidationFailure: + """Details about a specific validation failure.""" + + def __init__( + self, + validator_name: str, + actual_value: Any, + expected: str | None = None, + error_msg: str | None = None, + ): + self.validator_name = validator_name + self.actual_value = actual_value + self.expected = expected + self.error_msg = error_msg + + def __str__(self) -> str: + parts = [f"Validator '{self.validator_name}' failed"] + + if self.error_msg: + parts.append(f"Error: {self.error_msg}") + + # Add actual value info (but limit very long representations) + actual_str = self._format_value(self.actual_value) + parts.append(f"Actual: {actual_str}") + + if self.expected: + parts.append(f"Expected: {self.expected}") + + return ". ".join(parts) + + def _format_value(self, value: Any) -> str: + """Format a value for display in error messages.""" + if value is None: + return "None" + elif isinstance(value, torch.Tensor): + return f"tensor(shape={list(value.shape)}, dtype={value.dtype})" + elif isinstance(value, list): + if len(value) == 0: + return "[]" + elif len(value) <= 3: + item_strs = [self._format_value(item) for item in value] + return f"[{', '.join(item_strs)}]" + else: + return f"list(length={len(value)}, first_item={self._format_value(value[0])})" + elif isinstance(value, str): + if len(value) > 50: + return f"'{value[:47]}...'" + else: + return f"'{value}'" + else: + return f"{type(value).__name__}({value})" + + +class VerificationResult: + """Wrapper class for stage verification results.""" + + def __init__(self) -> None: + self._checks: dict[str, bool] = {} + self._failures: dict[str, list[ValidationFailure]] = {} + + def add_check( + self, + field_name: str, + value: Any, + validators: Callable[[Any], bool] | list[Callable[[Any], bool]], + ) -> "VerificationResult": + """Add a validation check for a field. + + Args: + field_name: Name of the field being checked + value: The actual value to validate + validators: Single validation function or list of validation functions. + Each function will be called with the value as its first argument. + + Returns: + Self for method chaining + + Examples: + # Single validator + result.add_check("tensor", my_tensor, V.is_tensor) + + # Multiple validators (all must pass) + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + + # Using partial functions for parameters + result.add_check("height", batch.height, [V.not_none, V.divisible(8)]) + """ + if not isinstance(validators, list): + validators = [validators] + + failures = [] + all_passed = True + + # Apply all validators and collect detailed failure info + for validator in validators: + try: + passed = validator(value) + if not passed: + all_passed = False + failure = self._create_validation_failure(validator, value) + failures.append(failure) + except Exception as e: + # If any validator raises an exception, consider the check failed + all_passed = False + validator_name = getattr(validator, "__name__", str(validator)) + failure = ValidationFailure( + validator_name=validator_name, + actual_value=value, + error_msg=f"Exception during validation: {str(e)}", + ) + failures.append(failure) + + self._checks[field_name] = all_passed + if not all_passed: + self._failures[field_name] = failures + + return self + + def _create_validation_failure( + self, validator: Callable, value: Any + ) -> ValidationFailure: + """Create a ValidationFailure with detailed information.""" + validator_name = getattr(validator, "__name__", str(validator)) + + # Try to extract meaningful expected value info based on validator type + expected = None + error_msg = None + + # Handle common validator patterns + if hasattr(validator, "__closure__") and validator.__closure__: + # This is likely a closure (like our helper functions) + if "dims" in validator_name or "with_dims" in str(validator): + if isinstance(value, torch.Tensor): + expected = f"tensor with {validator.__closure__[0].cell_contents} dimensions" + else: + expected = "tensor with specific dimensions" + elif "divisible" in str(validator): + expected = f"integer divisible by {validator.__closure__[0].cell_contents}" + + # Handle specific validator types and check for NaN values + if validator_name == "is_tensor": + expected = "torch.Tensor without NaN values" + if ( + isinstance(value, torch.Tensor) + and torch.isnan(value).any().item() + ): + error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values" + elif validator_name == "positive_int": + expected = "positive integer" + elif validator_name == "not_none": + expected = "non-None value" + elif validator_name == "list_not_empty": + expected = "non-empty list" + elif validator_name == "bool_value": + expected = "boolean value" + elif ( + "tensor_with_dims" in validator_name + or "tensor_min_dims" in validator_name + ): + if isinstance(value, torch.Tensor): + if torch.isnan(value).any().item(): + error_msg = f"tensor has {value.dim()} dimensions but contains {torch.isnan(value).sum().item()} NaN values" + else: + error_msg = f"tensor has {value.dim()} dimensions" + elif validator_name == "is_list": + expected = "list" + elif validator_name == "none_or_tensor": + expected = "None or tensor without NaN values" + if ( + isinstance(value, torch.Tensor) + and torch.isnan(value).any().item() + ): + error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values" + elif validator_name == "list_of_tensors": + expected = "non-empty list of tensors without NaN values" + if isinstance(value, list) and len(value) > 0: + nan_count = 0 + for item in value: + if ( + isinstance(item, torch.Tensor) + and torch.isnan(item).any().item() + ): + nan_count += torch.isnan(item).sum().item() + if nan_count > 0: + error_msg = f"list contains tensors with total {nan_count} NaN values" + elif "list_of_tensors_with_dims" in validator_name: + expected = "non-empty list of tensors with specific dimensions and no NaN values" + if isinstance(value, list) and len(value) > 0: + nan_count = 0 + for item in value: + if ( + isinstance(item, torch.Tensor) + and torch.isnan(item).any().item() + ): + nan_count += torch.isnan(item).sum().item() + if nan_count > 0: + error_msg = f"list contains tensors with total {nan_count} NaN values" + + return ValidationFailure( + validator_name=validator_name, + actual_value=value, + expected=expected, + error_msg=error_msg, + ) + + def is_valid(self) -> bool: + """Check if all validations passed.""" + return all(self._checks.values()) + + def get_failed_fields(self) -> list[str]: + """Get list of fields that failed validation.""" + return [field for field, passed in self._checks.items() if not passed] + + def get_detailed_failures(self) -> dict[str, list[ValidationFailure]]: + """Get detailed failure information for each failed field.""" + return self._failures.copy() + + def get_failure_summary(self) -> str: + """Get a comprehensive summary of all validation failures.""" + if self.is_valid(): + return "All validations passed" + + summary_parts = [] + for field_name, failures in self._failures.items(): + field_summary = f"\n Field '{field_name}':" + for i, failure in enumerate(failures, 1): + field_summary += f"\n {i}. {failure}" + summary_parts.append(field_summary) + + return "Validation failures:" + "".join(summary_parts) + + def to_dict(self) -> dict: + """Convert to dictionary for backward compatibility.""" + return self._checks.copy() + + +# Alias for convenience +V = StageValidators diff --git a/worldcompass/fastvideo/pipelines/training/__init__.py b/worldcompass/fastvideo/pipelines/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cc6d0f1f1965330689fcf6b671a1ab8b7eef84a2 --- /dev/null +++ b/worldcompass/fastvideo/pipelines/training/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Training pipelines for fastvideo.v1. + +This package contains pipelines for training diffusion models. +""" diff --git a/worldcompass/fastvideo/platforms/__init__.py b/worldcompass/fastvideo/platforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5e91a471364f57e2f45e499d2e8f2d39221e1706 --- /dev/null +++ b/worldcompass/fastvideo/platforms/__init__.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/__init__.py + +import traceback +from typing import TYPE_CHECKING + +from fastvideo.logger import init_logger + +# imported by other files, do not remove +from fastvideo.platforms.interface import AttentionBackendEnum # noqa: F401 +from fastvideo.platforms.interface import Platform, PlatformEnum +from fastvideo.utils import resolve_obj_by_qualname + +logger = init_logger(__name__) + + +def cuda_platform_plugin() -> str | None: + is_cuda = False + + try: + from fastvideo.utils import import_pynvml + + pynvml = import_pynvml() # type: ignore[no-untyped-call] + pynvml.nvmlInit() + try: + # NOTE: Edge case: fastvideo cpu build on a GPU machine. + # Third-party pynvml can be imported in cpu build, + # we need to check if fastvideo is built with cpu too. + # Otherwise, fastvideo will always activate cuda plugin + # on a GPU machine, even if in a cpu build. + is_cuda = pynvml.nvmlDeviceGetCount() > 0 + finally: + pynvml.nvmlShutdown() + except Exception as e: + if "nvml" not in e.__class__.__name__.lower(): + # If the error is not related to NVML, re-raise it. + raise e + + # CUDA is supported on Jetson, but NVML may not be. + import os + + def cuda_is_jetson() -> bool: + return os.path.isfile("/etc/nv_tegra_release") or os.path.exists( + "/sys/class/tegra-firmware" + ) + + if cuda_is_jetson(): + is_cuda = True + if is_cuda: + logger.info("CUDA is available") + + return "fastvideo.platforms.cuda.CudaPlatform" if is_cuda else None + + +def mps_platform_plugin() -> str | None: + """Detect if MPS (Metal Performance Shaders) is available on macOS.""" + is_mps = False + + try: + import torch + + if torch.backends.mps.is_available(): + is_mps = True + logger.info("MPS (Metal Performance Shaders) is available") + except Exception as e: + logger.info("MPS detection failed: %s", e) + + return "fastvideo.platforms.mps.MpsPlatform" if is_mps else None + + +def cpu_platform_plugin() -> str | None: + """Detect if CPU platform should be used.""" + # CPU is always available as a fallback + return "fastvideo.platforms.cpu.CpuPlatform" + + +def rocm_platform_plugin() -> str | None: + is_rocm = False + + try: + import amdsmi + + amdsmi.amdsmi_init() + try: + if len(amdsmi.amdsmi_get_processor_handles()) > 0: + is_rocm = True + logger.info("ROCm platform is available") + finally: + amdsmi.amdsmi_shut_down() + except Exception as e: + logger.info("ROCm platform is unavailable: %s", e) + + return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None + + +builtin_platform_plugins = { + "cuda": cuda_platform_plugin, + "rocm": rocm_platform_plugin, + "mps": mps_platform_plugin, + "cpu": cpu_platform_plugin, +} + + +def resolve_current_platform_cls_qualname() -> str: + # TODO(will): if we need to support other platforms, we should consider if + # vLLM's plugin architecture is suitable for our needs. + + # Try MPS first on macOS + platform_cls_qualname = mps_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to ROCm + platform_cls_qualname = rocm_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CUDA + platform_cls_qualname = cuda_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + # Fall back to CPU as last resort + platform_cls_qualname = cpu_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + + raise RuntimeError( + "No platform plugin found. Please check your installation." + ) + + +_current_platform = None +_init_trace: str = "" + +if TYPE_CHECKING: + current_platform: Platform + + +def __getattr__(name: str): + if name == "current_platform": + # lazy init current_platform. + # 1. out-of-tree platform plugins need `from fastvideo.platforms import + # Platform` so that they can inherit `Platform` class. Therefore, + # we cannot resolve `current_platform` during the import of + # `fastvideo.platforms`. + # 2. when users use out-of-tree platform plugins, they might run + # `import fastvideo`, some fastvideo internal code might access + # `current_platform` during the import, and we need to make sure + # `current_platform` is only resolved after the plugins are loaded + # (we have tests for this, if any developer violate this, they will + # see the test failures). + global _current_platform + if _current_platform is None: + platform_cls_qualname = resolve_current_platform_cls_qualname() + _current_platform = resolve_obj_by_qualname(platform_cls_qualname)() + global _init_trace + _init_trace = "".join(traceback.format_stack()) + return _current_platform + elif name in globals(): + return globals()[name] + else: + raise AttributeError( + f"No attribute named '{name}' exists in {__name__}." + ) + + +__all__ = ["Platform", "PlatformEnum", "current_platform", "_init_trace"] diff --git a/worldcompass/fastvideo/platforms/cpu.py b/worldcompass/fastvideo/platforms/cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..f778180074344f13e9a5f15807f771fbbca36969 --- /dev/null +++ b/worldcompass/fastvideo/platforms/cpu.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py + +import platform + +import torch + +from fastvideo.platforms.interface import CpuArchEnum, Platform, PlatformEnum + + +class CpuPlatform(Platform): + _enum = PlatformEnum.CPU + device_name = "CPU" + device_type = "cpu" + dispatch_key = "CPU" + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture.""" + machine = platform.machine().lower() + if machine in ("x86_64", "amd64", "i386", "i686"): + return CpuArchEnum.X86 + elif machine in ("arm64", "aarch64"): + return CpuArchEnum.ARM + else: + return CpuArchEnum.UNSPECIFIED + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return platform.processor() + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + return platform.machine() + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + # This is a rough estimate for CPU memory + # In practice, you might want to use psutil or similar + return 0 + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + return True + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + # For CPU, we can't easily get memory usage without additional libraries + return 0.0 + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cpu_communicator.CpuCommunicator" diff --git a/worldcompass/fastvideo/platforms/cuda.py b/worldcompass/fastvideo/platforms/cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..9af1c75f09c0ad59d9d8a3e92334b6b4c599e8b6 --- /dev/null +++ b/worldcompass/fastvideo/platforms/cuda.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py +"""Code inside this file can safely assume cuda platform, e.g. importing pynvml. + +However, it should not initialize cuda context. +""" + +import os +from collections.abc import Callable +from functools import lru_cache, wraps +from typing import TypeVar + +import torch +from typing_extensions import ParamSpec + +import fastvideo.envs as envs +from fastvideo.logger import init_logger +from fastvideo.platforms.interface import ( + AttentionBackendEnum, + DeviceCapability, + Platform, + PlatformEnum, +) +from fastvideo.utils import import_pynvml + +logger = init_logger(__name__) + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +pynvml = import_pynvml() # type: ignore[no-untyped-call] + +# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models +# see https://github.com/huggingface/diffusers/issues/9704 for details +torch.backends.cuda.enable_cudnn_sdp(False) + + +def device_id_to_physical_device_id(device_id: int) -> int: + if "CUDA_VISIBLE_DEVICES" in os.environ: + device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",") + if device_ids == [""]: + msg = ( + "CUDA_VISIBLE_DEVICES is set to empty string, which means" + " GPU support is disabled. If you are using ray, please unset" + " the environment variable `CUDA_VISIBLE_DEVICES` inside the" + " worker/actor. " + "Check https://github.com/vllm-project/vllm/issues/8402 for" + " more information." + ) + raise RuntimeError(msg) + physical_device_id = device_ids[device_id] + return int(physical_device_id) + else: + return device_id + + +def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: + + @wraps(fn) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + pynvml.nvmlInit() + try: + return fn(*args, **kwargs) + finally: + pynvml.nvmlShutdown() + + return wrapper + + +class CudaPlatformBase(Platform): + _enum = PlatformEnum.CUDA + device_name: str = "cuda" + device_type: str = "cuda" + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability( + cls, device_id: int = 0 + ) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used" + ) + return False + return True + + @classmethod + def is_full_nvlink(cls, device_ids: list[int]) -> bool: + raise NotImplementedError + + @classmethod + def log_warnings(cls) -> None: + pass + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + # TODO(will): maybe come up with a more general interface for local attention + # if distributed is False, we always try to use Flash attn + + logger.info( + "Trying FASTVIDEO_ATTENTION_BACKEND=%s", + envs.FASTVIDEO_ATTENTION_BACKEND, + ) + if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: + try: + from st_attn import sliding_tile_attention # noqa: F401 + + from fastvideo.attention.backends.sliding_tile_attn import ( # noqa: F401 + SlidingTileAttentionBackend, + ) + + logger.info("Using Sliding Tile Attention backend.") + + return "fastvideo.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Sliding Tile Attention backend: %s", + str(e), + ) + raise ImportError( + "Sliding Tile Attention backend is not installed. " + ) from e + elif selected_backend == AttentionBackendEnum.SAGE_ATTN: + try: + from sageattention import sageattn # noqa: F401 + + from fastvideo.attention.backends.sage_attn import ( # noqa: F401 + SageAttentionBackend, + ) + + logger.info("Using Sage Attention backend.") + + return "fastvideo.attention.backends.sage_attn.SageAttentionBackend" + except ImportError as e: + logger.info(e) + logger.info( + "Sage Attention backend is not installed. Fall back to Flash Attention." + ) + elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + try: + from vsa import block_sparse_attn # noqa: F401 + + from fastvideo.attention.backends.video_sparse_attn import ( # noqa: F401 + VideoSparseAttentionBackend, + ) + + logger.info("Using Video Sparse Attention backend.") + + return "fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Video Sparse Attention backend: %s", + str(e), + ) + raise ImportError( + "Video Sparse Attention backend is not installed. " + ) from e + elif selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + elif ( + selected_backend == AttentionBackendEnum.FLASH_ATTN + or selected_backend is None + ): + pass + elif selected_backend: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + + target_backend = AttentionBackendEnum.FLASH_ATTN + if not cls.has_device_capability(80): + logger.info( + "Cannot use FlashAttention-2 backend for Volta and Turing GPUs." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + elif dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention-2 backend for dtype other than " + "torch.float16 or torch.bfloat16." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + # FlashAttn is valid for the model, checking if the package is + # installed. + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from fastvideo.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend, + ) + + supported_sizes = ( + FlashAttentionBackend.get_supported_head_sizes() + ) + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info( + "Cannot use FlashAttention-2 backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "fastvideo.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + + +# NVML utils +# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, +# all the related functions work on real physical device ids. +# the major benefit of using NVML is that it will not initialize CUDA +class NvmlCudaPlatform(CudaPlatformBase): + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_capability( + cls, device_id: int = 0 + ) -> DeviceCapability | None: + try: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) + return DeviceCapability(major=major, minor=minor) + except RuntimeError: + return None + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + try: + return bool(super().has_device_capability(capability, device_id)) + except RuntimeError: + return False + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_name(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + return cls._get_physical_device_name(physical_device_id) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_uuid(cls, device_id: int = 0) -> str: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return str(pynvml.nvmlDeviceGetUUID(handle)) + + @classmethod + @lru_cache(maxsize=8) + @with_nvml_context + def get_device_total_memory(cls, device_id: int = 0) -> int: + physical_device_id = device_id_to_physical_device_id(device_id) + handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + + @classmethod + @with_nvml_context + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + """Query if the set of gpus are fully connected by nvlink (1 hop)""" + handles = [ + pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids + ] + for i, handle in enumerate(handles): + for j, peer_handle in enumerate(handles): + if i < j: + try: + p2p_status = pynvml.nvmlDeviceGetP2PStatus( + handle, + peer_handle, + pynvml.NVML_P2P_CAPS_INDEX_NVLINK, + ) + if p2p_status != pynvml.NVML_P2P_STATUS_OK: + return False + except pynvml.NVMLError: + logger.exception( + "NVLink detection failed. This is normal if" + " your machine has no NVLink equipped." + ) + return False + return True + + @classmethod + def _get_physical_device_name(cls, device_id: int = 0) -> str: + handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) + return str(pynvml.nvmlDeviceGetName(handle)) + + @classmethod + @with_nvml_context + def log_warnings(cls) -> None: + device_ids: int = pynvml.nvmlDeviceGetCount() + if device_ids > 1: + device_names = [ + cls._get_physical_device_name(i) for i in range(device_ids) + ] + if ( + len(set(device_names)) > 1 + and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID" + ): + logger.warning( + "Detected different devices in the system: %s. Please" + " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to " + "avoid unexpected behavior.", + ", ".join(device_names), + ) + + +class NonNvmlCudaPlatform(CudaPlatformBase): + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + device_props = torch.cuda.get_device_properties(device_id) + return int(device_props.total_memory) + + @classmethod + def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: + logger.exception( + "NVLink detection not possible, as context support was" + " not found. Assuming no NVLink available." + ) + return False + + +# Autodetect either NVML-enabled or non-NVML platform +# based on whether NVML is available. +nvml_available = False +try: + try: + pynvml.nvmlInit() + nvml_available = True + except Exception: + # On Jetson, NVML is not supported. + nvml_available = False +finally: + if nvml_available: + pynvml.nvmlShutdown() + +CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform + +try: + from sphinx.ext.autodoc.mock import _MockModule + + if not isinstance(pynvml, _MockModule): + CudaPlatform.log_warnings() +except ModuleNotFoundError: + CudaPlatform.log_warnings() diff --git a/worldcompass/fastvideo/platforms/interface.py b/worldcompass/fastvideo/platforms/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..036dd6dfb856adda6a750ea79ddbd26bdee15a9f --- /dev/null +++ b/worldcompass/fastvideo/platforms/interface.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py + +import enum +import random +from typing import NamedTuple + +import numpy as np +import torch + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class AttentionBackendEnum(enum.Enum): + FLASH_ATTN = enum.auto() + SLIDING_TILE_ATTN = enum.auto() + TORCH_SDPA = enum.auto() + SAGE_ATTN = enum.auto() + VIDEO_SPARSE_ATTN = enum.auto() + NO_ATTENTION = enum.auto() + + +class PlatformEnum(enum.Enum): + CUDA = enum.auto() + ROCM = enum.auto() + TPU = enum.auto() + CPU = enum.auto() + MPS = enum.auto() + OOT = enum.auto() + UNSPECIFIED = enum.auto() + + +class CpuArchEnum(enum.Enum): + X86 = enum.auto() + ARM = enum.auto() + UNSPECIFIED = enum.auto() + + +class DeviceCapability(NamedTuple): + major: int + minor: int + + def as_version_str(self) -> str: + return f"{self.major}.{self.minor}" + + def to_int(self) -> int: + """Express device capability as an integer ````. + + It is assumed that the minor version is always a single digit. + """ + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class Platform: + _enum: PlatformEnum + device_name: str + device_type: str + + # available dispatch keys: + # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa + # use "CPU" as a fallback for platforms not registered in PyTorch + dispatch_key: str = "CPU" + + # The torch.compile backend for compiling simple and + # standalone functions. The default value is "inductor" to keep + # the same behavior as PyTorch. + # NOTE: for the forward part of the model, vLLM has another separate + # compilation strategy. + simple_compile_backend: str = "inductor" + + supported_quantization: list[str] = [] + + def is_cuda(self) -> bool: + return self._enum == PlatformEnum.CUDA + + def is_rocm(self) -> bool: + return self._enum == PlatformEnum.ROCM + + def is_tpu(self) -> bool: + return self._enum == PlatformEnum.TPU + + def is_cpu(self) -> bool: + return self._enum == PlatformEnum.CPU + + def is_out_of_tree(self) -> bool: + return self._enum == PlatformEnum.OOT + + def is_cuda_alike(self) -> bool: + """Stateless version of :func:`torch.cuda.is_available`.""" + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + + def is_mps(self) -> bool: + return self._enum == PlatformEnum.MPS + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + """Get the attention backend class of a device.""" + return "" + + @classmethod + def get_device_capability( + cls, + device_id: int = 0, + ) -> DeviceCapability | None: + """Stateless version of :func:`torch.cuda.get_device_capability`.""" + return None + + @classmethod + def has_device_capability( + cls, + capability: tuple[int, int] | int, + device_id: int = 0, + ) -> bool: + """Test whether this platform is compatible with a device capability. + + The ``capability`` argument can either be: + + - A tuple ``(major, minor)``. + - An integer ````. (See :meth:`DeviceCapability.to_int`) + """ + current_capability = cls.get_device_capability(device_id=device_id) + if current_capability is None: + return False + + if isinstance(capability, tuple): + return current_capability >= capability + + return current_capability.to_int() >= capability + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + """Get the name of a device.""" + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + """Get the uuid of a device, e.g. the PCI bus ID.""" + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + """Get the total memory of a device in bytes.""" + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + """Check if the current platform supports async output.""" + raise NotImplementedError + + @classmethod + def inference_mode(cls): + """A device-specific wrapper of `torch.inference_mode`. + + This wrapper is recommended because some hardware backends such as TPU + do not support `torch.inference_mode`. In such a case, they will fall + back to `torch.no_grad` by overriding this method. + """ + return torch.inference_mode(mode=True) + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """Set the seed of each random module. `torch.manual_seed` will set seed on all devices. + + Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 + """ + if seed is not None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + @classmethod + def verify_model_arch(cls, model_arch: str) -> None: + """Verify whether the current platform supports the specified model architecture. + + - This will raise an Error or Warning based on the model support on + the current platform. + - By default all models are considered supported. + """ + pass + + @classmethod + def verify_quantization(cls, quant: str) -> None: + """Verify whether the quantization is supported by the current platform.""" + if ( + cls.supported_quantization + and quant not in cls.supported_quantization + ): + raise ValueError( + f"{quant} quantization is currently not supported in {cls.device_name}." + ) + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + """Return the memory usage in bytes.""" + raise NotImplementedError + + @classmethod + def get_device_communicator_cls(cls) -> str: + """Get device specific communicator class for distributed communication.""" + return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa + + @classmethod + def get_cpu_architecture(cls) -> CpuArchEnum: + """Get the CPU architecture of the current platform.""" + return CpuArchEnum.UNSPECIFIED + + +class UnspecifiedPlatform(Platform): + _enum = PlatformEnum.UNSPECIFIED + device_type = "" diff --git a/worldcompass/fastvideo/platforms/mps.py b/worldcompass/fastvideo/platforms/mps.py new file mode 100644 index 0000000000000000000000000000000000000000..80ac272ec76f33e647b63b8792c256e3f80bdba6 --- /dev/null +++ b/worldcompass/fastvideo/platforms/mps.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from fastvideo.logger import init_logger +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.platforms.interface import ( + DeviceCapability, + Platform, + PlatformEnum, +) + +logger = init_logger(__name__) + + +class MpsPlatform(Platform): + _enum = PlatformEnum.MPS + device_name: str = "mps" + device_type: str = "mps" + dispatch_key: str = "MPS" + device_control_env_var: str = "MPS_VISIBLE_DEVICES" + + @classmethod + def get_device_capability( + cls, device_id: int = 0 + ) -> DeviceCapability | None: + raise NotImplementedError + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_uuid(cls, device_id: int = 0) -> str: + raise NotImplementedError + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + raise NotImplementedError + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable MPS " + "graph. Since, enforce-eager is enabled, async output " + "processor cannot be used" + ) + return False + return True + + @classmethod + def get_current_memory_usage( + cls, device: torch.types.Device | None = None + ) -> float: + return 0.0 + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible + logger.info("Using Torch SDPA backend for MPS.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + # Use base communicator for MPS + return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + """Set the seed for MPS device.""" + if seed is not None: + import random + + import numpy as np + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + # MPS doesn't have manual_seed_all like CUDA + # The manual_seed above should be sufficient diff --git a/worldcompass/fastvideo/platforms/rocm.py b/worldcompass/fastvideo/platforms/rocm.py new file mode 100644 index 0000000000000000000000000000000000000000..60479664f11999280d6af6466b036230737b0b38 --- /dev/null +++ b/worldcompass/fastvideo/platforms/rocm.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py +"""This file is a platform abstraction for ROCm GPUs, adjusted to match the structure and interface +of `cuda.py`.""" + +import torch + +import fastvideo.envs as envs +from fastvideo.logger import init_logger +from fastvideo.platforms.interface import ( + AttentionBackendEnum, + DeviceCapability, + Platform, + PlatformEnum, +) + +logger = init_logger(__name__) + + +# ROCm uses the same torch.cuda interface +class RocmPlatform(Platform): + _enum = PlatformEnum.ROCM + device_name: str = "rocm" + device_type: str = "cuda" # torch uses 'cuda' backend string + dispatch_key: str = "CUDA" + device_control_env_var: str = "CUDA_VISIBLE_DEVICES" + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return str(torch.cuda.get_device_name(device_id)) + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + return torch.cuda.get_device_properties(device_id).total_memory + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + if enforce_eager: + logger.warning( + "To see benefits of async output processing, enable CUDA graph. " + "Since enforce-eager is enabled, async output processor cannot be used" + ) + return False + return True + + @classmethod + def log_warnings(cls) -> None: + pass # ROCm-specific warnings can be added here + + @classmethod + def get_current_memory_usage( + cls, device: torch.device | None = None + ) -> float: + torch.cuda.reset_peak_memory_stats(device) + return float(torch.cuda.max_memory_allocated(device)) + + @classmethod + def get_attn_backend_cls( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + logger.info( + "Trying FASTVIDEO_ATTENTION_BACKEND=%s", + envs.FASTVIDEO_ATTENTION_BACKEND, + ) + + if selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + + elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None): + pass + + elif selected_backend in ( + AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.SAGE_ATTN, + ): + raise ValueError( + f"{selected_backend.name} is not supported on {cls.device_name}." + ) + elif selected_backend: + raise ValueError( + f"Invalid attention backend for {cls.device_name}: {selected_backend}" + ) + + target_backend = AttentionBackendEnum.FLASH_ATTN + if dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention backend for dtype other than " + "torch.float16 or torch.bfloat16." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.FLASH_ATTN: + try: + import flash_attn # noqa: F401 + + from fastvideo.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend, + ) + + supported_sizes = ( + FlashAttentionBackend.get_supported_head_sizes() + ) + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info( + "Cannot use FlashAttention backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend.") + + return "fastvideo.attention.backends.sdpa.SDPABackend" + + logger.info("Using Flash Attention backend.") + + return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend" + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too diff --git a/worldcompass/fastvideo/third_party/__init__.py b/worldcompass/fastvideo/third_party/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/fastvideo/third_party/pynvml.py b/worldcompass/fastvideo/third_party/pynvml.py new file mode 100644 index 0000000000000000000000000000000000000000..9dd8526eab7f218f373ec9469bbad35554be2ddf --- /dev/null +++ b/worldcompass/fastvideo/third_party/pynvml.py @@ -0,0 +1,7404 @@ +# SPDX-License-Identifier: Apache-2.0 +# copied from https://pypi.org/project/nvidia-ml-py +# version 12.570.86 + +##### +# Copyright (c) 2011-2023, NVIDIA Corporation. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA Corporation nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +##### + +## +# Python bindings for the NVML library +## +from ctypes import * +from ctypes.util import find_library +from functools import wraps +import sys +import os +import threading +import string + +## C Type mappings ## +## Enums +_nvmlEnableState_t = c_uint +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 + +_nvmlBrandType_t = c_uint +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 17 + +_nvmlTemperatureThresholds_t = c_uint +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 + +_nvmlTemperatureSensors_t = c_uint +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 + + +_nvmlComputeMode_t = c_uint +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 +NVML_COMPUTEMODE_COUNT = 4 + +_nvmlMemoryLocation_t = c_uint +NVML_MEMORY_LOCATION_L1_CACHE = 0 +NVML_MEMORY_LOCATION_L2_CACHE = 1 +NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 +NVML_MEMORY_LOCATION_DRAM = 2 +NVML_MEMORY_LOCATION_REGISTER_FILE = 3 +NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 +NVML_MEMORY_LOCATION_TEXTURE_SHM = 5 +NVML_MEMORY_LOCATION_CBU = 6 +NVML_MEMORY_LOCATION_SRAM = 7 +NVML_MEMORY_LOCATION_COUNT = 8 + +NVML_NVLINK_MAX_LINKS = 18 + +# For backwards compatibility, maintain the incorrectly-named "LANES" define +NVML_NVLINK_MAX_LANES = NVML_NVLINK_MAX_LINKS + +_nvmlNvLinkErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_REPLAY = 0 +NVML_NVLINK_ERROR_DL_RECOVERY = 1 +NVML_NVLINK_ERROR_DL_CRC_FLIT = 2 +NVML_NVLINK_ERROR_DL_CRC_DATA = 3 +NVML_NVLINK_ERROR_DL_ECC_DATA = 4 +NVML_NVLINK_ERROR_COUNT = 5 + +_nvmlNvLinkEccLaneErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_ECC_LANE0 = 0 +NVML_NVLINK_ERROR_DL_ECC_LANE1 = 1 +NVML_NVLINK_ERROR_DL_ECC_LANE2 = 2 +NVML_NVLINK_ERROR_DL_ECC_LANE3 = 3 +NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 + +_nvmlNvLinkCapability_t = c_uint +NVML_NVLINK_CAP_P2P_SUPPORTED = 0 +NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 + +_nvmlNvLinkUtilizationCountPktTypes_t = c_uint +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF + +_nvmlNvLinkUtilizationCountUnits_t = c_uint +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 + +_nvmlNvLinkDeviceType_t = c_uint +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF + +# These are deprecated, instead use _nvmlMemoryErrorType_t +_nvmlEccBitType_t = c_uint +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 +NVML_ECC_ERROR_TYPE_COUNT = 2 + +_nvmlEccCounterType_t = c_uint +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 +NVML_ECC_COUNTER_TYPE_COUNT = 2 + +_nvmlMemoryErrorType_t = c_uint +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 + +_nvmlClockType_t = c_uint +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 + +_nvmlClockId_t = c_uint +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 +NVML_CLOCK_ID_COUNT = 4 + +_nvmlDriverModel_t = c_uint +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 + +NVML_MAX_GPU_PERF_PSTATES = 16 + +_nvmlPstates_t = c_uint +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 + +_nvmlInforomObject_t = c_uint +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 + +_nvmlReturn_t = c_uint +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_UNKNOWN = 999 + +_nvmlFanState_t = c_uint +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 + +_nvmlFanControlPolicy_t = c_uint +NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 +NVML_FAN_POLICY_MANUAL = 1 + +_nvmlLedColor_t = c_uint +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 + +_nvmlGpuOperationMode_t = c_uint +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 + +_nvmlPageRetirementCause_t = c_uint +NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 + +_nvmlRestrictedAPI_t = c_uint +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 + +_nvmlBridgeChipType_t = c_uint +NVML_BRIDGE_CHIP_PLX = 0 +NVML_BRIDGE_CHIP_BRO4 = 1 +NVML_MAX_PHYSICAL_BRIDGE = 128 + +_nvmlValueType_t = c_uint +NVML_VALUE_TYPE_DOUBLE = 0 +NVML_VALUE_TYPE_UNSIGNED_INT = 1 +NVML_VALUE_TYPE_UNSIGNED_LONG = 2 +NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 +NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4 +NVML_VALUE_TYPE_SIGNED_INT = 5 +NVML_VALUE_TYPE_UNSIGNED_SHORT = 6 +NVML_VALUE_TYPE_COUNT = 7 + +_nvmlNvlinkVersion_t = c_uint +NVML_NVLINK_VERSION_INVALID = 0 +NVML_NVLINK_VERSION_1_0 = 1 +NVML_NVLINK_VERSION_2_0 = 2 +NVML_NVLINK_VERSION_2_2 = 3 +NVML_NVLINK_VERSION_3_0 = 4 +NVML_NVLINK_VERSION_3_1 = 5 +NVML_NVLINK_VERSION_4_0 = 6 +NVML_NVLINK_VERSION_5_0 = 7 + +_nvmlPerfPolicyType_t = c_uint +NVML_PERF_POLICY_POWER = 0 +NVML_PERF_POLICY_THERMAL = 1 +NVML_PERF_POLICY_SYNC_BOOST = 2 +NVML_PERF_POLICY_BOARD_LIMIT = 3 +NVML_PERF_POLICY_LOW_UTILIZATION = 4 +NVML_PERF_POLICY_RELIABILITY = 5 +NVML_PERF_POLICY_TOTAL_APP_CLOCKS = 10 +NVML_PERF_POLICY_TOTAL_BASE_CLOCKS = 11 +NVML_PERF_POLICY_COUNT = 12 + +_nvmlEncoderQueryType_t = c_uint +NVML_ENCODER_QUERY_H264 = 0 +NVML_ENCODER_QUERY_HEVC = 1 +NVML_ENCODER_QUERY_AV1 = 2 +NVML_ENCODER_QUERY_UNKNOWN = 255 + +_nvmlFBCSessionType_t = c_uint +NVML_FBC_SESSION_TYPE_UNKNOWN = 0 +NVML_FBC_SESSION_TYPE_TOSYS = 1 +NVML_FBC_SESSION_TYPE_CUDA = 2 +NVML_FBC_SESSION_TYPE_VID = 3 +NVML_FBC_SESSION_TYPE_HWENC = 4 + +_nvmlDetachGpuState_t = c_uint +NVML_DETACH_GPU_KEEP = 0 +NVML_DETACH_GPU_REMOVE = 1 + +_nvmlPcieLinkState_t = c_uint +NVML_PCIE_LINK_KEEP = 0 +NVML_PCIE_LINK_SHUT_DOWN = 1 + +_nvmlSamplingType_t = c_uint +NVML_TOTAL_POWER_SAMPLES = 0 +NVML_GPU_UTILIZATION_SAMPLES = 1 +NVML_MEMORY_UTILIZATION_SAMPLES = 2 +NVML_ENC_UTILIZATION_SAMPLES = 3 +NVML_DEC_UTILIZATION_SAMPLES = 4 +NVML_PROCESSOR_CLK_SAMPLES = 5 +NVML_MEMORY_CLK_SAMPLES = 6 +NVML_MODULE_POWER_SAMPLES = 7 +NVML_JPG_UTILIZATION_SAMPLES = 8 +NVML_OFA_UTILIZATION_SAMPLES = 9 +NVML_SAMPLINGTYPE_COUNT = 10 + +_nvmlPcieUtilCounter_t = c_uint +NVML_PCIE_UTIL_TX_BYTES = 0 +NVML_PCIE_UTIL_RX_BYTES = 1 +NVML_PCIE_UTIL_COUNT = 2 + +_nvmlGpuTopologyLevel_t = c_uint +NVML_TOPOLOGY_INTERNAL = 0 +NVML_TOPOLOGY_SINGLE = 10 +NVML_TOPOLOGY_MULTIPLE = 20 +NVML_TOPOLOGY_HOSTBRIDGE = 30 +NVML_TOPOLOGY_NODE = 40 +NVML_TOPOLOGY_CPU = NVML_TOPOLOGY_NODE +NVML_TOPOLOGY_SYSTEM = 50 + +_nvmlGpuP2PCapsIndex_t = c_uint +NVML_P2P_CAPS_INDEX_READ = (0,) +NVML_P2P_CAPS_INDEX_WRITE = 1 +NVML_P2P_CAPS_INDEX_NVLINK = 2 +NVML_P2P_CAPS_INDEX_ATOMICS = 3 +# +# NVML_P2P_CAPS_INDEX_PROP is deprecated. +# Use NVML_P2P_CAPS_INDEX_PCI instead. +# +NVML_P2P_CAPS_INDEX_PROP = 4 +NVML_P2P_CAPS_INDEX_PCI = 4 +NVML_P2P_CAPS_INDEX_UNKNOWN = 5 + +_nvmlGpuP2PStatus_t = c_uint +NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED +NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4 +NVML_P2P_STATUS_NOT_SUPPORTED = 5 +NVML_P2P_STATUS_UNKNOWN = 6 + +_nvmlDeviceArchitecture_t = c_uint +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_T23X = 11 +NVML_DEVICE_ARCH_UNKNOWN = 0xFFFFFFFF + +# PCI bus Types +_nvmlBusType_t = c_uint +NVML_BUS_TYPE_UNKNOWN = 0 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 + +_nvmlPowerSource_t = c_uint +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 + +_nvmlAdaptiveClockInfoStatus_t = c_uint +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000 +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 + +_nvmlClockLimitId_t = c_uint +NVML_CLOCK_LIMIT_ID_RANGE_START = 0xFFFFFF00 +NVML_CLOCK_LIMIT_ID_TDP = 0xFFFFFF01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xFFFFFF02 + +_nvmlPcieLinkMaxSpeed_t = c_uint +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 +NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 +NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 + +_nvmlPcieAtomicsCapability_t = c_uint +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 + +_nvmlAffinityScope_t = c_uint +NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_SOCKET = 1 + +_nvmlDeviceGpuRecoveryAction_t = c_uint +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 + +# C preprocessor defined values +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 + +NVML_MAX_GPC_COUNT = 32 + +# buffer size +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = ( + 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +) +NVML_PERF_MODES_BUFFER_SIZE = 2048 + +# Format strings +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" + +NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) +NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) +""" + Field Identifiers. + + All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. +""" +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive + +# ECC Count Totals +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = ( + 5 # Total single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = ( + 6 # Total double bit aggregate (persistent) ECC errors +) +# Individual ECC locations +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = ( + 15 # Texture memory single bit volatile ECC errors +) +NVML_FI_DEV_ECC_DBE_VOL_TEX = ( + 16 # Texture memory double bit volatile ECC errors +) +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = ( + 18 # L1 cache single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_L1 = ( + 19 # L1 cache double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_L2 = ( + 20 # L2 cache single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_L2 = ( + 21 # L2 cache double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_DEV = ( + 22 # Device memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_DEV = ( + 23 # Device memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_REG = ( + 24 # Register File single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_REG = ( + 25 # Register File double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_SBE_AGG_TEX = ( + 26 # Texture memory single bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_TEX = ( + 27 # Texture memory double bit aggregate (persistent) ECC errors +) +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors + +# Page Retirement +NVML_FI_DEV_RETIRED_SBE = ( + 29 # Number of retired pages because of single bit errors +) +NVML_FI_DEV_RETIRED_DBE = ( + 30 # Number of retired pages because of double bit errors +) +NVML_FI_DEV_RETIRED_PENDING = ( + 31 # If any pages are pending retirement. 1=yes. 0=no. +) + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = ( + 32 # NVLink flow control CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = ( + 33 # NVLink flow control CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = ( + 34 # NVLink flow control CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = ( + 35 # NVLink flow control CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = ( + 36 # NVLink flow control CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = ( + 37 # NVLink flow control CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = ( + 38 # NVLink flow control CRC Error Counter total for all Lanes +) + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = ( + 39 # NVLink data CRC Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = ( + 40 # NVLink data CRC Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = ( + 41 # NVLink data CRC Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = ( + 42 # NVLink data CRC Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = ( + 43 # NVLink data CRC Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = ( + 44 # NVLink data CRC Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = ( + 45 # NvLink data CRC Error Counter total for all Lanes +) + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = ( + 46 # NVLink Replay Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = ( + 47 # NVLink Replay Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = ( + 48 # NVLink Replay Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = ( + 49 # NVLink Replay Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = ( + 50 # NVLink Replay Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = ( + 51 # NVLink Replay Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = ( + 52 # NVLink Replay Error Counter total for all Lanes +) + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = ( + 53 # NVLink Recovery Error Counter for Lane 0 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = ( + 54 # NVLink Recovery Error Counter for Lane 1 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = ( + 55 # NVLink Recovery Error Counter for Lane 2 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = ( + 56 # NVLink Recovery Error Counter for Lane 3 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = ( + 57 # NVLink Recovery Error Counter for Lane 4 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = ( + 58 # NVLink Recovery Error Counter for Lane 5 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = ( + 59 # NVLink Recovery Error Counter total for all Lanes +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = ( + 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = ( + 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = ( + 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = ( + 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = ( + 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = ( + 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = ( + 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = ( + 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = ( + 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = ( + 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = ( + 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = ( + 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = ( + 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = ( + 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes +) + +# Perf Policy Counters +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = ( + 76 # Perf Policy Counter for Sync boost Policy +) +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = ( + 78 # Perf Policy Counter for Low GPU Utilization Policy +) +NVML_FI_DEV_PERF_POLICY_RELIABILITY = ( + 79 # Perf Policy Counter for Reliability Policy +) +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = ( + 80 # Perf Policy Counter for Total App Clock Policy +) +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = ( + 81 # Perf Policy Counter for Total Base Clocks Policy +) + +# Memory temperatures +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device + +# Energy Counter +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 + +# NVLink Link Count +NVML_FI_DEV_NVLINK_LINK_COUNT = 91 + +# Page Retirement pending fields +NVML_FI_DEV_RETIRED_PENDING_SBE = 92 +NVML_FI_DEV_RETIRED_PENDING_DBE = 93 + +# PCIe replay and replay rollover counters +NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 +NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = ( + 96 # NVLink flow control CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = ( + 97 # NVLink flow control CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = ( + 98 # NVLink flow control CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = ( + 99 # NVLink flow control CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = ( + 100 # NVLink flow control CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = ( + 101 # NVLink flow control CRC Error Counter for Lane 11 +) + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = ( + 102 # NVLink data CRC Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = ( + 103 # NVLink data CRC Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = ( + 104 # NVLink data CRC Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = ( + 105 # NVLink data CRC Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = ( + 106 # NVLink data CRC Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = ( + 107 # NVLink data CRC Error Counter for Lane 11 +) + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = ( + 108 # NVLink Replay Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = ( + 109 # NVLink Replay Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = ( + 110 # NVLink Replay Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = ( + 111 # NVLink Replay Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = ( + 112 # NVLink Replay Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = ( + 113 # NVLink Replay Error Counter for Lane 11 +) + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = ( + 114 # NVLink Recovery Error Counter for Lane 6 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = ( + 115 # NVLink Recovery Error Counter for Lane 7 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = ( + 116 # NVLink Recovery Error Counter for Lane 8 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = ( + 117 # NVLink Recovery Error Counter for Lane 9 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = ( + 118 # NVLink Recovery Error Counter for Lane 10 +) +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = ( + 119 # NVLink Recovery Error Counter for Lane 11 +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = ( + 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = ( + 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = ( + 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = ( + 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = ( + 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = ( + 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 +) + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = ( + 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = ( + 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = ( + 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = ( + 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = ( + 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +) +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = ( + 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 +) + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 + +# NVLink Throughput Counters +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = ( + 140 # NVLink TX Data + protocol overhead in KiB +) +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = ( + 141 # NVLink RX Data + protocol overhead in KiB +) + +# Row Remapper +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 + +# Remote device NVLink ID +NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 + +# Number of NVLinks connected to NVSwitch +NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 + +# NvLink ECC Data Error Counters +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = ( + 148 # < NVLink data ECC Error Counter for Link 0 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = ( + 149 # < NVLink data ECC Error Counter for Link 1 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = ( + 150 # < NVLink data ECC Error Counter for Link 2 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = ( + 151 # < NVLink data ECC Error Counter for Link 3 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = ( + 152 # < NVLink data ECC Error Counter for Link 4 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = ( + 153 # < NVLink data ECC Error Counter for Link 5 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = ( + 154 # < NVLink data ECC Error Counter for Link 6 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = ( + 155 # < NVLink data ECC Error Counter for Link 7 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = ( + 156 # < NVLink data ECC Error Counter for Link 8 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = ( + 157 # < NVLink data ECC Error Counter for Link 9 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = ( + 158 # < NVLink data ECC Error Counter for Link 10 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = ( + 159 # < NVLink data ECC Error Counter for Link 11 +) +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = ( + 160 # < NvLink data ECC Error Counter total for all Links +) + +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 + +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 + +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 + +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 + +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 + +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 + +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 + +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 + +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 + +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 + +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 + +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 + +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = ( + 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +) +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 + +NVML_FI_DEV_RESET_STATUS = ( + 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = ( + 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +) +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 + +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = ( + 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = ( + 256 # Max % TMP Floor value +) +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = ( + 257 # Min % TMP Floor value +) +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = ( + 258 # HW Circuitry % lifetime remaining +) +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = ( + 259 # Max number of preset profiles +) +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = ( + 260 # % TMP floor for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = ( + 261 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = ( + 262 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = ( + 263 # Ramp down hysteresis value in ms for a given profile +) +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = ( + 264 # Active preset profile number +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = ( + 265 # % TMP floor for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = ( + 266 # Ramp up rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = ( + 267 # Ramp down rate in mW/s for a given profile +) +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = ( + 268 # Ramp down hysteresis value in ms for a given profile +) + +NVML_FI_MAX = 269 # One greater than the largest field ID defined above + +# NVML_FI_DEV_NVLINK_GET_STATE state enums +NVML_NVLINK_STATE_INACTIVE = 0x0 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 + +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = ( + 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = ( + 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +) + +## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = ( + 1 # Device is associated with GPU-Passthorugh +) +NVML_GPU_VIRTUALIZATION_MODE_VGPU = ( + 2 # Device is associated with vGPU inside virtual machine. +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = ( + 3 # Device is associated with VGX hypervisor in vGPU mode +) +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = ( + 4 # Device is associated with VGX hypervisor in vSGA mode +) + +## Lib loading ## +nvmlLib = None +libLoadLock = threading.Lock() +_nvmlLib_refcount = ( + 0 # Incremented on each nvmlInit and decremented on nvmlShutdown +) + +## vGPU Management +_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuInstance_t = c_uint + +_nvmlVgpuVmIdType_t = c_uint +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 + +_nvmlGridLicenseFeatureCode_t = c_uint +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = ( + 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +) +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 + +_nvmlGridLicenseExpiryStatus_t = c_uint8 +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = ( + 0, +) # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = ( + 1, +) # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = (2,) # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = (3,) # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = (4,) # Permanent expiry + +_nvmlVgpuCapability_t = c_uint +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = ( + 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +) +NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +NVML_VGPU_CAP_COUNT = 5 + +_nvmlVgpuDriverCapability_t = c_uint +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = ( + 0 # Supports mixing of different vGPU profiles within one guest VM +) +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 + +_nvmlDeviceVgpuCapability_t = c_uint +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = ( + 6 # Set/Get support of mini-quarter vGPU profiles +) +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = ( + 7 # Set/Get support for compute media engine vGPU profiles +) +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = ( + 8 # Query whether the GPU supports FSR and warm update +) +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_COUNT = 10 + +_nvmlVgpuGuestInfoState_t = c_uint +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 + +_nvmlVgpuVmCompatibility_t = c_uint +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 + +_nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 + +_nvmlHostVgpuMode_t = c_uint +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 + +_nvmlConfComputeGpusReadyState_t = c_uint +NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 +NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1 + +_nvmlConfComputeGpuCaps_t = c_uint +NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0 +NVML_CC_SYSTEM_GPUS_CC_CAPABLE = 1 + +_nvmlConfComputeCpuCaps_t = c_uint +NVML_CC_SYSTEM_CPU_CAPS_NONE = 0 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1 +NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV_SNP = 3 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SNP_VTOM = 4 + +_nvmlConfComputeDevToolsMode_t = c_uint +NVML_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0 +NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 + +NVML_CC_SYSTEM_MULTIGPU_NONE = 0 +NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 + +NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 +NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 +NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 + +_nvmlConfComputeCcFeature_t = c_uint +NVML_CC_SYSTEM_FEATURE_DISABLED = 0 +NVML_CC_SYSTEM_FEATURE_ENABLED = 1 + +_nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50 +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 + +# GSP firmware +NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + + +class NVMLLibraryMismatchError(Exception): + pass + + +## Error Checking ## +class NVMLError(Exception): + _valClassMapping = dict() + # List of currently known error codes + _errcode_to_string = { + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + + def __new__(typ, value): + """Maps value to a proper subclass of NVMLError. + + See _extractNVMLErrorsAsClasses function for more details + """ + if typ == NVMLError: + typ = NVMLError._valClassMapping.get(value, typ) + obj = Exception.__new__(typ) + obj.value = value + return obj + + def __str__(self): + try: + if self.value not in NVMLError._errcode_to_string: + NVMLError._errcode_to_string[self.value] = str( + nvmlErrorString(self.value) + ) + return NVMLError._errcode_to_string[self.value] + except NVMLError: + return "NVML Error with code %d" % self.value + + def __eq__(self, other): + return self.value == other.value + + +def nvmlExceptionClass(nvmlErrorCode): + if nvmlErrorCode not in NVMLError._valClassMapping: + raise ValueError("nvmlErrorCode %s is not valid" % nvmlErrorCode) + return NVMLError._valClassMapping[nvmlErrorCode] + + +def _extractNVMLErrorsAsClasses(): + """Generates a hierarchy of classes on top of NVMLError class. + + Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate + exceptions more easily. + + NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. e.g. + NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized + """ + this_module = sys.modules[__name__] + nvmlErrorsNames = [ + x for x in dir(this_module) if x.startswith("NVML_ERROR_") + ] + for err_name in nvmlErrorsNames: + # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized + class_name = "NVMLError_" + string.capwords( + err_name.replace("NVML_ERROR_", ""), "_" + ).replace("_", "") + err_val = getattr(this_module, err_name) + + def gen_new(val): + def new(typ): + obj = NVMLError.__new__(typ, val) + return obj + + return new + + new_error_class = type( + class_name, (NVMLError,), {"__new__": gen_new(err_val)} + ) + new_error_class.__module__ = __name__ + setattr(this_module, class_name, new_error_class) + NVMLError._valClassMapping[err_val] = new_error_class + + +_extractNVMLErrorsAsClasses() + + +def _nvmlCheckReturn(ret): + if ret != NVML_SUCCESS: + raise NVMLError(ret) + return ret + + +## Function access ## +_nvmlGetFunctionPointer_cache = ( + dict() +) # function pointers are cached to prevent unnecessary libLoadLock locking + + +def _nvmlGetFunctionPointer(name): + global nvmlLib + + if name in _nvmlGetFunctionPointer_cache: + return _nvmlGetFunctionPointer_cache[name] + + libLoadLock.acquire() + try: + # ensure library was loaded + if nvmlLib == None: + raise NVMLError(NVML_ERROR_UNINITIALIZED) + try: + _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) + return _nvmlGetFunctionPointer_cache[name] + except AttributeError: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + + +## Alternative object +# Allows the object to be printed +# Allows mismatched types to be assigned +# - like None when the Structure variant requires c_uint +class nvmlFriendlyObject(object): + def __init__(self, dictionary): + for x in dictionary: + setattr(self, x, dictionary[x]) + + def __str__(self): + return self.__dict__.__str__() + + +def nvmlStructToFriendlyObject(struct): + d = {} + for x in struct._fields_: + key = x[0] + value = getattr(struct, key) + # only need to convert from bytes if bytes, no need to check python version. + d[key] = value.decode() if isinstance(value, bytes) else value + obj = nvmlFriendlyObject(d) + return obj + + +# pack the object so it can be passed to the NVML library +def nvmlFriendlyObjectToStruct(obj, model): + for x in model._fields_: + key = x[0] + value = obj.__dict__[key] + # any c_char_p in python3 needs to be bytes, default encoding works fine. + if sys.version_info >= (3,): + setattr(model, key, value.encode()) + else: + setattr(model, key, value) + return model + + +## Unit structures +class struct_c_nvmlUnit_t(Structure): + pass # opaque handle + + +c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + + +class _PrintableStructure(Structure): + """Abstract class that produces nicer __str__ output than ctypes.Structure. e.g. instead of: >>> + print str(obj) this class will print + class_name(field_name: formatted_value, field_name: formatted_value) + + _fmt_ dictionary of -> e.g. class that has _field_ 'hex_value', + c_uint could be formatted with _fmt_ = {"hex_value" : "%08X"} to produce nicer output. Default + fomratting string for all fields can be set with key "" like: _fmt_ = {"" : + "%d MHz"} # e.g all values are numbers in MHz. If not set it's assumed to be just "%s" + + Exact format of returned str from this class is subject to change in the future. + """ + + _fmt_ = {} + + def __str__(self): + result = [] + for x in self._fields_: + key = x[0] + value = getattr(self, key) + fmt = "%s" + if key in self._fmt_: + fmt = self._fmt_[key] + elif "" in self._fmt_: + fmt = self._fmt_[""] + result.append(("%s: " + fmt) % (key, value)) + return self.__class__.__name__ + "(" + ", ".join(result) + ")" + + def __getattribute__(self, name): + res = super(_PrintableStructure, self).__getattribute__(name) + # need to convert bytes to unicode for python3 don't need to for python2 + # Python 2 strings are of both str and bytes + # Python 3 strings are not of type bytes + # ctypes should convert everything to the correct values otherwise + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + def __setattr__(self, name, value): + if isinstance(value, str): + # encoding a python2 string returns the same value, since python2 strings are bytes already + # bytes passed in python3 will be ignored. + value = value.encode() + super(_PrintableStructure, self).__setattr__(name, value) + + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), + ] + + +class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): + _fields_ = [("isC2cEnabled", c_uint)] + + +nvmlC2cModeInfo_v1 = 0x1000008 + + +class c_nvmlLedState_t(_PrintableStructure): + _fields_ = [ + ("cause", c_char * 256), + ("color", _nvmlLedColor_t), + ] + + +class c_nvmlPSUInfo_t(_PrintableStructure): + _fields_ = [ + ("state", c_char * 256), + ("current", c_uint), + ("voltage", c_uint), + ("power", c_uint), + ] + + +class c_nvmlUnitFanInfo_t(_PrintableStructure): + _fields_ = [ + ("speed", c_uint), + ("state", _nvmlFanState_t), + ] + + +class c_nvmlUnitFanSpeeds_t(_PrintableStructure): + _fields_ = [("fans", c_nvmlUnitFanInfo_t * 24), ("count", c_uint)] + + +## Device structures +class struct_c_nvmlDevice_t(Structure): + pass # opaque handle + + +c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + + +class nvmlPciInfoExt_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + ("pciSubSystemId", c_uint), + ("baseClass", c_uint), + ("subClass", c_uint), + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + "version": "0x%04X", + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + "baseClass": "0x%01X", + "subClass": "0x%01X", + } + + +nvmlPciInfoExt_v1 = 0x1000040 + + +# Legacy pciInfo used for _v1 and _v2 +class nvmlPciInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + # Added in 2.285 + ("pciSubSystemId", c_uint), + ("reserved0", c_uint), + ("reserved1", c_uint), + ("reserved2", c_uint), + ("reserved3", c_uint), + ] + _fmt_ = { + "domain": "0x%04X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + + +class nvmlPciInfo_t(_PrintableStructure): + _fields_ = [ + # Moved to the new busId location below + ("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ("domain", c_uint), + ("bus", c_uint), + ("device", c_uint), + ("pciDeviceId", c_uint), + # Added in 2.285 + ("pciSubSystemId", c_uint), + # New busId replaced the long deprecated and reserved fields with a + # field of the same size in 9.0 + ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + "domain": "0x%08X", + "bus": "0x%02X", + "device": "0x%02X", + "pciDeviceId": "0x%08X", + "pciSubSystemId": "0x%08X", + } + + +class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ] + + +SystemDriverBranchInfo_v1 = 0x1000054 + + +class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): + _fields_ = [ + ("pci", nvmlPciInfo_t), + ("uuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), + ] + + +class nvmlNvLinkUtilizationControl_t(_PrintableStructure): + _fields_ = [ + ("units", _nvmlNvLinkUtilizationCountUnits_t), + ("pktfilter", _nvmlNvLinkUtilizationCountPktTypes_t), + ] + + +class c_nvmlMemory_t(_PrintableStructure): + _fields_ = [ + ("total", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +class c_nvmlMemory_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("total", c_ulonglong), + ("reserved", c_ulonglong), + ("free", c_ulonglong), + ("used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +nvmlMemory_v2 = 0x02000028 + + +class c_nvmlBAR1Memory_t(_PrintableStructure): + _fields_ = [ + ("bar1Total", c_ulonglong), + ("bar1Free", c_ulonglong), + ("bar1Used", c_ulonglong), + ] + _fmt_ = {"": "%d B"} + + +class nvmlClkMonFaultInfo_t(Structure): + _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)] + + +MAX_CLK_DOMAINS = 32 + + +class nvmlClkMonStatus_t(Structure): + _fields_ = [ + ("bGlobalStatus", c_uint), + ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS), + ] + + +# On Windows with the WDDM driver, usedGpuMemory is reported as None +# Code that processes this structure should check for None, I.E. +# +# if (info.usedGpuMemory == None): +# # TODO handle the error +# pass +# else: +# print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) +# endif +# +# See NVML documentation for more information +class c_nvmlProcessInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ] + _fmt_ = {"usedGpuMemory": "%d B"} + + +c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t + +c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t + +_nvmlProcessMode_t = c_uint +NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_GRAPHICS = 1 +NVML_PROCESS_MODE_MPS = 2 + + +class c_nvmlProcessDetail_v1_t(Structure): + _fields_ = [ + ("pid", c_uint), + ("usedGpuMemory", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ("usedGpuCcProtectedMemory", c_ulonglong), + ] + + +class c_nvmlProcessDetailList_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("mode", _nvmlProcessMode_t), + ("numProcArrayEntries", c_uint), + ("procArray", POINTER(c_nvmlProcessDetail_v1_t)), + ] + _fmt_ = {"numProcArrayEntries": "%d B"} + + +c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t + +nvmlProcessDetailList_v1 = 0x1000018 + + +class c_nvmlBridgeChipInfo_t(_PrintableStructure): + _fields_ = [ + ("type", _nvmlBridgeChipType_t), + ("fwVersion", c_uint), + ] + + +class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): + _fields_ = [ + ("bridgeCount", c_uint), + ("bridgeChipInfo", c_nvmlBridgeChipInfo_t * 128), + ] + + +class c_nvmlEccErrorCounts_t(_PrintableStructure): + _fields_ = [ + ("l1Cache", c_ulonglong), + ("l2Cache", c_ulonglong), + ("deviceMemory", c_ulonglong), + ("registerFile", c_ulonglong), + ] + + +class c_nvmlUtilization_t(_PrintableStructure): + _fields_ = [ + ("gpu", c_uint), + ("memory", c_uint), + ] + _fmt_ = {"": "%d %%"} + + +# Added in 2.285 +class c_nvmlHwbcEntry_t(_PrintableStructure): + _fields_ = [ + ("hwbcId", c_uint), + ("firmwareVersion", c_char * 32), + ] + + +class c_nvmlValue_t(Union): + _fields_ = [ + ("dVal", c_double), + ("uiVal", c_uint), + ("ulVal", c_ulong), + ("ullVal", c_ulonglong), + ("sllVal", c_longlong), + ("siVal", c_int), + ("usVal", c_ushort), + ] + + +class c_nvmlSample_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("sampleValue", c_nvmlValue_t), + ] + + +class c_nvmlViolationTime_t(_PrintableStructure): + _fields_ = [ + ("referenceTime", c_ulonglong), + ("violationTime", c_ulonglong), + ] + + +class c_nvmlFieldValue_t(_PrintableStructure): + _fields_ = [ + ("fieldId", c_uint32), + ("scopeId", c_uint32), + ("timestamp", c_int64), + ("latencyUsec", c_int64), + ("valueType", _nvmlValueType_t), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_nvmlValue_t), + ] + + +NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 + +nvmlNvlinkSupportedBwModes_v1 = 0x100001C + + +class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("bwModes", c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ("totalBwModes", c_uint8), + ] + + def __init__(self): + super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__( + version=nvmlNvlinkSupportedBwModes_v1 + ) + + +nvmlNvlinkGetBwMode_v1 = 0x100000C + + +class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): + _fields_ = [("version", c_uint), ("bIsBest", c_uint), ("bwMode", c_uint8)] + + def __init__(self): + super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__( + version=nvmlNvlinkGetBwMode_v1 + ) + + +nvmlNvlinkSetBwMode_v1 = 0x100000C + + +class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): + _fields_ = [("version", c_uint), ("bSetBest", c_uint), ("bwMode", c_uint8)] + + def __init__(self): + super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__( + version=nvmlNvlinkSetBwMode_v1 + ) + + +class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("mode", c_uint), + ] + + +VgpuHeterogeneousMode_v1 = 0x1000008 + + +class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("placementId", c_uint), + ] + + +VgpuPlacementId_v1 = 0x1000008 + + +class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("count", c_uint), + ("placementSize", c_uint), + ("placementIds", POINTER(c_uint)), + ] + + +VgpuPlacementList_v1 = 0x1000018 + +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + + +class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("placementSize", c_uint), + ("count", c_uint), + ("placementIds", POINTER(c_uint)), + ("mode", c_uint), + ] + + +VgpuPlacementList_v2 = 0x2000020 + + +class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("bar1Size", c_ulonglong), + ] + + +VgpuTypeBar1Info_v1 = 0x1000010 + + +class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("vgpuInstance", _nvmlVgpuInstance_t), + ("timeStamp", c_ulonglong), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), + ] + + +class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("smUtil", c_nvmlValue_t), + ("memUtil", c_nvmlValue_t), + ("encUtil", c_nvmlValue_t), + ("decUtil", c_nvmlValue_t), + ("jpgUtil", c_nvmlValue_t), + ("ofaUtil", c_nvmlValue_t), + ] + + +class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("sampleValType", _nvmlValueType_t), + ("vgpuInstanceCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuUtilArray", POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), + ] + + +VgpuInstancesUtilizationInfo_v1 = 0x01000020 + + +class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ] + + +class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ("timeStamp", c_ulonglong), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), + ] + + +class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("vgpuProcessCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("vgpuProcUtilArray", POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), + ] + + +VgpuProcessesUtilizationInfo_v1 = 0x01000018 + + +class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("size", c_ulonglong), + ] + + +VgpuRuntimeState_v1 = 0x1000010 + + +class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), + ] + + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 +NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + + +class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): + _fields_ = [ + ("isLicensed", c_uint8), + ("licenseExpiry", c_nvmlVgpuLicenseExpiry_t), + ("currentState", c_uint), + ] + + +class c_nvmlEncoderSession_t(_PrintableStructure): + _fields_ = [ + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("codecType", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFps", c_uint), + ("encodeLatency", c_uint), + ] + + +class c_nvmlProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ("pid", c_uint), + ("timeStamp", c_ulonglong), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ] + + +class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("timeStamp", c_ulonglong), + ("pid", c_uint), + ("smUtil", c_uint), + ("memUtil", c_uint), + ("encUtil", c_uint), + ("decUtil", c_uint), + ("jpgUtil", c_uint), + ("ofaUtil", c_uint), + ] + + +class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("processSamplesCount", c_uint), + ("lastSeenTimeStamp", c_ulonglong), + ("procUtilArray", POINTER(c_nvmlProcessUtilizationInfo_v1_t)), + ] + + +ProcessesUtilizationInfo_v1 = 0x01000018 + + +class c_nvmlGridLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ("year", c_uint32), + ("month", c_uint16), + ("day", c_uint16), + ("hour", c_uint16), + ("min", c_uint16), + ("sec", c_uint16), + ("status", c_uint8), + ] + + +class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), + ("licenseExpiry", c_nvmlGridLicenseExpiry_t), + ] + + +class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v4_t + * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("featureEnabled", c_uint), + ] + + +class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v3_t + * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + + +class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_v2_t + * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlGridLicensableFeature_t(_PrintableStructure): + _fields_ = [ + ("featureCode", _nvmlGridLicenseFeatureCode_t), + ("featureState", c_uint), + ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + + +class c_nvmlGridLicensableFeatures_t(_PrintableStructure): + _fields_ = [ + ("isGridLicenseSupported", c_int), + ("licensableFeaturesCount", c_uint), + ( + "gridLicensableFeatures", + c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, + ), + ] + + +class c_nvmlMarginTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("marginTemperature", c_int), + ] + + +nvmlMarginTemperature_v1 = 0x1000008 + + +## Event structures +class struct_c_nvmlEventSet_t(Structure): + pass # opaque handle + + +c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) + +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = ( + nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction +) + +## Clock Event Reasons defines +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = ( + nvmlClocksEventReasonNone + | nvmlClocksEventReasonGpuIdle + | nvmlClocksEventReasonApplicationsClocksSetting + | nvmlClocksEventReasonSwPowerCap + | nvmlClocksEventReasonHwSlowdown + | nvmlClocksEventReasonSyncBoost + | nvmlClocksEventReasonSwThermalSlowdown + | nvmlClocksEventReasonHwThermalSlowdown + | nvmlClocksEventReasonHwPowerBrakeSlowdown + | nvmlClocksEventReasonDisplayClockSetting +) + +## Following have been deprecated +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = ( + nvmlClocksThrottleReasonNone + | nvmlClocksThrottleReasonGpuIdle + | nvmlClocksThrottleReasonApplicationsClocksSetting + | nvmlClocksThrottleReasonSwPowerCap + | nvmlClocksThrottleReasonHwSlowdown + | nvmlClocksThrottleReasonSyncBoost + | nvmlClocksThrottleReasonSwThermalSlowdown + | nvmlClocksThrottleReasonHwThermalSlowdown + | nvmlClocksThrottleReasonHwPowerBrakeSlowdown + | nvmlClocksThrottleReasonDisplayClockSetting +) + + +class c_nvmlEventData_t(_PrintableStructure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("eventType", c_ulonglong), + ("eventData", c_ulonglong), + ("gpuInstanceId", c_uint), + ("computeInstanceId", c_uint), + ] + _fmt_ = {"eventType": "0x%08X"} + + +class c_nvmlAccountingStats_t(_PrintableStructure): + _fields_ = [ + ("gpuUtilization", c_uint), + ("memoryUtilization", c_uint), + ("maxMemoryUsage", c_ulonglong), + ("time", c_ulonglong), + ("startTime", c_ulonglong), + ("isRunning", c_uint), + ("reserved", c_uint * 5), + ] + + +class c_nvmlVgpuVersion_t(Structure): + _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)] + + +class c_nvmlVgpuMetadata_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("guestInfoState", _nvmlVgpuGuestInfoState_t), + ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), + ("guestVgpuVersion", c_uint), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE), + ] + + +class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("revision", c_uint), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("pgpuVirtualizationCaps", c_uint), + ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE), + ] + + +class c_nvmlVgpuPgpuCompatibility_t(Structure): + _fields_ = [ + ("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t), + ] + + +## vGPU scheduler policy defines +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 + +## Supported vGPU scheduler policy count +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 + +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 + + +class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): + _fields_ = [ + ("avgFactor", c_uint), + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedData_t(_PrintableStructure): + _fields_ = [ + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedulerParams_t(Union): + _fields_ = [ + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedData_t), + ] + + +class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): + _fields_ = [ + ("timestamp", c_ulonglong), + ("timeRunTotal", c_ulonglong), + ("timeRun", c_ulonglong), + ("swRunlistId", c_uint), + ("targetTimeSlice", c_ulonglong), + ("cumulativePreemptionTime", c_ulonglong), + ] + + +class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): + _fields_ = [ + ("engineId", c_uint), + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), + ("entriesCount", c_uint), + ( + "logEntries", + c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES, + ), + ] + + +class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): + _fields_ = [ + ("schedulerPolicy", c_uint), + ("arrMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerParams_t), + ] + + +class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): + _fields_ = [ + ("avgFactor", c_uint), + ("frequency", c_uint), + ] + + +class c_nvmlVgpuSchedSetData_t(_PrintableStructure): + _fields_ = [ + ("timeslice", c_uint), + ] + + +class c_nvmlVgpuSchedulerSetParams_t(Union): + _fields_ = [ + ("vgpuSchedDataWithARR", c_nvmlVgpuSchedSetDataWithARR_t), + ("vgpuSchedData", c_nvmlVgpuSchedSetData_t), + ] + + +class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): + _fields_ = [ + ("schedulerPolicy", c_uint), + ("enableARRMode", c_uint), + ("schedulerParams", c_nvmlVgpuSchedulerSetParams_t), + ] + + +class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): + _fields_ = [ + ( + "supportedSchedulers", + c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT, + ), + ("maxTimeslice", c_uint), + ("minTimeslice", c_uint), + ("isArrModeSupported", c_uint), + ("maxFrequencyForARR", c_uint), + ("minFrequencyForARR", c_uint), + ("maxAvgFactorForARR", c_uint), + ("minAvgFactorForARR", c_uint), + ] + + +class c_nvmlFBCStats_t(Structure): + _fields_ = [ + ("sessionsCount", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), + ] + + +class c_nvmlFBCSession_t(_PrintableStructure): + _fields_ = [ + ("sessionId", c_uint), + ("pid", c_uint), + ("vgpuInstance", _nvmlVgpuInstance_t), + ("displayOrdinal", c_uint), + ("sessionType", c_uint), + ("sessionFlags", c_uint), + ("hMaxResolution", c_uint), + ("vMaxResolution", c_uint), + ("hResolution", c_uint), + ("vResolution", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint), + ] + + +NVML_DEVICE_MIG_DISABLE = 0x0 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD + + +class c_nvmlGpuInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), ("size", c_uint)] + + +class c_nvmlGpuInstanceProfileInfo_t(Structure): + _fields_ = [ + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + + +nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + + +class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] + + def __init__(self): + super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__( + version=nvmlGpuInstanceProfileInfo_v2 + ) + + +class c_nvmlGpuInstanceInfo_t(Structure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t), + ] + + +class struct_c_nvmlGpuInstance_t(Structure): + pass # opaque handle + + +c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) + +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 + +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + + +class c_nvmlComputeInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), ("size", c_uint)] + + +class c_nvmlComputeInstanceProfileInfo_t(Structure): + _fields_ = [ + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ] + + +nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + + +class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), + ] + + def __init__(self): + super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__( + version=nvmlComputeInstanceProfileInfo_v2 + ) + + +class c_nvmlComputeInstanceInfo_t(Structure): + _fields_ = [ + ("device", c_nvmlDevice_t), + ("gpuInstance", c_nvmlGpuInstance_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t), + ] + + +NVML_MAX_GPU_UTILIZATIONS = 8 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 + + +class c_nvmlGpuDynamicPstatesUtilization_t(Structure): + _fields_ = [ + ("bIsPresent", c_uint, 1), + ("percentage", c_uint), + ("incThreshold", c_uint), + ("decThreshold", c_uint), + ] + + +class c_nvmlGpuDynamicPstatesInfo_t(Structure): + _fields_ = [ + ("flags", c_uint), + ( + "utilization", + c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS, + ), + ] + + +NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 + +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 + +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + + +class c_nvmlGpuThermalSensor_t(Structure): + _fields_ = [ + ("controller", c_int), + ("defaultMinTemp", c_int), + ("defaultMaxTemp", c_int), + ("currentTemp", c_int), + ("target", c_int), + ] + + +class c_nvmlGpuThermalSettings_t(Structure): + _fields_ = [ + ("count", c_uint), + ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU), + ] + + +_nvmlCoolerControl_t = c_uint +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 + +_nvmlCoolerTarget_t = c_uint +NVML_THERMAL_COOLER_TARGET_NONE = 1 << 0 +NVML_THERMAL_COOLER_TARGET_GPU = 1 << 1 +NVML_THERMAL_COOLER_TARGET_MEMORY = 1 << 2 +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = 1 << 3 +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = ( + NVML_THERMAL_COOLER_TARGET_GPU + | NVML_THERMAL_COOLER_TARGET_MEMORY + | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY +) + + +class c_nvmlCoolerInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("index", c_uint), + ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t), + ] + + +nvmlCoolerInfo_v1 = 0x1000010 + + +def nvmlDeviceGetCoolerInfo(handle): + c_coolerInfo = c_nvmlCoolerInfo_t() + c_coolerInfo.version = nvmlCoolerInfo_v1 + c_coolerInfo.index = 0 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCoolerInfo") + ret = fn(handle, byref(c_coolerInfo)) + _nvmlCheckReturn(ret) + return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + + +class struct_c_nvmlComputeInstance_t(Structure): + pass # opaque handle + + +c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + + +class c_nvmlDeviceAttributes(Structure): + _fields_ = [ + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + + +class c_nvmlRowRemapperHistogramValues(Structure): + _fields_ = [ + ("max", c_uint), + ("high", c_uint), + ("partial", c_uint), + ("low", c_uint), + ("none", c_uint), + ] + + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 +NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 +NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + + +class c_nvmlConfComputeSystemState_t(Structure): + _fields_ = [ + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ] + + +nvmlSystemConfComputeSettings_v1 = 0x1000014 + + +class c_nvmlSystemConfComputeSettings_v1_t(Structure): + _fields_ = [ + ("version", c_uint), + ("environment", c_uint), + ("ccFeature", c_uint), + ("devToolsMode", c_uint), + ("multiGpuMode", c_uint), + ] + + def __init__(self): + super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__( + version=nvmlSystemConfComputeSettings_v1 + ) + + +class c_nvmlConfComputeSystemCaps_t(Structure): + _fields_ = [ + ("cpuCaps", c_uint), + ("gpusCaps", c_uint), + ] + + +class c_nvmlConfComputeMemSizeInfo_t(Structure): + _fields_ = [ + ("protectedMemSizeKib", c_ulonglong), + ("unprotectedMemSizeKib", c_ulonglong), + ] + + +class c_nvmlConfComputeGpuCertificate_t(Structure): + _fields_ = [ + ("certChainSize", c_uint), + ("attestationCertChainSize", c_uint), + ("certChain", c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ( + "attestationCertChain", + c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE, + ), + ] + + +class c_nvmlConfComputeGpuAttestationReport_t(Structure): + _fields_ = [ + ("isCecAttestationReportPresent", c_uint), + ("attestationReportSize", c_uint), + ("cecAttestationReportSize", c_uint), + ("nonce", c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), + ("attestationReport", c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), + ( + "cecAttestationReport", + c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE, + ), + ] + + +class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): + _fields_ = [ + ("version", c_uint), + ("maxAttackerAdvantage", c_ulong), + ] + + +ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + + +class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): + _fields_ = [ + ("version", c_uint), + ("attackerAdvantage", c_ulong), + ] + + +ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 + + +## string/bytes conversion for ease of use +def convertStrBytes(func): + """ + In python 3, strings are unicode instead of bytes, and need to be converted for ctypes + Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) + Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> + ---- + Returned from function: b'returned string' + Returned to caller: 'returned string' + """ + + @wraps(func) + def wrapper(*args, **kwargs): + # encoding a str returns bytes in python 2 and 3 + args = [arg.encode() if isinstance(arg, str) else arg for arg in args] + res = func(*args, **kwargs) + # In python 2, str and bytes are the same + # In python 3, str is unicode and should be decoded. + # Ctypes handles most conversions, this only effects c_char and char arrays. + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + if sys.version_info >= (3,): + return wrapper + return func + + +def throwOnVersionMismatch(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except NVMLError_FunctionNotFound: + raise NVMLLibraryMismatchError( + "Unversioned function called and the " + "pyNVML version does not match the NVML lib version. " + "Either use matching pyNVML and NVML lib versions or " + "use a versioned function such as " + func.__name__ + "_v2" + ) + + return wrapper + + +## C function wrappers ## +def nvmlInitWithFlags(flags): + _LoadNvmlLibrary() + + # + # Initialize the library + # + fn = _nvmlGetFunctionPointer("nvmlInitWithFlags") + ret = fn(flags) + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + _nvmlLib_refcount += 1 + libLoadLock.release() + return None + + +def nvmlInit(): + nvmlInitWithFlags(0) + return None + + +def _LoadNvmlLibrary(): + """Load the library if it isn't loaded already.""" + global nvmlLib + + if nvmlLib == None: + # lock to ensure only one caller loads the library + libLoadLock.acquire() + + try: + # ensure the library still isn't loaded + if nvmlLib == None: + try: + if sys.platform[:3] == "win": + # cdecl calling convention + try: + # Check for nvml.dll in System32 first for DCH drivers + nvmlLib = CDLL( + os.path.join( + os.getenv("WINDIR", "C:/Windows"), + "System32/nvml.dll", + ) + ) + except OSError as ose: + # If nvml.dll is not found in System32, it should be in ProgramFiles + # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll + nvmlLib = CDLL( + os.path.join( + os.getenv( + "ProgramFiles", "C:/Program Files" + ), + "NVIDIA Corporation/NVSMI/nvml.dll", + ) + ) + else: + # assume linux + nvmlLib = CDLL("libnvidia-ml.so.1") + except OSError as ose: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + if nvmlLib == None: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + + +def nvmlShutdown(): + # + # Leave the library loaded, but shutdown the interface + # + fn = _nvmlGetFunctionPointer("nvmlShutdown") + ret = fn() + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + if 0 < _nvmlLib_refcount: + _nvmlLib_refcount -= 1 + libLoadLock.release() + return None + + +# Added in 2.285 +@convertStrBytes +def nvmlErrorString(result): + fn = _nvmlGetFunctionPointer("nvmlErrorString") + fn.restype = c_char_p # otherwise return is an int + ret = fn(result) + return ret + + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetNVMLVersion(): + c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +def nvmlSystemGetCudaDriverVersion(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + + +def nvmlSystemGetCudaDriverVersion_v2(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetProcessName(pid): + c_name = create_string_buffer(1024) + fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") + ret = fn(c_uint(pid), c_name, c_uint(1024)) + _nvmlCheckReturn(ret) + return c_name.value + + +@convertStrBytes +def nvmlSystemGetDriverVersion(): + c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 2.285 +def nvmlSystemGetHicVersion(): + c_count = c_uint(0) + hics = None + fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") + + # get the count + ret = fn(byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # If there are no hics + if c_count.value == 0: + return [] + + hic_array = c_nvmlHwbcEntry_t * c_count.value + hics = hic_array() + ret = fn(byref(c_count), hics) + _nvmlCheckReturn(ret) + return hics + + +def nvmlSystemGetDriverBranch(): + c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) + c_branchInfo.version = SystemDriverBranchInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + ret = fn( + byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + ) + _nvmlCheckReturn(ret) + return c_branchInfo + + +## Unit get functions +def nvmlUnitGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlUnitGetHandleByIndex(index): + c_index = c_uint(index) + unit = c_nvmlUnit_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") + ret = fn(c_index, byref(unit)) + _nvmlCheckReturn(ret) + return unit + + +def nvmlUnitGetUnitInfo(unit): + c_info = c_nvmlUnitInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlUnitGetLedState(unit): + c_state = c_nvmlLedState_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") + ret = fn(unit, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + + +def nvmlUnitGetPsuInfo(unit): + c_info = c_nvmlPSUInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlUnitGetTemperature(unit, type): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") + ret = fn(unit, c_uint(type), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlUnitGetFanSpeedInfo(unit): + c_speeds = c_nvmlUnitFanSpeeds_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") + ret = fn(unit, byref(c_speeds)) + _nvmlCheckReturn(ret) + return c_speeds + + +# added to API +def nvmlUnitGetDeviceCount(unit): + c_count = c_uint(0) + # query the unit to determine device count + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), None) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = NVML_SUCCESS + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlUnitGetDevices(unit): + c_count = c_uint(nvmlUnitGetDeviceCount(unit)) + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return c_devices + + +## Device get functions +def nvmlDeviceGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetHandleByIndex(index): + c_index = c_uint(index) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") + ret = fn(c_index, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleBySerial(serial): + c_serial = c_char_p(serial) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") + ret = fn(c_serial, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleByUUID(uuid): + c_uuid = c_char_p(uuid) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") + ret = fn(c_uuid, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetHandleByPciBusId(pciBusId): + c_busId = c_char_p(pciBusId) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") + ret = fn(c_busId, byref(device)) + _nvmlCheckReturn(ret) + return device + + +@convertStrBytes +def nvmlDeviceGetName(handle): + c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") + ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_name.value + + +class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + + +nvmlDevicePerfModes_v1 = 0x1000804 + + +@convertStrBytes +def nvmlDeviceGetPerformanceModes(handle): + perfModes = c_nvmlDevicePerfModes_v1_t() + perfModes.version = nvmlDevicePerfModes_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceModes") + ret = fn(handle, byref(perfModes)) + _nvmlCheckReturn(ret) + return perfModes.str + + +class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + + +nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + + +@convertStrBytes +def nvmlDeviceGetCurrentClockFreqs(handle): + currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() + currentClockFreqs.version = nvmlDeviceCurrentClockFreqs_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClockFreqs") + ret = fn(handle, byref(currentClockFreqs)) + _nvmlCheckReturn(ret) + return currentClockFreqs.str + + +def nvmlDeviceGetBoardId(handle): + c_id = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") + ret = fn(handle, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + + +def nvmlDeviceGetMultiGpuBoard(handle): + c_multiGpu = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") + ret = fn(handle, byref(c_multiGpu)) + _nvmlCheckReturn(ret) + return c_multiGpu.value + + +def nvmlDeviceGetBrand(handle): + c_type = _nvmlBrandType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") + ret = fn(handle, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + + +def nvmlDeviceGetC2cModeInfoV1(handle): + c_info = c_nvmlC2cModeInfo_v1_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlDeviceGetC2cModeInfoV(handle): + return nvmlDeviceGetC2cModeInfoV1(handle) + + +@convertStrBytes +def nvmlDeviceGetBoardPartNumber(handle): + c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardPartNumber") + ret = fn(handle, c_part_number, c_uint(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_part_number.value + + +@convertStrBytes +def nvmlDeviceGetSerial(handle): + c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") + ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_serial.value + + +def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): + isReference = type(moduleId) is not c_uint + moduleIdRef = moduleId if isReference else byref(moduleId) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetModuleId") + ret = fn(handle, moduleIdRef) + if isReference: + return ret + else: + _nvmlCheckReturn(ret) + return moduleId.value + + +def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): + affinity_array = c_ulonglong * nodeSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryAffinity") + ret = fn( + handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope) + ) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinityWithinScope") + ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") + ret = fn(handle, cpuSetSize, byref(c_affinity)) + _nvmlCheckReturn(ret) + return c_affinity + + +def nvmlDeviceSetCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNumaNodeId(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") + node = c_int() + ret = fn(handle, byref(node)) + _nvmlCheckReturn(ret) + return node.value + + +def nvmlDeviceGetMinorNumber(handle): + c_minor_number = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") + ret = fn(handle, byref(c_minor_number)) + _nvmlCheckReturn(ret) + return c_minor_number.value + + +@convertStrBytes +def nvmlDeviceGetUUID(handle): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") + ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlDeviceGetInforomVersion(handle, infoRomObject): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") + ret = fn( + handle, + _nvmlInforomObject_t(infoRomObject), + c_version, + c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE), + ) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 4.304 +@convertStrBytes +def nvmlDeviceGetInforomImageVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 4.304 +def nvmlDeviceGetInforomConfigurationChecksum(handle): + c_checksum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") + ret = fn(handle, byref(c_checksum)) + _nvmlCheckReturn(ret) + return c_checksum.value + + +# Added in 4.304 +def nvmlDeviceValidateInforom(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetLastBBXFlushTime(handle): + c_timestamp = c_ulonglong() + c_durationUs = c_ulong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetLastBBXFlushTime") + ret = fn(handle, byref(c_timestamp), byref(c_durationUs)) + _nvmlCheckReturn(ret) + return [c_timestamp.value, c_durationUs.value] + + +def nvmlDeviceGetDisplayMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetDisplayActive(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetPersistenceMode(handle): + c_state = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") + ret = fn(handle, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +def nvmlDeviceGetPciInfoExt(handle, c_info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") + ret = fn(handle, c_info) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetPciInfo_v3(handle): + c_info = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlDeviceGetPciInfo(handle): + return nvmlDeviceGetPciInfo_v3(handle) + + +def nvmlDeviceGetClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 2.285 +def nvmlDeviceGetMaxClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 4.304 +def nvmlDeviceGetApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +def nvmlDeviceGetMaxCustomerBoostClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +def nvmlDeviceGetClock(handle, type, id): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") + ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 5.319 +def nvmlDeviceGetDefaultApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + + +# Added in 4.304 +def nvmlDeviceGetSupportedMemoryClocks(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no clocks + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + + +# Added in 4.304 +def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no clocks + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetFanSpeed(handle): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") + ret = fn(handle, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetFanSpeed_v2(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +class c_nvmlFanSpeedInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("fan", c_uint), + ("speed", c_uint), + ] + + +nvmlFanSpeedInfo_v1 = 0x100000C + + +def nvmlDeviceGetFanSpeedRPM(handle): + c_fanSpeed = c_nvmlFanSpeedInfo_t() + c_fanSpeed.fan = 0 + c_fanSpeed.version = nvmlFanSpeedInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeedRPM") + ret = fn(handle, byref(c_fanSpeed)) + _nvmlCheckReturn(ret) + return c_fanSpeed.speed + + +def nvmlDeviceGetTargetFanSpeed(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetNumFans(device): + c_numFans = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") + ret = fn(device, byref(c_numFans)) + _nvmlCheckReturn(ret) + return c_numFans.value + + +def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2") + ret = fn(handle, index) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): + isReference = (type(minSpeed) is not c_uint) or ( + type(maxSpeed) is not c_uint + ) + minSpeedRef = minSpeed if isReference else byref(minSpeed) + maxSpeedRef = maxSpeed if isReference else byref(maxSpeed) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxFanSpeed") + ret = fn(handle, minSpeedRef, maxSpeedRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + + +def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): + isReference = type(fanControlPolicy) is not c_uint + fanControlPolicyRef = ( + fanControlPolicy if isReference else byref(fanControlPolicy) + ) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanControlPolicy_v2") + ret = fn(handle, fan, fanControlPolicyRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else fanControlPolicy.value + + +def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") + ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +class c_nvmlTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("sensorType", _nvmlTemperatureSensors_t), + ("temperature", c_int), + ] + + +nvmlTemperature_v1 = 0x100000C + + +def nvmlDeviceGetTemperatureV1(handle, sensor): + c_temp = c_nvmlTemperature_v1_t() + c_temp.version = nvmlTemperature_v1 + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") + ret = fn(handle, byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.temperature + + +def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): + if version == nvmlTemperature_v1: + return nvmlDeviceGetTemperatureV1(handle, sensor) + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + +# DEPRECATED use nvmlDeviceGetTemperatureV instead +def nvmlDeviceGetTemperature(handle, sensor): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") + ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlDeviceGetTemperatureThreshold(handle, threshold): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + + +def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): + c_temp = c_uint() + c_temp.value = temp + fn = _nvmlGetFunctionPointer("nvmlDeviceSetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetMarginTemperature(handle): + c_marginTempInfo = c_nvmlMarginTemperature_v1_t() + c_marginTempInfo.version = nvmlMarginTemperature_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMarginTemperature") + ret = fn(handle, byref(c_marginTempInfo)) + _nvmlCheckReturn(ret) + return c_marginTempInfo.marginTemperature + + +# DEPRECATED use nvmlDeviceGetPerformanceState +def nvmlDeviceGetPowerState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + + +def nvmlDeviceGetPerformanceState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + + +def nvmlDeviceGetPowerManagementMode(handle): + c_pcapMode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") + ret = fn(handle, byref(c_pcapMode)) + _nvmlCheckReturn(ret) + return c_pcapMode.value + + +def nvmlDeviceGetPowerManagementLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 4.304 +def nvmlDeviceGetPowerManagementLimitConstraints(handle): + c_minLimit = c_uint() + c_maxLimit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") + ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) + _nvmlCheckReturn(ret) + return [c_minLimit.value, c_maxLimit.value] + + +# Added in 4.304 +def nvmlDeviceGetPowerManagementDefaultLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 331 +def nvmlDeviceGetEnforcedPowerLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +def nvmlDeviceGetPowerUsage(handle): + c_watts = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") + ret = fn(handle, byref(c_watts)) + _nvmlCheckReturn(ret) + return c_watts.value + + +def nvmlDeviceGetTotalEnergyConsumption(handle): + c_millijoules = c_uint64() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") + ret = fn(handle, byref(c_millijoules)) + _nvmlCheckReturn(ret) + return c_millijoules.value + + +# Added in 4.304 +def nvmlDeviceGetGpuOperationMode(handle): + c_currState = _nvmlGpuOperationMode_t() + c_pendingState = _nvmlGpuOperationMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + + +# Added in 4.304 +def nvmlDeviceGetCurrentGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[0] + + +# Added in 4.304 +def nvmlDeviceGetPendingGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[1] + + +def nvmlDeviceGetMemoryInfo(handle, version=None): + if not version: + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") + else: + c_memory = c_nvmlMemory_v2_t() + c_memory.version = version + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo_v2") + ret = fn(handle, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + + +def nvmlDeviceGetBAR1MemoryInfo(handle): + c_bar1_memory = c_nvmlBAR1Memory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") + ret = fn(handle, byref(c_bar1_memory)) + _nvmlCheckReturn(ret) + return c_bar1_memory + + +def nvmlDeviceGetComputeMode(handle): + c_mode = _nvmlComputeMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetCudaComputeCapability(handle): + c_major = c_int() + c_minor = c_int() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") + ret = fn(handle, byref(c_major), byref(c_minor)) + _nvmlCheckReturn(ret) + return (c_major.value, c_minor.value) + + +def nvmlDeviceGetEccMode(handle): + c_currState = _nvmlEnableState_t() + c_pendingState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + + +# added to API +def nvmlDeviceGetCurrentEccMode(handle): + return nvmlDeviceGetEccMode(handle)[0] + + +# added to API +def nvmlDeviceGetPendingEccMode(handle): + return nvmlDeviceGetEccMode(handle)[1] + + +def nvmlDeviceGetDefaultEccMode(handle): + c_defaultState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") + ret = fn(handle, byref(c_defaultState)) + _nvmlCheckReturn(ret) + return [c_defaultState.value] + + +def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_count), + ) + _nvmlCheckReturn(ret) + return c_count.value + + +# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter +def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): + c_counts = c_nvmlEccErrorCounts_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + byref(c_counts), + ) + _nvmlCheckReturn(ret) + return c_counts + + +# Added in 4.304 +def nvmlDeviceGetMemoryErrorCounter( + handle, errorType, counterType, locationType +): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") + ret = fn( + handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), + byref(c_count), + ) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetUtilizationRates(handle): + c_util = c_nvmlUtilization_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") + ret = fn(handle, byref(c_util)) + _nvmlCheckReturn(ret) + return c_util + + +def nvmlDeviceGetEncoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetDecoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetJpgUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetJpgUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetOfaUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetOfaUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + + +def nvmlDeviceGetPcieReplayCounter(handle): + c_replay = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") + ret = fn(handle, byref(c_replay)) + _nvmlCheckReturn(ret) + return c_replay.value + + +def nvmlDeviceGetDriverModel(handle): + c_currModel = _nvmlDriverModel_t() + c_pendingModel = _nvmlDriverModel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") + ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) + _nvmlCheckReturn(ret) + return [c_currModel.value, c_pendingModel.value] + + +# added to API +def nvmlDeviceGetCurrentDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[0] + + +# added to API +def nvmlDeviceGetPendingDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[1] + + +# Added in 2.285 +@convertStrBytes +def nvmlDeviceGetVbiosVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +@throwOnVersionMismatch +def nvmlDeviceGetComputeRunningProcesses(handle): + return nvmlDeviceGetComputeRunningProcesses_v3(handle) + + +def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +@throwOnVersionMismatch +def nvmlDeviceGetGraphicsRunningProcesses(handle): + return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + + +@throwOnVersionMismatch +def nvmlDeviceGetMPSComputeRunningProcesses(handle): + return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + + +def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): + c_processDetailList = c_nvmlProcessDetailList_t() + c_processDetailList.version = version + c_processDetailList.mode = mode + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRunningProcessDetailList") + + # first call to get the size + ret = fn(handle, byref(c_processDetailList)) + if ret == NVML_SUCCESS: + # special case, no running processes + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + c_procs = ( + c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries + ) + c_processDetailList.procArray = cast( + (c_procs)(), POINTER(c_nvmlProcessDetail_v1_t) + ) + + # make the call again + ret = fn(handle, byref(c_processDetailList)) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_processDetailList.numProcArrayEntries): + # use an alternative struct for this object + obj = c_processDetailList.procArray[i] + if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + obj.usedGpuMemory = None + if ( + obj.usedGpuCcProtectedMemory + == NVML_VALUE_NOT_AVAILABLE_ulonglong.value + ): + obj.usedGpuCcProtectedMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetAutoBoostedClocksEnabled(handle): + c_isEnabled = _nvmlEnableState_t() + c_defaultIsEnabled = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") + ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) + _nvmlCheckReturn(ret) + return [c_isEnabled.value, c_defaultIsEnabled.value] + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +## Set functions +def nvmlUnitSetLedState(unit, color): + fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") + ret = fn(unit, _nvmlLedColor_t(color)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetPersistenceMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetComputeMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") + ret = fn(handle, _nvmlComputeMode_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetEccMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearEccErrorCounts(handle, counterType): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") + ret = fn(handle, _nvmlEccCounterType_t(counterType)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetDriverModel(handle, model): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") + ret = fn(handle, _nvmlDriverModel_t(model)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled)) + _nvmlCheckReturn(ret) + return None + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) + _nvmlCheckReturn(ret) + return None + # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + + +def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") + ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetGpuLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") + ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetMemoryLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): + isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t + c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClkMonStatus") + ret = fn(handle, c_clkMonInfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_clkMonInfo + + +# Added in 4.304 +def nvmlDeviceSetApplicationsClocks( + handle, maxMemClockMHz, maxGraphicsClockMHz +): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") + ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceResetApplicationsClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceSetPowerManagementLimit(handle, limit): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") + ret = fn(handle, c_uint(limit)) + _nvmlCheckReturn(ret) + return None + + +# Added in 4.304 +def nvmlDeviceSetGpuOperationMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") + ret = fn(handle, _nvmlGpuOperationMode_t(mode)) + _nvmlCheckReturn(ret) + return None + + +# Added in 2.285 +def nvmlEventSetCreate(): + fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") + eventSet = c_nvmlEventSet_t() + ret = fn(byref(eventSet)) + _nvmlCheckReturn(ret) + return eventSet + + +# Added in 2.285 +def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): + fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") + ret = fn(handle, c_ulonglong(eventTypes), eventSet) + _nvmlCheckReturn(ret) + return None + + +# Added in 2.285 +def nvmlDeviceGetSupportedEventTypes(handle): + c_eventTypes = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") + ret = fn(handle, byref(c_eventTypes)) + _nvmlCheckReturn(ret) + return c_eventTypes.value + + +# raises NVML_ERROR_TIMEOUT exception on timeout +def nvmlEventSetWait_v2(eventSet, timeoutms): + fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") + data = c_nvmlEventData_t() + ret = fn(eventSet, byref(data), c_uint(timeoutms)) + _nvmlCheckReturn(ret) + return data + + +def nvmlEventSetWait(eventSet, timeoutms): + return nvmlEventSetWait_v2(eventSet, timeoutms) + + +# Added in 2.285 +def nvmlEventSetFree(eventSet): + fn = _nvmlGetFunctionPointer("nvmlEventSetFree") + ret = fn(eventSet) + _nvmlCheckReturn(ret) + return None + + +# Added in 3.295 +def nvmlDeviceOnSameBoard(handle1, handle2): + fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") + onSameBoard = c_int() + ret = fn(handle1, handle2, byref(onSameBoard)) + _nvmlCheckReturn(ret) + return onSameBoard.value != 0 + + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + + +def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + + +# Added in 4.304 +def nvmlDeviceGetSupportedClocksThrottleReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +def nvmlDeviceGetSupportedClocksEventReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +# Added in 4.304 +def nvmlDeviceGetCurrentClocksThrottleReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +def nvmlDeviceGetCurrentClocksEventReasons(handle): + c_reasons = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + + +# Added in 5.319 +def nvmlDeviceGetIndex(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") + c_index = c_uint() + ret = fn(handle, byref(c_index)) + _nvmlCheckReturn(ret) + return c_index.value + + +# Added in 5.319 +def nvmlDeviceGetAccountingMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceSetAccountingMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceClearAccountingPids(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetAccountingStats(handle, pid): + stats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") + ret = fn(handle, c_uint(pid), byref(stats)) + _nvmlCheckReturn(ret) + if stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: + # special case for WDDM on Windows, see comment above + stats.maxMemoryUsage = None + return stats + + +def nvmlDeviceGetAccountingPids(handle): + count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) + pids = (c_uint * count.value)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") + ret = fn(handle, byref(count), pids) + _nvmlCheckReturn(ret) + return list(map(int, pids[0 : count.value])) + + +def nvmlDeviceGetAccountingBufferSize(handle): + bufferSize = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") + ret = fn(handle, byref(bufferSize)) + _nvmlCheckReturn(ret) + return int(bufferSize.value) + + +def nvmlDeviceGetRetiredPages(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + ret = fn(device, c_source, byref(c_count), c_pages) + _nvmlCheckReturn(ret) + return list(map(int, c_pages[0 : c_count.value])) + + +def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages_v2") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + times_array = c_ulonglong * c_count.value + c_times = times_array() + ret = fn(device, c_source, byref(c_count), c_pages, c_times) + _nvmlCheckReturn(ret) + return [ + {"address": int(c_pages[i]), "timestamp": int(c_times[i])} + for i in range(c_count.value) + ] + + +def nvmlDeviceGetRetiredPagesPendingStatus(device): + c_pending = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") + ret = fn(device, byref(c_pending)) + _nvmlCheckReturn(ret) + return int(c_pending.value) + + +def nvmlDeviceGetAPIRestriction(device, apiType): + c_permission = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") + ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) + _nvmlCheckReturn(ret) + return int(c_permission.value) + + +def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") + ret = fn( + handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted) + ) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetBridgeChipInfo(handle): + bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") + ret = fn(handle, byref(bridgeHierarchy)) + _nvmlCheckReturn(ret) + return bridgeHierarchy + + +def nvmlDeviceGetSamples(device, sampling_type, timeStamp): + c_sampling_type = _nvmlSamplingType_t(sampling_type) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_count = c_uint(0) + c_sample_value_type = _nvmlValueType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") + + ## First Call gets the size + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + None, + ) + + # Stop if this fails + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + sampleArray = c_sample_count.value * c_nvmlSample_t + c_samples = sampleArray() + ret = fn( + device, + c_sampling_type, + c_time_stamp, + byref(c_sample_value_type), + byref(c_sample_count), + c_samples, + ) + _nvmlCheckReturn(ret) + return (c_sample_value_type.value, c_samples[0 : c_sample_count.value]) + + +def nvmlDeviceGetViolationStatus(device, perfPolicyType): + c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) + c_violTime = c_nvmlViolationTime_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") + + ## Invoke the method to get violation time + ret = fn(device, c_perfPolicy_type, byref(c_violTime)) + _nvmlCheckReturn(ret) + return c_violTime + + +def nvmlDeviceGetPcieThroughput(device, counter): + c_util = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") + ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) + _nvmlCheckReturn(ret) + return c_util.value + + +def nvmlSystemGetTopologyGpuSet(cpuNumber): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") + + # First call will get the size + ret = fn(cpuNumber, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(cpuNumber, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0 : c_count.value]) + + +def nvmlDeviceGetTopologyNearestGpus(device, level): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") + + # First call will get the size + ret = fn(device, level, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(device, level, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0 : c_count.value]) + + +def nvmlDeviceGetTopologyCommonAncestor(device1, device2): + c_level = _nvmlGpuTopologyLevel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") + ret = fn(device1, device2, byref(c_level)) + _nvmlCheckReturn(ret) + return c_level.value + + +def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): + c_rxcounter = c_ulonglong() + c_txcounter = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationCounter") + ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter)) + _nvmlCheckReturn(ret) + return (c_rxcounter.value, c_txcounter.value) + + +def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): + fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") + ret = fn(device, link, counter, freeze) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") + ret = fn(device, link, counter) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceSetNvLinkUtilizationControl( + device, link, counter, control, reset +): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(control), reset) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): + c_control = nvmlNvLinkUtilizationControl_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(c_control)) + _nvmlCheckReturn(ret) + return c_control + + +def nvmlDeviceGetNvLinkCapability(device, link, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") + ret = fn(device, link, capability, byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): + c_result = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") + ret = fn(device, link, counter, byref(c_result)) + _nvmlCheckReturn(ret) + return c_result.value + + +def nvmlDeviceResetNvLinkErrorCounters(device, link): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") + ret = fn(device, link) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetNvLinkRemotePciInfo(device, link): + c_pci = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") + ret = fn(device, link, byref(c_pci)) + _nvmlCheckReturn(ret) + return c_pci + + +def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): + c_type = _nvmlNvLinkDeviceType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") + ret = fn(handle, link, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + + +def nvmlDeviceGetNvLinkState(device, link): + c_isActive = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") + ret = fn(device, link, byref(c_isActive)) + _nvmlCheckReturn(ret) + return c_isActive.value + + +def nvmlDeviceGetNvLinkVersion(device, link): + c_version = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") + ret = fn(device, link, byref(c_version)) + _nvmlCheckReturn(ret) + return c_version.value + + +def nvmlDeviceModifyDrainState(pciInfo, newState): + fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") + ret = fn(pointer(pciInfo), newState) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceQueryDrainState(pciInfo): + c_newState = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") + ret = fn(pointer(pciInfo), byref(c_newState)) + _nvmlCheckReturn(ret) + return c_newState.value + + +def nvmlDeviceRemoveGpu(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceDiscoverGpus(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + + +def nvmlDeviceGetFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + values[i].fieldId, values[i].scopeId = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + + +def nvmlDeviceClearFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceClearFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + values[i].fieldId, values[i].scopeId = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + + +def nvmlDeviceGetVirtualizationMode(handle): + c_virtualization_mode = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") + ret = fn(handle, byref(c_virtualization_mode)) + _nvmlCheckReturn(ret) + return c_virtualization_mode.value + + +def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") + return fn(handle, virtualization_mode) + + +def nvmlDeviceGetVgpuHeterogeneousMode(handle): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return c_vgpuHeterogeneousMode.mode + + +def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + c_vgpuHeterogeneousMode.mode = heterogeneous_mode + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlVgpuInstanceGetPlacementId(vgpuInstance): + c_placement = c_nvmlVgpuPlacementId_v1_t(0) + c_placement.version = VgpuPlacementId_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetPlacementId") + ret = fn(vgpuInstance, byref(c_placement)) + _nvmlCheckReturn(ret) + return c_placement.placementId + + +def nvmlDeviceGetVgpuTypeSupportedPlacements( + handle, vgpuTypeId, mode=0, version=1 +): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + c_vgpu_placements.mode = mode + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + + +def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + + +def nvmlGetVgpuDriverCapabilities(capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") + ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceGetVgpuCapabilities(handle, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") + ret = fn( + handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult) + ) + _nvmlCheckReturn(ret) + return c_capResult.value + + +def nvmlDeviceSetVgpuCapabilities(handle, capability, state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetSupportedVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no supported vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetCreatableVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no supported vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): + c_profile_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + ret = fn(vgpuTypeId, byref(c_profile_id)) + _nvmlCheckReturn(ret) + return c_profile_id.value + + +@convertStrBytes +def nvmlVgpuTypeGetClass(vgpuTypeId): + c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_class.value + + +@convertStrBytes +def nvmlVgpuTypeGetName(vgpuTypeId): + c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_name.value + + +def nvmlVgpuTypeGetDeviceID(vgpuTypeId): + c_device_id = c_ulonglong(0) + c_subsystem_id = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) + _nvmlCheckReturn(ret) + return (c_device_id.value, c_subsystem_id.value) + + +def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): + c_fb_size = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + ret = fn(vgpuTypeId, byref(c_fb_size)) + _nvmlCheckReturn(ret) + return c_fb_size.value + + +def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): + c_num_heads = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + ret = fn(vgpuTypeId, byref(c_num_heads)) + _nvmlCheckReturn(ret) + return c_num_heads.value + + +def nvmlVgpuTypeGetResolution(vgpuTypeId): + c_xdim = c_uint(0) + c_ydim = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) + _nvmlCheckReturn(ret) + return (c_xdim.value, c_ydim.value) + + +@convertStrBytes +def nvmlVgpuTypeGetLicense(vgpuTypeId): + c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + ret = fn(vgpuTypeId, c_license, c_buffer_size) + _nvmlCheckReturn(ret) + return c_license.value + + +def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): + c_frl_config = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + ret = fn(vgpuTypeId, byref(c_frl_config)) + _nvmlCheckReturn(ret) + return c_frl_config.value + + +def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): + c_gsp_heap = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + ret = fn(vgpuTypeId, byref(c_gsp_heap)) + _nvmlCheckReturn(ret) + return c_gsp_heap.value + + +def nvmlVgpuTypeGetFbReservation(vgpuTypeId): + c_fb_reservation = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + ret = fn(vgpuTypeId, byref(c_fb_reservation)) + _nvmlCheckReturn(ret) + return c_fb_reservation.value + + +def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): + c_runtime_state = nvmlVgpuRuntimeState_v1_t() + c_runtime_state.version = VgpuRuntimeState_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + ret = fn(vgpuInstance, byref(c_runtime_state)) + _nvmlCheckReturn(ret) + return c_runtime_state + + +def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + return c_max_instances.value + + +def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): + c_max_instances_per_vm = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) + _nvmlCheckReturn(ret) + return c_max_instances_per_vm.value + + +def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): + c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) + c_bar1Info.version = VgpuTypeBar1Info_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + ret = fn(vgpuTypeId, byref(c_bar1Info)) + _nvmlCheckReturn(ret) + return c_bar1Info + + +def nvmlDeviceGetActiveVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value + c_vgpu_instances = vgpu_instance_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_instances[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + + +@convertStrBytes +def nvmlVgpuInstanceGetVmID(vgpuInstance): + c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) + _nvmlCheckReturn(ret) + return (c_vm_id.value, c_vm_id_type.value) + + +@convertStrBytes +def nvmlVgpuInstanceGetUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + + +@convertStrBytes +def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): + c_driver_version = create_string_buffer( + NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE + ) + c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) + _nvmlCheckReturn(ret) + return c_driver_version.value + + +def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): + c_license_status = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + ret = fn(vgpuInstance, byref(c_license_status)) + _nvmlCheckReturn(ret) + return c_license_status.value + + +def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + c_license_info = c_nvmlVgpuLicenseInfo_t() + ret = fn(vgpuInstance, byref(c_license_info)) + _nvmlCheckReturn(ret) + return c_license_info + + +def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): + return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + + +def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): + c_frl = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + ret = fn(vgpuInstance, byref(c_frl)) + _nvmlCheckReturn(ret) + return c_frl.value + + +def nvmlVgpuInstanceGetEccMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlVgpuInstanceGetType(vgpuInstance): + c_vgpu_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + ret = fn(vgpuInstance, byref(c_vgpu_type)) + _nvmlCheckReturn(ret) + return c_vgpu_type.value + + +def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): + c_encoder_capacity = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + ret = fn(vgpuInstance, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + + +def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + return fn(vgpuInstance, encoder_capacity) + + +def nvmlVgpuInstanceGetFbUsage(vgpuInstance): + c_fb_usage = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + ret = fn(vgpuInstance, byref(c_fb_usage)) + _nvmlCheckReturn(ret) + return c_fb_usage.value + + +def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): + c_cap_result = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) + _nvmlCheckReturn(ret) + return c_cap_result.value + + +def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): + c_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + ret = fn(vgpuInstance, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + + +@convertStrBytes +def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): + c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") + ret = fn( + vgpuInstance, + c_vgpuPciId, + byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)), + ) + _nvmlCheckReturn(ret) + return c_vgpuPciId.value + + +def nvmlDeviceGetVgpuUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_value_type = _nvmlValueType_t() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + ret = fn( + handle, + c_time_stamp, + byref(c_sample_value_type), + byref(c_vgpu_count), + None, + ) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn( + handle, + c_time_stamp, + byref(c_sample_value_type), + byref(c_vgpu_count), + c_samples, + ) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0) + c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1 + c_vgpuUtilInfo.sampleValType = _nvmlValueType_t() + c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) + c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + ret = fn(handle, byref(c_vgpuUtilInfo)) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_vgpuUtilInfo.vgpuInstanceCount + * c_nvmlVgpuInstanceUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_vgpuUtilInfo.vgpuUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpuUtilInfo.vgpuInstanceCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): + c_p2pstatus = _nvmlGpuP2PStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") + ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus)) + _nvmlCheckReturn(ret) + return c_p2pstatus.value + + +def nvmlDeviceGetGridLicensableFeatures_v4(handle): + c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") + ret = fn(handle, byref(c_get_grid_licensable_features)) + _nvmlCheckReturn(ret) + + return c_get_grid_licensable_features + + +def nvmlDeviceGetGridLicensableFeatures(handle): + return nvmlDeviceGetGridLicensableFeatures_v4(handle) + + +def nvmlDeviceGetGspFirmwareVersion(handle, version=None): + isUserDefined = version is not None + if not isUserDefined: + version = (c_char * NVML_GSP_FIRMWARE_VERSION_BUF_SIZE)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareVersion") + ret = fn(handle, version) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else version.value + + +def nvmlDeviceGetGspFirmwareMode( + handle, isEnabled=c_uint(), defaultMode=c_uint() +): + isReference = type(isEnabled) is not c_uint + isEnabledRef = isEnabled if isReference else byref(isEnabled) + defaultModeRef = defaultMode if isReference else byref(defaultMode) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareMode") + ret = fn(handle, isEnabledRef, defaultModeRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + + +def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): + c_encoder_capacity = c_ulonglong(0) + c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderCapacity") + ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + + +def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0) + c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1 + c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) + c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + + if ret == NVML_SUCCESS: + # special case, no active vGPUs + return [] + elif ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_vgpuProcUtilInfo.vgpuProcessCount + * c_nvmlVgpuProcessUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_vgpuProcUtilInfo.vgpuProcessCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetEncoderStats(handle): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderStats") + ret = fn( + handle, + byref(c_encoderCount), + byref(c_encodeFps), + byref(c_encoderLatency), + ) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + + +def nvmlDeviceGetEncoderSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + ret = fn(handle, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetFBCStats(handle): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") + ret = fn(handle, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + + +def nvmlDeviceGetFBCSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + ret = fn(handle, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + ret = fn( + vgpuInstance, + byref(c_encoderCount), + byref(c_encodeFps), + byref(c_encoderLatency), + ) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + + +def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetFBCStats(vgpuInstance): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") + ret = fn(vgpuInstance, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + + +def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if ret == NVML_SUCCESS: + if c_session_count.value != 0: + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetProcessUtilization(handle, timeStamp): + # first call to get the size + c_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + ret = fn(handle, None, byref(c_count), c_time_stamp) + + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_samples, byref(c_count), c_time_stamp) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_count.value] + else: + # error case + raise NVMLError(ret) + + +def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0) + c_processesUtilInfo.version = ProcessesUtilizationInfo_v1 + c_processesUtilInfo.processSamplesCount = c_uint(0) + c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + ret = fn(handle, byref(c_processesUtilInfo)) + + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + # typical case + sampleArray = ( + c_processesUtilInfo.processSamplesCount + * c_nvmlProcessUtilizationInfo_v1_t + ) + c_samples = sampleArray() + c_processesUtilInfo.procUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_processesUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0 : c_processesUtilInfo.processSamplesCount] + else: + # error case + raise NVMLError(ret) + + +def nvmlVgpuInstanceGetMetadata(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") + c_vgpuMetadata = c_nvmlVgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuMetadata + + +def nvmlDeviceGetVgpuMetadata(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") + c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuPgpuMetadata + + +def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): + fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") + c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() + ret = fn( + byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility) + ) + _nvmlCheckReturn(ret) + return c_vgpuPgpuCompatibility + + +@convertStrBytes +def nvmlDeviceGetPgpuMetadataString(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") + c_pgpuMetadata = create_string_buffer( + NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE + ) + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pgpuMetadata.value, c_bufferSize.value) + + +def nvmlDeviceGetVgpuSchedulerLog(handle): + c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") + ret = fn(handle, byref(c_vgpu_sched_log)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_log + + +def nvmlDeviceGetVgpuSchedulerState(handle): + c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") + ret = fn(handle, byref(c_vgpu_sched_state)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_state + + +def nvmlDeviceGetVgpuSchedulerCapabilities(handle): + c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") + ret = fn(handle, byref(c_vgpu_sched_caps)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_caps + + +def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") + ret = fn(handle, byref(sched_state)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSetVgpuVersion(vgpuVersion): + fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") + ret = fn(byref(vgpuVersion)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGetVgpuVersion(supported=None, current=None): + isUserDefined = (supported is not None) or (current is not None) + if not isUserDefined: + supported = c_nvmlVgpuVersion_t() + current = c_nvmlVgpuVersion_t() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") + ret = fn(byref(supported), byref(current)) + _nvmlCheckReturn(ret) + return ( + NVML_SUCCESS + if isUserDefined + else [ + (supported.minVersion, supported.maxVersion), + (current.minVersion, current.maxVersion), + ] + ) + + +def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): + c_pidCount = c_uint() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") + ret = fn(vgpuInstance, byref(c_pidCount), None) + if ret == NVML_ERROR_INSUFFICIENT_SIZE: + sampleArray = c_pidCount.value * c_uint + c_pidArray = sampleArray() + ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pidCount, c_pidArray) + + +def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): + c_accountingStats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") + ret = fn(vgpuInstance, pid, byref(c_accountingStats)) + _nvmlCheckReturn(ret) + return c_accountingStats + + +def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") + ret = fn(vgpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGetExcludedDeviceCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlGetExcludedDeviceInfoByIndex(index): + c_index = c_uint(index) + info = c_nvmlExcludedDeviceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceInfoByIndex") + ret = fn(c_index, byref(info)) + _nvmlCheckReturn(ret) + return info + + +def nvmlDeviceGetHostVgpuMode(handle): + c_host_vgpu_mode = _nvmlHostVgpuMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") + ret = fn(handle, byref(c_host_vgpu_mode)) + _nvmlCheckReturn(ret) + return c_host_vgpu_mode.value + + +def nvmlDeviceSetMigMode(device, mode): + c_activationStatus = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") + ret = fn(device, mode, byref(c_activationStatus)) + _nvmlCheckReturn(ret) + return c_activationStatus.value + + +def nvmlDeviceGetMigMode(device): + c_currentMode = c_uint() + c_pendingMode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigMode") + ret = fn(device, byref(c_currentMode), byref(c_pendingMode)) + _nvmlCheckReturn(ret) + return [c_currentMode.value, c_pendingMode.value] + + +def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): + if version == 2: + c_info = c_nvmlGpuInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlGpuInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +# Define function alias for the API exposed by NVML +nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo + + +def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") + ret = fn(device, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetGpuInstancePossiblePlacements( + device, profileId, placementsRef, countRef +): + fn = _nvmlGetFunctionPointer( + "nvmlDeviceGetGpuInstancePossiblePlacements_v2" + ) + ret = fn(device, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceCreateGpuInstance(device, profileId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") + ret = fn(device, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") + ret = fn(device, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceDestroy(gpuInstance): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") + ret = fn(gpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") + ret = fn(device, profileId, gpuInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") + ret = fn(device, gpuInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceGetInfo(gpuInstance): + c_info = c_nvmlGpuInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") + ret = fn(gpuInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlGpuInstanceGetComputeInstanceProfileInfo( + device, profile, engProfile, version=2 +): + if version == 2: + c_info = c_nvmlComputeInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer( + "nvmlGpuInstanceGetComputeInstanceProfileInfoV" + ) + elif version == 1: + c_info = c_nvmlComputeInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer( + "nvmlGpuInstanceGetComputeInstanceProfileInfo" + ) + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, engProfile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +# Define function alias for the API exposed by NVML +nvmlGpuInstanceGetComputeInstanceProfileInfoV = ( + nvmlGpuInstanceGetComputeInstanceProfileInfo +) + + +def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer( + "nvmlGpuInstanceGetComputeInstanceRemainingCapacity" + ) + ret = fn(gpuInstance, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlGpuInstanceGetComputeInstancePossiblePlacements( + gpuInstance, profileId, placementsRef, countRef +): + fn = _nvmlGetFunctionPointer( + "nvmlGpuInstanceGetComputeInstancePossiblePlacements" + ) + ret = fn(gpuInstance, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") + ret = fn(gpuInstance, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlGpuInstanceCreateComputeInstanceWithPlacement( + gpuInstance, profileId, placement +): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer( + "nvmlGpuInstanceCreateComputeInstanceWithPlacement" + ) + ret = fn(gpuInstance, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlComputeInstanceDestroy(computeInstance): + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") + ret = fn(computeInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceGetComputeInstances( + gpuInstance, profileId, computeInstancesRef, countRef +): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") + ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") + ret = fn(gpuInstance, computeInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + + +def nvmlComputeInstanceGetInfo_v2(computeInstance): + c_info = c_nvmlComputeInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") + ret = fn(computeInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + + +def nvmlComputeInstanceGetInfo(computeInstance): + return nvmlComputeInstanceGetInfo_v2(computeInstance) + + +def nvmlDeviceIsMigDeviceHandle(device): + c_isMigDevice = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") + ret = fn(device, byref(c_isMigDevice)) + _nvmlCheckReturn(ret) + return c_isMigDevice + + +def nvmlDeviceGetGpuInstanceId(device): + c_gpuInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") + ret = fn(device, byref(c_gpuInstanceId)) + _nvmlCheckReturn(ret) + return c_gpuInstanceId.value + + +def nvmlDeviceGetComputeInstanceId(device): + c_computeInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") + ret = fn(device, byref(c_computeInstanceId)) + _nvmlCheckReturn(ret) + return c_computeInstanceId.value + + +def nvmlDeviceGetMaxMigDeviceCount(device): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") + ret = fn(device, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + + +def nvmlDeviceGetMigDeviceHandleByIndex(device, index): + c_index = c_uint(index) + migDevice = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigDeviceHandleByIndex") + ret = fn(device, c_index, byref(migDevice)) + _nvmlCheckReturn(ret) + return migDevice + + +def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") + ret = fn(migDevice, byref(device)) + _nvmlCheckReturn(ret) + return device + + +def nvmlDeviceGetAttributes_v2(device): + c_attrs = c_nvmlDeviceAttributes() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") + ret = fn(device, byref(c_attrs)) + _nvmlCheckReturn(ret) + return c_attrs + + +def nvmlDeviceGetAttributes(device): + return nvmlDeviceGetAttributes_v2(device) + + +def nvmlDeviceGetRemappedRows(device): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") + c_corr = c_uint() + c_unc = c_uint() + c_bpending = c_uint() + c_bfailure = c_uint() + ret = fn( + device, + byref(c_corr), + byref(c_unc), + byref(c_bpending), + byref(c_bfailure), + ) + _nvmlCheckReturn(ret) + return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + + +def nvmlDeviceGetRowRemapperHistogram(device): + c_vals = c_nvmlRowRemapperHistogramValues() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") + ret = fn(device, byref(c_vals)) + _nvmlCheckReturn(ret) + return c_vals + + +def nvmlDeviceGetArchitecture(device): + arch = _nvmlDeviceArchitecture_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") + ret = fn(device, byref(arch)) + _nvmlCheckReturn(ret) + return arch.value + + +def nvmlDeviceGetBusType(device): + c_busType = _nvmlBusType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") + ret = fn(device, byref(c_busType)) + _nvmlCheckReturn(ret) + return c_busType.value + + +def nvmlDeviceGetIrqNum(device): + c_irqNum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") + ret = fn(device, byref(c_irqNum)) + _nvmlCheckReturn(ret) + return c_irqNum.value + + +def nvmlDeviceGetNumGpuCores(device): + c_numCores = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") + ret = fn(device, byref(c_numCores)) + _nvmlCheckReturn(ret) + return c_numCores.value + + +def nvmlDeviceGetPowerSource(device): + c_powerSource = _nvmlPowerSource_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") + ret = fn(device, byref(c_powerSource)) + _nvmlCheckReturn(ret) + return c_powerSource.value + + +def nvmlDeviceGetMemoryBusWidth(device): + c_memBusWidth = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") + ret = fn(device, byref(c_memBusWidth)) + _nvmlCheckReturn(ret) + return c_memBusWidth.value + + +def nvmlDeviceGetPcieLinkMaxSpeed(device): + c_speed = _nvmlPcieLinkMaxSpeed_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetAdaptiveClockInfoStatus(device): + c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") + ret = fn(device, byref(c_adaptiveClockInfoStatus)) + _nvmlCheckReturn(ret) + return c_adaptiveClockInfoStatus.value + + +def nvmlDeviceGetPcieSpeed(device): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + + +def nvmlDeviceGetDynamicPstatesInfo( + device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t() +): + isReference = ( + type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t + ) + dynamicpstatesinfoRef = ( + c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + ) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo") + ret = fn(device, dynamicpstatesinfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + + +def nvmlDeviceSetFanSpeed_v2(handle, index, speed): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2") + ret = fn(handle, index, speed) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetThermalSettings( + device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t() +): + isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t + thermalsettingsRef = ( + c_thermalsettings if isReference else byref(c_thermalsettings) + ) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings") + ret = fn(device, sensorindex, thermalsettingsRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] + + +def nvmlDeviceGetMinMaxClockOfPState( + device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint() +): + isReference = (type(minClockMHz) is not c_uint) or ( + type(maxClockMHz) is not c_uint + ) + minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) + maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState") + ret = fn( + device, + _nvmlClockType_t(clockType), + _nvmlClockType_t(pstate), + minClockMHzRef, + maxClockMHzRef, + ) + _nvmlCheckReturn(ret) + return ( + NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + ) + + +class c_nvmlClockOffset_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("type", _nvmlClockType_t), + ("pstate", _nvmlPstates_t), + ("clockOffsetMHz", c_int), + ("minClockOffsetMHz", c_int), + ("maxClockOffsetMHz", c_int), + ] + + +nvmlClockOffset_v1 = 0x1000018 + + +def nvmlDeviceGetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets") + ret = fn(device, info) + return NVML_SUCCESS + + +def nvmlDeviceSetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets") + ret = fn(device, info) + return NVML_SUCCESS + + +def nvmlDeviceGetSupportedPerformanceStates(device): + pstates = [] + c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) + c_size = sizeof(c_uint) * c_count.value + + # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. + pstates_array = _nvmlPstates_t * c_count.value + c_pstates = pstates_array() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedPerformanceStates") + ret = fn(device, c_pstates, c_size) + _nvmlCheckReturn(ret) + + for value in c_pstates: + if value != NVML_PSTATE_UNKNOWN: + pstates.append(value) + + return pstates + + +def nvmlDeviceGetGpcClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + + +def nvmlDeviceSetGpcClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpcClkMinMaxVfOffset( + device, minOffset=c_int(), maxOffset=c_int() +): + isReference = (type(minOffset) is not c_int) or ( + type(maxOffset) is not c_int + ) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + + +def nvmlDeviceGetMemClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + + +def nvmlDeviceSetMemClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetMemClkMinMaxVfOffset( + device, minOffset=c_int(), maxOffset=c_int() +): + isReference = (type(minOffset) is not c_int) or ( + type(maxOffset) is not c_int + ) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + + +def nvmlSystemSetConfComputeGpusReadyState(state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") + ret = fn(c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetConfComputeGpusReadyState(): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +def nvmlSystemGetConfComputeCapabilities(): + c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") + ret = fn(byref(c_ccSysCaps)) + _nvmlCheckReturn(ret) + return c_ccSysCaps + + +def nvmlSystemGetConfComputeState(): + c_state = c_nvmlConfComputeSystemState_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + + +def nvmlSystemGetConfComputeSettings(settings): + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") + return fn(settings) + + +def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") + ret = fn(device, c_ccMemSize) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetConfComputeMemSizeInfo(device): + c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") + ret = fn(device, byref(c_ccMemSize)) + _nvmlCheckReturn(ret) + return c_ccMemSize + + +def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") + ret = fn(device, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + + +def nvmlDeviceGetConfComputeGpuCertificate(device): + c_cert = c_nvmlConfComputeGpuCertificate_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") + ret = fn(device, byref(c_cert)) + _nvmlCheckReturn(ret) + return c_cert + + +def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): + c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() + c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) + setattr(c_attestReport, "nonce", c_nonce_arr) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") + ret = fn(device, byref(c_attestReport)) + _nvmlCheckReturn(ret) + return c_attestReport + + +def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): + c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 + c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv + fn = _nvmlGetFunctionPointer( + "nvmlSystemSetConfComputeKeyRotationThresholdInfo" + ) + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): + c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 + fn = _nvmlGetFunctionPointer( + "nvmlSystemGetConfComputeKeyRotationThresholdInfo" + ) + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return c_keyRotationThrInfo + + +## GPM ## +######### + +## Enums/defines + +#### GPM Metric Identifiers +NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DRAM_BW_UTIL = ( + 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +) +NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = ( + 60 # NvLink read bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = ( + 61 # NvLink write bandwidth for all links in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = ( + 62 # NvLink read bandwidth for link 0 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = ( + 63 # NvLink write bandwidth for link 0 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = ( + 64 # NvLink read bandwidth for link 1 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = ( + 65 # NvLink write bandwidth for link 1 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = ( + 66 # NvLink read bandwidth for link 2 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = ( + 67 # NvLink write bandwidth for link 2 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = ( + 68 # NvLink read bandwidth for link 3 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = ( + 69 # NvLink write bandwidth for link 3 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = ( + 70 # NvLink read bandwidth for link 4 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = ( + 71 # NvLink write bandwidth for link 4 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = ( + 72 # NvLink read bandwidth for link 5 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = ( + 73 # NvLink write bandwidth for link 5 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = ( + 74 # NvLink read bandwidth for link 6 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = ( + 75 # NvLink write bandwidth for link 6 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = ( + 76 # NvLink read bandwidth for link 7 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = ( + 77 # NvLink write bandwidth for link 7 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = ( + 78 # NvLink read bandwidth for link 8 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = ( + 79 # NvLink write bandwidth for link 8 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = ( + 80 # NvLink read bandwidth for link 9 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = ( + 81 # NvLink write bandwidth for link 9 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = ( + 82 # NvLink read bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = ( + 83 # NvLink write bandwidth for link 10 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = ( + 84 # NvLink read bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = ( + 85 # NvLink write bandwidth for link 11 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = ( + 86 # NvLink read bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = ( + 87 # NvLink write bandwidth for link 12 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = ( + 88 # NvLink read bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = ( + 89 # NvLink write bandwidth for link 13 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = ( + 90 # NvLink read bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = ( + 91 # NvLink write bandwidth for link 14 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = ( + 92 # NvLink read bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = ( + 93 # NvLink write bandwidth for link 15 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = ( + 94 # NvLink read bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = ( + 95 # NvLink write bandwidth for link 16 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = ( + 96 # NvLink read bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = ( + 97 # NvLink write bandwidth for link 17 in MiB/sec +) +NVML_GPM_METRIC_MAX = 98 + +## Structs + + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ("name", c_char * 96), + ("id", c_char * 96), + ("serial", c_char * 96), + ("firmwareVersion", c_char * 96), + ] + + +class struct_c_nvmlGpmSample_t(Structure): + pass # opaque handle + + +c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + + +class c_metricInfo_t(Structure): + _fields_ = [ + ("shortName", c_char_p), + ("longName", c_char_p), + ("unit", c_char_p), + ] + + +class c_nvmlGpmMetric_t(_PrintableStructure): + _fields_ = [ + ("metricId", c_uint), + ("nvmlReturn", _nvmlReturn_t), + ("value", c_double), + ("metricInfo", c_metricInfo_t), + ] + + +class c_nvmlGpmMetricsGet_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("numMetrics", c_uint), + ("sample1", c_nvmlGpmSample_t), + ("sample2", c_nvmlGpmSample_t), + ("metrics", c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX), + ] + + +NVML_GPM_METRICS_GET_VERSION = 1 + + +class c_nvmlGpmSupport_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("isSupportedDevice", c_uint), + ] + + +NVML_GPM_SUPPORT_VERSION = 1 + +## Functions + + +def nvmlGpmMetricsGet(metricsGet): + fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") + ret = fn(byref(metricsGet)) + _nvmlCheckReturn(ret) + return metricsGet + + +def nvmlGpmSampleFree(gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") + ret = fn(gpmSample) + _nvmlCheckReturn(ret) + return + + +def nvmlGpmSampleAlloc(): + gpmSample = c_nvmlGpmSample_t() + fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") + ret = fn(byref(gpmSample)) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmSampleGet(device, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") + ret = fn(device, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") + ret = fn(device, gpuInstanceId, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + + +def nvmlGpmQueryDeviceSupport(device): + gpmSupport = c_nvmlGpmSupport_t() + gpmSupport.version = NVML_GPM_SUPPORT_VERSION + fn = _nvmlGetFunctionPointer("nvmlGpmQueryDeviceSupport") + ret = fn(device, byref(gpmSupport)) + _nvmlCheckReturn(ret) + return gpmSupport + + +def nvmlGpmSetStreamingEnabled(device, state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") + ret = fn(device, c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlGpmQueryIfStreamingEnabled(device): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") + ret = fn(device, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + + +# Low Power Structure and Function + +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 + +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + + +class c_nvmlNvLinkPowerThres_t(Structure): + _fields_ = [ + ("lowPwrThreshold", c_uint), + ] + + +def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): + c_info = c_nvmlNvLinkPowerThres_t() + c_info.lowPwrThreshold = l1threshold + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkDeviceLowPowerThreshold") + ret = fn(device, byref(c_info)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +NVML_GPU_FABRIC_UUID_LEN = 16 + +_nvmlGpuFabricState_t = c_uint +NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + + +class c_nvmlGpuFabricInfo_t(_PrintableStructure): + _fields_ = [ + ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ] + + +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 + +nvmlGpuFabricInfo_v2 = 0x02000024 + + +class c_nvmlGpuFabricInfoV_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ("healthMask", c_uint32), + ] + + def __init__(self): + super(c_nvmlGpuFabricInfoV_t, self).__init__( + version=nvmlGpuFabricInfo_v2 + ) + + +def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo") + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV") + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +###################### +## Enums/defines +#### NVML GPU NVLINK BW MODE +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + + +def nvmlSystemSetNvlinkBwMode(mode): + fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") + ret = fn(mode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlSystemGetNvlinkBwMode(): + mode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") + ret = fn(byref(mode)) + _nvmlCheckReturn(ret) + return mode.value + + +_nvmlPowerScopeType_t = c_uint +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + + +class c_nvmlPowerValue_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("powerScope", _nvmlPowerScopeType_t), + ("powerValueMw", c_uint), + ] + _fmt_ = {"": "%d B"} + + +nvmlPowerValue_v2 = 0x0200000C + + +def nvmlDeviceSetPowerManagementLimit_v2( + device, powerScope, powerLimit, version=nvmlPowerValue_v2 +): + c_powerScope = _nvmlPowerScopeType_t(powerScope) + c_powerValue = c_nvmlPowerValue_v2_t() + c_powerValue.version = c_uint(version) + c_powerValue.powerScope = c_powerScope + c_powerValue.powerValueMw = c_uint(powerLimit) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit_v2") + ret = fn(device, byref(c_powerValue)) + return NVML_SUCCESS + + +class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("aggregateUncParity", c_ulonglong), + ("aggregateUncSecDed", c_ulonglong), + ("aggregateCor", c_ulonglong), + ("volatileUncParity", c_ulonglong), + ("volatileUncSecDed", c_ulonglong), + ("volatileCor", c_ulonglong), + ("aggregateUncBucketL2", c_ulonglong), + ("aggregateUncBucketSm", c_ulonglong), + ("aggregateUncBucketPcie", c_ulonglong), + ("aggregateUncBucketMcu", c_ulonglong), + ("aggregateUncBucketOther", c_ulonglong), + ("bThresholdExceeded", c_uint), + ] + + def __init__(self): + super(c_nvmlEccSramErrorStatus_v1_t, self).__init__( + version=nvmlEccSramErrorStatus_v1 + ) + + +nvmlEccSramErrorStatus_v1 = 0x1000068 + + +def nvmlDeviceGetSramEccErrorStatus(device, status): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") + ret = fn(device, status) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +NVML_DEV_CAP_EGM = 1 << 0 +nvmlDeviceCapabilities_v1 = 0x1000008 + + +class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("capMask", c_uint), + ] + + def __init__(self): + super(c_nvmlDeviceCapabilities_v1_t, self).__init__( + version=nvmlDeviceCapabilities_v1 + ) + + +def nvmlDeviceGetCapabilities(device, caps): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") + return fn(device, caps) + + +class c_nvmlPlatformInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("ibGuid", c_char * 16), + ("rackGuid", c_char * 16), + ("chassisPhysicalSlotNumber", c_char), + ("computeSlotIndex", c_char), + ("nodeIndex", c_char), + ("peerType", c_char), + ("moduleId", c_char), + ] + + def __init__(self): + super(c_nvmlPlatformInfo_v1_t, self).__init__( + version=nvmlPlatformInfo_v1 + ) + + +nvmlPlatformInfo_v1 = 0x100002C + + +def nvmlDeviceGetPlatformInfo(device, platformInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") + ret = fn(device, platformInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +class c_nvmlMask255_t(_PrintableStructure): + _fields_ = [ + ("mask", c_uint * 8), + ] + + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 + +nvmlWorkloadPowerProfileInfo_v1 = 0x100002C + + +class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("profileId", c_uint), + ("priority", c_uint), + ("conflictingmask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileInfo_v1 + ) + + +nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002BF8 + + +class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ( + "perfProfile", + c_nvmlWorkloadPowerProfileInfo_v1_t + * NVML_WORKLOAD_POWER_MAX_PROFILES, + ), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileProfilesInfo_v1 + ) + + +nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 + + +class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("perfProfilesMask", c_nvmlMask255_t), + ("requestedProfilesMask", c_nvmlMask255_t), + ("enforcedProfilesMask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileCurrentProfiles_v1 + ) + + +nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 + + +class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("requestedProfilesMask", c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__( + version=nvmlWorkloadPowerProfileRequestedProfiles_v1 + ) + + +def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): + fn = _nvmlGetFunctionPointer( + "nvmlDeviceWorkloadPowerProfileGetProfilesInfo" + ) + ret = fn(device, profilesInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): + fn = _nvmlGetFunctionPointer( + "nvmlDeviceWorkloadPowerProfileGetCurrentProfiles" + ) + ret = fn(device, currentProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles( + device, requestedProfiles +): + fn = _nvmlGetFunctionPointer( + "nvmlDeviceWorkloadPowerProfileSetRequestedProfiles" + ) + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles( + device, requestedProfiles +): + fn = _nvmlGetFunctionPointer( + "nvmlDeviceWorkloadPowerProfileClearRequestedProfiles" + ) + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") + ret = fn(device, supportedBwModes) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceGetNvlinkBwMode(device, getBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") + ret = fn(device, getBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +def nvmlDeviceSetNvlinkBwMode(device, setBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") + ret = fn(device, setBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + + +nvmlDramEncryptionInfo_v1 = 0x01000008 + + +class c_nvmlDramEncryptionInfo_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("encryptionState", _nvmlEnableState_t), + ] + + def __init__(self): + super(c_nvmlDramEncryptionInfo_t, self).__init__( + version=nvmlDramEncryptionInfo_v1 + ) + + +def nvmlDeviceGetDramEncryptionMode(handle): + c_currState = c_nvmlDramEncryptionInfo_t() + c_pendingState = c_nvmlDramEncryptionInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDramEncryptionMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.encryptionState, c_pendingState.encryptionState] + + +# added to API +def nvmlDeviceGetCurrentDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[0] + + +# added to API +def nvmlDeviceGetPendingDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[1] + + +def nvmlDeviceSetDramEncryptionMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") + c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() + c_dramEncryptionMode.encryptionState = mode + ret = fn(handle, byref(c_dramEncryptionMode)) + _nvmlCheckReturn(ret) + return None + + +# Power Smoothing defines +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 + +nvmlPowerSmoothingState_v1 = 0x1000008 + + +class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("state", c_uint), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingState_v1_t, self).__init__( + version=nvmlPowerSmoothingState_v1 + ) + + +nvmlPowerSmoothingProfile_v1 = 0x1000018 + + +class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("profileId", c_uint), + ("paramId", c_uint), + ("value", c_double), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__( + version=nvmlPowerSmoothingProfile_v1 + ) + + +def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): + fn = _nvmlGetFunctionPointer( + "nvmlDevicePowerSmoothingActivatePresetProfile" + ) + ret = fn(device, profile) + _nvmlCheckReturn(ret) + + +def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): + fn = _nvmlGetFunctionPointer( + "nvmlDevicePowerSmoothingUpdatePresetProfileParam" + ) + ret = fn(device, profile) + _nvmlCheckReturn(ret) + + +def nvmlDevicePowerSmoothingSetState(device, state): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") + ret = fn(device, state) + _nvmlCheckReturn(ret) diff --git a/worldcompass/fastvideo/training/__init__.py b/worldcompass/fastvideo/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/worldcompass/fastvideo/training/__init__.py @@ -0,0 +1 @@ + diff --git a/worldcompass/fastvideo/training/activation_checkpoint.py b/worldcompass/fastvideo/training/activation_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..56ffcdb759c1e1df2baba98313fbcd0740001216 --- /dev/null +++ b/worldcompass/fastvideo/training/activation_checkpoint.py @@ -0,0 +1,105 @@ +import collections +from enum import Enum + +import torch +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + checkpoint_wrapper, +) + +TRANSFORMER_BLOCK_NAMES = [ + "blocks", + "double_blocks", + "single_blocks", + "transformer_blocks", + "temporal_transformer_blocks", + "transformer_double_blocks", + "transformer_single_blocks", +] + + +class CheckpointType(str, Enum): + FULL = "full" + OPS = "ops" + BLOCK_SKIP = "block_skip" + + +_SELECTIVE_ACTIVATION_CHECKPOINTING_OPS = { + torch.ops.aten.mm.default, + torch.ops.aten._scaled_dot_product_efficient_attention.default, + torch.ops.aten._scaled_dot_product_flash_attention.default, + torch.ops._c10d_functional.reduce_scatter_tensor.default, +} + + +def apply_activation_checkpointing( + module: torch.nn.Module, + checkpointing_type: str = CheckpointType.FULL, + n_layer: int = 1, +) -> torch.nn.Module: + if checkpointing_type == CheckpointType.FULL: + module = _apply_activation_checkpointing_blocks(module) + elif checkpointing_type == CheckpointType.OPS: + module = _apply_activation_checkpointing_ops( + module, _SELECTIVE_ACTIVATION_CHECKPOINTING_OPS + ) + elif checkpointing_type == CheckpointType.BLOCK_SKIP: + module = _apply_activation_checkpointing_blocks(module, n_layer) + else: + raise ValueError( + f"Checkpointing type '{checkpointing_type}' not supported. " + f"Supported types are {CheckpointType.__members__.keys()}" + ) + return module + + +def _apply_activation_checkpointing_blocks( + module: torch.nn.Module, n_layer: int | None = None +) -> torch.nn.Module: + for transformer_block_name in TRANSFORMER_BLOCK_NAMES: + blocks: torch.nn.Module = getattr(module, transformer_block_name, None) + if blocks is None: + continue + for index, (layer_id, block) in enumerate(blocks.named_children()): + if n_layer is None or index % n_layer == 0: + block = checkpoint_wrapper(block, preserve_rng_state=False) + blocks.register_module(layer_id, block) + return module + + +def _apply_activation_checkpointing_ops( + module: torch.nn.Module, ops +) -> torch.nn.Module: + from torch.utils.checkpoint import ( + CheckpointPolicy, + create_selective_checkpoint_contexts, + ) + + def _get_custom_policy(meta: dict[str, int]) -> CheckpointPolicy: + + def _custom_policy(ctx, func, *args, **kwargs): + mode = "recompute" if ctx.is_recompute else "forward" + mm_count_key = f"{mode}_mm_count" + if func == torch.ops.aten.mm.default: + meta[mm_count_key] += 1 + # Saves output of all compute ops, except every second mm + to_save = func in ops and not ( + func == torch.ops.aten.mm.default + and meta[mm_count_key] % 2 == 0 + ) + return ( + CheckpointPolicy.MUST_SAVE + if to_save + else CheckpointPolicy.PREFER_RECOMPUTE + ) + + return _custom_policy + + def selective_checkpointing_context_fn(): + meta: dict[str, int] = collections.defaultdict(int) + return create_selective_checkpoint_contexts(_get_custom_policy(meta)) + + return checkpoint_wrapper( + module, + context_fn=selective_checkpointing_context_fn, + preserve_rng_state=False, + ) diff --git a/worldcompass/fastvideo/training/checkpointing_utils.py b/worldcompass/fastvideo/training/checkpointing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1684c7a8ecffbe203aa018302dda7af00f7ffd09 --- /dev/null +++ b/worldcompass/fastvideo/training/checkpointing_utils.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +import random +from typing import Any + +import numpy as np +import torch +import torch.distributed.checkpoint.stateful +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) + + +class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful): + def __init__(self, model: torch.nn.Module) -> None: + self.model = model + + def state_dict(self) -> dict[str, Any]: + state_dict = get_model_state_dict(self.model) # type: ignore[no-any-return] + # filter out non-trainable parameters + param_requires_grad = set( + [ + k + for k, v in dict(self.model.named_parameters()).items() + if v.requires_grad + ] + ) + state_dict = { + k: v for k, v in state_dict.items() if k in param_requires_grad + } + return state_dict # type: ignore + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + set_model_state_dict( + self.model, + model_state_dict=state_dict, + options=StateDictOptions(strict=False), + ) + + +class OptimizerWrapper(torch.distributed.checkpoint.stateful.Stateful): + def __init__( + self, model: torch.nn.Module, optimizer: torch.optim.Optimizer + ) -> None: + self.model = model + self.optimizer = optimizer + + def state_dict(self) -> dict[str, Any]: + return get_optimizer_state_dict( # type: ignore[no-any-return] + self.model, + self.optimizer, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + set_optimizer_state_dict( + self.model, + self.optimizer, + optim_state_dict=state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + + +class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful): + def __init__(self, scheduler) -> None: + self.scheduler = scheduler + + def state_dict(self) -> dict[str, Any]: + return {"scheduler": self.scheduler.state_dict()} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.scheduler.load_state_dict(state_dict["scheduler"]) + + +class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful): + def __init__(self, noise_generator: torch.Generator | None = None) -> None: + self.noise_generator = noise_generator + + def state_dict(self) -> dict[str, Any]: + state = { + "torch_rng_state": torch.get_rng_state(), + "numpy_rng_state": np.random.get_state(), + "python_rng_state": random.getstate(), + } + + if torch.cuda.is_available(): + state["cuda_rng_state"] = torch.cuda.get_rng_state() + if torch.cuda.device_count() > 1: + state["cuda_rng_state_all"] = torch.cuda.get_rng_state_all() + + if self.noise_generator is not None: + state["noise_generator_state"] = self.noise_generator.get_state() + + return state + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if "torch_rng_state" in state_dict: + torch.set_rng_state(state_dict["torch_rng_state"]) + + if "numpy_rng_state" in state_dict: + np.random.set_state(state_dict["numpy_rng_state"]) + + if "python_rng_state" in state_dict: + random.setstate(state_dict["python_rng_state"]) + + # Restore CUDA random state + if torch.cuda.is_available(): + if "cuda_rng_state" in state_dict: + torch.cuda.set_rng_state(state_dict["cuda_rng_state"]) + if "cuda_rng_state_all" in state_dict: + torch.cuda.set_rng_state_all(state_dict["cuda_rng_state_all"]) + + # Restore noise generator state + if ( + "noise_generator_state" in state_dict + and self.noise_generator is not None + ): + self.noise_generator.set_state(state_dict["noise_generator_state"]) diff --git a/worldcompass/fastvideo/training/muon.py b/worldcompass/fastvideo/training/muon.py new file mode 100644 index 0000000000000000000000000000000000000000..f534498d632d177623b36bdc6473169ff29cee77 --- /dev/null +++ b/worldcompass/fastvideo/training/muon.py @@ -0,0 +1,243 @@ +import math +import torch + +try: + from torch.distributed.tensor import DTensor + from torch.distributed.tensor.placement_types import ( + Partial, + Placement, + Replicate, + Shard, + ) +except ImportError: + # handle old pytorch versions + Dtensor = None + + +# This code is modified from the GitHub repository of KellerJordan: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +def zeropower_via_newtonschulz5(G, steps=5): + """Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. + + We opt to use a quintic iteration whose coefficients are selected to maximize the slope at zero. + For the purpose of minimizing steps, it turns out to be empirically effective to keep increasing + the slope at zero even beyond the point where the iteration no longer converges all the way to + one everywhere on the interval. This iteration therefore does not produce UV^T but rather + something like US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not + to hurt model performance at all relative to UV^T, where USV^T = G is the SVD. + """ + if isinstance(G, DTensor): + device_mesh = G.device_mesh + G = G.full_tensor() + else: + device_mesh = None + + assert len(G.shape) >= 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G + if G.size(-2) > G.size(-1): + X = X.mT + # Ensure spectral norm is at most 1 + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + B = ( + b * A + c * A @ A + ) # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng + X = a * X + B @ X + + if G.size(-2) > G.size(-1): + X = X.mT + + if device_mesh != None: + return DTensor.from_local(X, device_mesh) + else: + return X + + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + lr=1e-3, + wd=0.1, + muon_params=None, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_params=None, + adamw_betas=(0.95, 0.95), + adamw_eps=1e-8, + ): + + defaults = dict( + lr=lr, + wd=wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) + + params = list(muon_params) + adamw_params = list(adamw_params) if adamw_params is not None else [] + params.extend(adamw_params) + super().__init__(params, defaults) + # Sort parameters into those for which we will use Muon, and those for which we will not + for p in muon_params: + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim >= 2, p.ndim + self.state[p]["use_muon"] = True + for p in adamw_params: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + # generate weight updates in distributed fashion + for p in params: + # sanity check + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + g = g.bfloat16() + u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u.view(p.shape), alpha=-adjusted_lr) + + ############################ + # AdamW backup # + ############################ + + params = [ + p for p in group["params"] if not self.state[p]["use_muon"] + ] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss + + +# help function to create the Muon optimizer +def get_muon_optimizer( + model, + lr=1e-3, + weight_decay=0.1, + momentum=0.95, + adamw_betas=(0.95, 0.95), + adamw_eps=1e-8, +): + muon_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and p.ndim >= 2 + ] + adamw_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and not (p.ndim >= 2) + ] + + return Muon( + lr=lr, + wd=weight_decay, + muon_params=muon_params, + momentum=momentum, + adamw_params=adamw_params, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) diff --git a/worldcompass/fastvideo/training/training_utils.py b/worldcompass/fastvideo/training/training_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3b09051578a49ac0d5043145bdc1fd24c724b546 --- /dev/null +++ b/worldcompass/fastvideo/training/training_utils.py @@ -0,0 +1,1729 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import math +import os +import time +from collections.abc import Callable, Iterator +from enum import Enum +from typing import Any +import shutil + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from einops import rearrange +from safetensors.torch import save_file +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LambdaLR + +from fastvideo.distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_world_size, +) +from fastvideo.logger import init_logger +from fastvideo.training.checkpointing_utils import ( + ModelWrapper, + OptimizerWrapper, + RandomStateWrapper, + SchedulerWrapper, +) + +logger = init_logger(__name__) + +_HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False + + +def gather_state_dict_on_cpu_rank0( + model, + device: torch.device | None = None, +) -> dict[str, Any]: + rank = dist.get_rank() + cpu_state_dict = {} + sharded_sd = model.state_dict() + param_requires_grad = set( + [ + k.replace("._checkpoint_wrapped_module.", ".") + for k, v in dict(model.named_parameters()).items() + if v.requires_grad + ] + ) + for param_name, param in sharded_sd.items(): + if param_name not in param_requires_grad: + continue + if hasattr(param, "_local_tensor"): + # DTensor case + if param.is_cpu: + # Gather directly on CPU + param = param.full_tensor() + else: + if device is not None: + param = param.to(device) + param = param.full_tensor() + else: + # Regular tensor case + if param.is_cpu: + pass + else: + if device is not None: + param = param.to(device) + + if rank == 0: + cpu_state_dict[param_name] = param.cpu() + + return cpu_state_dict + + +def compute_density_for_timestep_sampling( + weighting_scheme: str, + batch_size: int, + generator, + logit_mean: float | None = None, + logit_std: float | None = None, + mode_scale: float | None = None, +): + """Compute the density for sampling the timesteps when doing SD3 training. + + Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. + + SD3 paper reference: https://arxiv.org/abs/2403.03206v1. + """ + if weighting_scheme == "logit_normal": + # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). + u = torch.normal( + mean=logit_mean, + std=logit_std, + size=(batch_size,), + device="cpu", + generator=generator, + ) + u = torch.nn.functional.sigmoid(u) + elif weighting_scheme == "mode": + u = torch.rand(size=(batch_size,), device="cpu", generator=generator) + u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u) + else: + u = torch.rand(size=(batch_size,), device="cpu", generator=generator) + return u + + +def get_sigmas( + noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32 +) -> torch.Tensor: + sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(device) + timesteps = timesteps.to(device) + step_indices = [ + (schedule_timesteps == t).nonzero().item() for t in timesteps + ] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + +def remove_previous_state(output_dir, rank): + checkpoint_dir_list = os.listdir(output_dir) + for checkpoint_dir in checkpoint_dir_list: + save_dir = os.path.join(output_dir, checkpoint_dir) + dcp_dir = os.path.join(save_dir, "distributed_checkpoint") + if os.path.exists(dcp_dir) and rank == 0: + shutil.rmtree(dcp_dir) + + +def save_checkpoint( + transformer, + rank, + output_dir, + step, + optimizer=None, + dataloader=None, + scheduler=None, + noise_generator=None, +) -> None: + """Save checkpoint following finetrainer's distributed checkpoint approach. + + Saves both distributed checkpoint and consolidated model weights. + """ + if os.path.exists(output_dir): + remove_previous_state(output_dir, rank) + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + + states = { + "model": ModelWrapper(transformer), + "random_state": RandomStateWrapper(noise_generator), + } + + if optimizer is not None: + states["optimizer"] = OptimizerWrapper(transformer, optimizer) + + if dataloader is not None: + states["dataloader"] = dataloader + + if scheduler is not None: + states["scheduler"] = SchedulerWrapper(scheduler) + dcp_dir = os.path.join(save_dir, "distributed_checkpoint") + logger.info( + "rank: %s, saving distributed checkpoint to %s", + rank, + dcp_dir, + local_main_process_only=False, + ) + + if rank == 0: + os.makedirs(dcp_dir, exist_ok=True) + + time.sleep(1) + + begin_time = time.perf_counter() + dcp.save(states, checkpoint_id=dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, distributed checkpoint saved in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + + cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None) + if rank == 0: + # Save model weights (consolidated) + transformer_save_dir = os.path.join(save_dir, "transformer") + os.makedirs(transformer_save_dir, exist_ok=True) + weight_path = os.path.join( + transformer_save_dir, "diffusion_pytorch_model.safetensors" + ) + logger.info( + "rank: %s, saving consolidated checkpoint to %s", + rank, + weight_path, + local_main_process_only=False, + ) + + # Convert training format to diffusers format and save + # diffusers_state_dict = custom_to_hf_state_dict( + # cpu_state, transformer.reverse_param_names_mapping) + save_file(cpu_state, weight_path) + + logger.info( + "rank: %s, consolidated checkpoint saved to %s", + rank, + weight_path, + local_main_process_only=False, + ) + + # Save model config + config_dict = transformer.config + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(transformer_save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + logger.info("--> checkpoint saved at step %s to %s", step, weight_path) + + +def save_distillation_checkpoint( + generator_transformer, + fake_score_transformer, + rank, + output_dir, + step, + generator_optimizer=None, + fake_score_optimizer=None, + dataloader=None, + generator_scheduler=None, + fake_score_scheduler=None, + noise_generator=None, + only_save_generator_weight=False, +) -> None: + """Save distillation checkpoint with both generator and fake_score models. Saves both + distributed checkpoint and consolidated model weights. Only saves the generator model for + inference (consolidated weights). + + Args: + only_save_generator_weight: If True, only save the generator model weights for inference + without saving distributed checkpoint for training resume. + """ + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + + # Create directories for models + inference_save_dir = os.path.join( + save_dir, "generator_inference_transformer" + ) + + # Only save distributed checkpoint if not only saving generator weight + if not only_save_generator_weight: + # Save generator distributed checkpoint + generator_states = { + "model": ModelWrapper(generator_transformer), + } + if generator_optimizer is not None: + generator_states["optimizer"] = OptimizerWrapper( + generator_transformer, generator_optimizer + ) + if dataloader is not None: + generator_states["dataloader"] = dataloader + if generator_scheduler is not None: + generator_states["scheduler"] = SchedulerWrapper( + generator_scheduler + ) + + generator_dcp_dir = os.path.join( + save_dir, "distributed_checkpoint", "generator" + ) + logger.info( + "rank: %s, saving generator distributed checkpoint to %s", + rank, + generator_dcp_dir, + local_main_process_only=False, + ) + + begin_time = time.perf_counter() + dcp.save(generator_states, checkpoint_id=generator_dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, generator distributed checkpoint saved in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + + # Save critic distributed checkpoint + critic_states = { + "model": ModelWrapper(fake_score_transformer), + } + if fake_score_optimizer is not None: + critic_states["optimizer"] = OptimizerWrapper( + fake_score_transformer, fake_score_optimizer + ) + if dataloader is not None: + critic_states["dataloader"] = dataloader + if fake_score_scheduler is not None: + critic_states["scheduler"] = SchedulerWrapper(fake_score_scheduler) + + critic_dcp_dir = os.path.join( + save_dir, "distributed_checkpoint", "critic" + ) + logger.info( + "rank: %s, saving critic distributed checkpoint to %s", + rank, + critic_dcp_dir, + local_main_process_only=False, + ) + + begin_time = time.perf_counter() + dcp.save(critic_states, checkpoint_id=critic_dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, critic distributed checkpoint saved in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + + # Save shared random state separately + shared_states = { + "random_state": RandomStateWrapper(noise_generator), + } + shared_dcp_dir = os.path.join( + save_dir, "distributed_checkpoint", "shared" + ) + + dcp.save(shared_states, checkpoint_id=shared_dcp_dir) + + else: + logger.info( + "rank: %s, skipping distributed checkpoint save (only_save_generator_weight=True)", + rank, + local_main_process_only=False, + ) + + # Save generator model weights (consolidated) for inference + cpu_state = gather_state_dict_on_cpu_rank0( + generator_transformer, device=None + ) + + if rank == 0: + # Save generator model weights (consolidated) for inference + os.makedirs(inference_save_dir, exist_ok=True) + weight_path = os.path.join( + inference_save_dir, "diffusion_pytorch_model.safetensors" + ) + logger.info( + "rank: %s, saving consolidated generator inference checkpoint to %s", + rank, + weight_path, + local_main_process_only=False, + ) + + # Convert training format to diffusers format and save + diffusers_state_dict = custom_to_hf_state_dict( + cpu_state, generator_transformer.reverse_param_names_mapping + ) + save_file(diffusers_state_dict, weight_path) + + logger.info( + "rank: %s, consolidated generator inference checkpoint saved to %s", + rank, + weight_path, + local_main_process_only=False, + ) + + # Save model config + config_dict = generator_transformer.hf_config + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(inference_save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + logger.info( + "--> distillation checkpoint saved at step %s to %s", + step, + weight_path, + ) + + +def load_checkpoint( + transformer, + rank, + checkpoint_path, + optimizer=None, + dataloader=None, + scheduler=None, + noise_generator=None, +) -> int: + """Load checkpoint following finetrainer's distributed checkpoint approach. + + Returns the step number from which training should resume. + """ + if not os.path.exists(checkpoint_path): + logger.warning("Checkpoint path %s does not exist", checkpoint_path) + return 0 + + # Extract step number from checkpoint path + step = int(os.path.basename(checkpoint_path).split("-")[-1]) + + if rank == 0: + logger.info("Loading checkpoint from step %s", step) + + dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint") + + if not os.path.exists(dcp_dir): + logger.warning( + "Distributed checkpoint directory %s does not exist", dcp_dir + ) + return 0 + + states = { + "model": ModelWrapper(transformer), + "random_state": RandomStateWrapper(noise_generator), + } + + if optimizer is not None: + states["optimizer"] = OptimizerWrapper(transformer, optimizer) + + if dataloader is not None: + states["dataloader"] = dataloader + + if scheduler is not None: + states["scheduler"] = SchedulerWrapper(scheduler) + + logger.info( + "rank: %s, loading distributed checkpoint from %s", + rank, + dcp_dir, + local_main_process_only=False, + ) + + begin_time = time.perf_counter() + dcp.load(states, checkpoint_id=dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, distributed checkpoint loaded in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + logger.info("--> checkpoint loaded from step %s", step) + + return step + + +def load_distillation_checkpoint( + generator_transformer, + fake_score_transformer, + rank, + checkpoint_path, + generator_optimizer=None, + fake_score_optimizer=None, + dataloader=None, + generator_scheduler=None, + fake_score_scheduler=None, + noise_generator=None, +) -> int: + """Load distillation checkpoint with both generator and fake_score models. + + Returns the step number from which training should resume. + """ + if not os.path.exists(checkpoint_path): + logger.warning( + "Distillation checkpoint path %s does not exist", checkpoint_path + ) + return 0 + + # Extract step number from checkpoint path + step = int(os.path.basename(checkpoint_path).split("-")[-1]) + + if rank == 0: + logger.info("Loading distillation checkpoint from step %s", step) + + # Load generator distributed checkpoint + generator_dcp_dir = os.path.join( + checkpoint_path, "distributed_checkpoint", "generator" + ) + if not os.path.exists(generator_dcp_dir): + logger.warning( + "Generator distributed checkpoint directory %s does not exist", + generator_dcp_dir, + ) + return 0 + + generator_states = { + "model": ModelWrapper(generator_transformer), + } + + if generator_optimizer is not None: + generator_states["optimizer"] = OptimizerWrapper( + generator_transformer, generator_optimizer + ) + + if dataloader is not None: + generator_states["dataloader"] = dataloader + + if generator_scheduler is not None: + generator_states["scheduler"] = SchedulerWrapper(generator_scheduler) + + logger.info( + "rank: %s, loading generator distributed checkpoint from %s", + rank, + generator_dcp_dir, + local_main_process_only=False, + ) + + begin_time = time.perf_counter() + dcp.load(generator_states, checkpoint_id=generator_dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, generator distributed checkpoint loaded in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + + # Load critic distributed checkpoint + critic_dcp_dir = os.path.join( + checkpoint_path, "distributed_checkpoint", "critic" + ) + if not os.path.exists(critic_dcp_dir): + logger.warning( + "Critic distributed checkpoint directory %s does not exist", + critic_dcp_dir, + ) + return 0 + + critic_states = { + "model": ModelWrapper(fake_score_transformer), + } + + if fake_score_optimizer is not None: + critic_states["optimizer"] = OptimizerWrapper( + fake_score_transformer, fake_score_optimizer + ) + + if dataloader is not None: + critic_states["dataloader"] = dataloader + + if fake_score_scheduler is not None: + critic_states["scheduler"] = SchedulerWrapper(fake_score_scheduler) + + logger.info( + "rank: %s, loading critic distributed checkpoint from %s", + rank, + critic_dcp_dir, + local_main_process_only=False, + ) + + begin_time = time.perf_counter() + dcp.load(critic_states, checkpoint_id=critic_dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, critic distributed checkpoint loaded in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + + # Load shared random state + shared_dcp_dir = os.path.join( + checkpoint_path, "distributed_checkpoint", "shared" + ) + if not os.path.exists(shared_dcp_dir): + logger.warning( + "Shared random state directory %s does not exist", shared_dcp_dir + ) + return 0 + + shared_states = { + "random_state": RandomStateWrapper(noise_generator), + } + + begin_time = time.perf_counter() + dcp.load(shared_states, checkpoint_id=shared_dcp_dir) + end_time = time.perf_counter() + + logger.info( + "rank: %s, shared random state loaded in %.2f seconds", + rank, + end_time - begin_time, + local_main_process_only=False, + ) + logger.info("--> distillation checkpoint loaded from step %s", step) + return step + + +def normalize_dit_input(model_type, latents, vae) -> torch.Tensor: + if model_type == "hunyuan_hf" or model_type == "hunyuan": + return latents * 0.476986 + elif model_type == "wan": + latents_mean = torch.tensor(vae.latents_mean) + latents_std = 1.0 / torch.tensor(vae.latents_std) + + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to( + device=latents.device + ) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device) + latents = ((latents.float() - latents_mean) * latents_std).to(latents) + return latents + else: + raise NotImplementedError(f"model_type {model_type} not supported") + + +def shard_latents_across_sp( + latents: torch.Tensor, num_latent_t: int +) -> torch.Tensor: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + latents = latents[:, :, :num_latent_t] + if sp_world_size > 1: + latents = rearrange( + latents, "b c (n s) h w -> b c n s h w", n=sp_world_size + ).contiguous() + latents = latents[:, :, rank_in_sp_group, :, :, :] + return latents + + +def shard_latents_dim_across_sp( + latents: torch.Tensor, total_dim +) -> torch.Tensor: + sp_world_size = get_sp_world_size() + rank_in_sp_group = get_sp_parallel_rank() + # latents = latents[:, :, :num_latent_t] + if sp_world_size > 1: + if total_dim == 5: + latents = rearrange( + latents, "b c (n s) h w -> b c n s h w", n=sp_world_size + ).contiguous() + latents = latents[:, :, rank_in_sp_group, :, :, :] + elif total_dim == 4: + latents = rearrange( + latents, "b (n s) h w -> b n s h w", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :, :, :] + elif total_dim == 3: + latents = rearrange( + latents, "b (n s) c -> b n s c", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :, :] + elif total_dim == 2: + latents = rearrange( + latents, "b (n s)-> b n s", n=sp_world_size + ).contiguous() + latents = latents[:, rank_in_sp_group, :] + elif total_dim == 1: + latents = rearrange( + latents, "(n s)-> n s", n=sp_world_size + ).contiguous() + latents = latents[rank_in_sp_group, :] + else: + raise NotImplementedError(f"total_dim {total_dim} not supported") + return latents + + +def clip_grad_norm_while_handling_failing_dtensor_cases( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, + pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, +) -> torch.Tensor | None: + global _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES + + if not _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES: + try: + return clip_grad_norm_( + parameters, + max_norm, + norm_type, + error_if_nonfinite, + foreach, + pp_mesh, + ) + except NotImplementedError as e: + if "DTensor does not support cross-mesh operation" in str(e): + # https://github.com/pytorch/pytorch/issues/134212 + logger.warning( + "DTensor does not support cross-mesh operation. If you haven't fully tensor-parallelized your " + "model, while combining other parallelisms such as FSDP, it could be the reason for this error. " + "Gradient clipping will be skipped and gradient norm will not be logged." + ) + except Exception as e: + logger.warning( + "An error occurred while clipping gradients: %s. Gradient clipping will be skipped and gradient " + "norm will not be logged.", + e, + ) + _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = ( + True + ) + return None + + +# Copied from https://github.com/pytorch/torchtitan/blob/4a169701555ab9bd6ca3769f9650ae3386b84c6e/torchtitan/utils.py#L362 +@torch.no_grad() +def clip_grad_norm_( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, + pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, +) -> torch.Tensor: + r"""Clip the gradient norm of parameters. + + Gradient norm clipping requires computing the gradient norm over the entire model. + `torch.nn.utils.clip_grad_norm_` only computes gradient norm along DP/FSDP/TP dimensions. + We need to manually reduce the gradient norm across PP stages. + See https://github.com/pytorch/torchtitan/issues/596 for details. + + Args: + parameters (`torch.Tensor` or `List[torch.Tensor]`): + Tensors that will have gradients normalized. + max_norm (`float`): + Maximum norm of the gradients after clipping. + norm_type (`float`, defaults to `2.0`): + Type of p-norm to use. Can be `inf` for infinity norm. + error_if_nonfinite (`bool`, defaults to `False`): + If `True`, an error is thrown if the total norm of the gradients from `parameters` is `nan`, `inf`, or `-inf`. + foreach (`bool`, defaults to `None`): + Use the faster foreach-based implementation. If `None`, use the foreach implementation for CUDA and CPU native tensors + and silently fall back to the slow implementation for other device types. + pp_mesh (`torch.distributed.device_mesh.DeviceMesh`, defaults to `None`): + Pipeline parallel device mesh. If not `None`, will reduce gradient norm across PP stages. + + Returns: + `torch.Tensor`: + Total norm of the gradients + """ + grads = [p.grad for p in parameters if p.grad is not None] + + # TODO(aryan): Wait for next Pytorch release to use `torch.nn.utils.get_total_norm` + # total_norm = torch.nn.utils.get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + + # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. + # We can simply reduce the DTensor to get the total norm in this tensor's process group + # and then convert it to a local tensor. + # It has two purposes: + # 1. to make sure the total norm is computed correctly when PP is used (see below) + # 2. to return a reduced total_norm tensor whose .item() would return the correct value + if isinstance(total_norm, torch.distributed.tensor.DTensor): + # Will reach here if any non-PP parallelism is used. + # If only using PP, total_norm will be a local tensor. + total_norm = total_norm.full_tensor() + + if pp_mesh is not None: + raise NotImplementedError("Pipeline parallel is not supported") + if math.isinf(norm_type): + dist.all_reduce( + total_norm, op=dist.ReduceOp.MAX, group=pp_mesh.get_group() + ) + else: + total_norm **= norm_type + dist.all_reduce( + total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group() + ) + total_norm **= 1.0 / norm_type + + _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) + return total_norm + + +@torch.no_grad() +def _clip_grads_with_norm_( + parameters: torch.Tensor | list[torch.Tensor], + max_norm: float, + total_norm: torch.Tensor, + foreach: bool | None = None, +) -> None: + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + grads = [p.grad for p in parameters if p.grad is not None] + max_norm = float(max_norm) + if len(grads) == 0: + return + grouped_grads: dict[ + tuple[torch.device, torch.dtype], + tuple[list[list[torch.Tensor]], list[int]], + ] = _group_tensors_by_device_and_dtype( + [grads] + ) # type: ignore[assignment] + + clip_coef = max_norm / (total_norm + 1e-6) + + # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so + # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization + # when the gradients do not reside in CPU memory. + clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + for (device, _), ([device_grads], _) in grouped_grads.items(): + if (foreach is None and _has_foreach_support(device_grads, device)) or ( + foreach and _device_has_foreach_support(device) + ): + torch._foreach_mul_(device_grads, clip_coef_clamped.to(device)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + clip_coef_clamped_device = clip_coef_clamped.to(device) + for g in device_grads: + g.mul_(clip_coef_clamped_device) + + +def _get_total_norm( + tensors: torch.Tensor | list[torch.Tensor], + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, +) -> torch.Tensor: + tensors = [tensors] if isinstance(tensors, torch.Tensor) else list(tensors) + norm_type = float(norm_type) + if len(tensors) == 0: + return torch.tensor(0.0) + first_device = tensors[0].device + grouped_tensors: dict[ + tuple[torch.device, torch.dtype], + tuple[list[list[torch.Tensor]], list[int]], + ] = _group_tensors_by_device_and_dtype( + [tensors] # type: ignore[list-item] + ) # type: ignore[assignment] + + norms: list[torch.Tensor] = [] + for (device, _), ([device_tensors], _) in grouped_tensors.items(): + local_tensors = [ + ( + t.to_local() + if isinstance(t, torch.distributed.tensor.DTensor) + else t + ) + for t in device_tensors + ] + if ( + foreach is None and _has_foreach_support(local_tensors, device) + ) or (foreach and _device_has_foreach_support(device)): + norms.extend(torch._foreach_norm(local_tensors, norm_type)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + norms.extend( + [torch.linalg.vector_norm(g, norm_type) for g in local_tensors] + ) + + total_norm = torch.linalg.vector_norm( + torch.stack([norm.to(first_device) for norm in norms]), norm_type + ) + + if error_if_nonfinite and torch.logical_or( + total_norm.isnan(), total_norm.isinf() + ): + raise RuntimeError( + f"The total norm of order {norm_type} for gradients from " + "`parameters` is non-finite, so it cannot be clipped. To disable " + "this error and scale the gradients by the non-finite norm anyway, " + "set `error_if_nonfinite=False`" + ) + return total_norm + + +def _get_foreach_kernels_supported_devices() -> list[str]: + r"""Return the device type list that supports foreach kernels.""" + return ["cuda", "xpu", torch._C._get_privateuse1_backend_name()] + + +@torch.no_grad() +def _group_tensors_by_device_and_dtype( + tensorlistlist: list[list[torch.Tensor | None]], + with_indices: bool = False, +) -> dict[ + tuple[torch.device, torch.dtype], + tuple[list[list[torch.Tensor | None]], list[int]], +]: + return torch._C._group_tensors_by_device_and_dtype( # type: ignore[no-any-return] + tensorlistlist, with_indices + ) + + +def _device_has_foreach_support(device: torch.device) -> bool: + return ( + device.type in (_get_foreach_kernels_supported_devices() + ["cpu"]) + and not torch.jit.is_scripting() + ) + + +def _has_foreach_support( + tensors: list[torch.Tensor], device: torch.device +) -> bool: + return _device_has_foreach_support(device) and all( + t is None or type(t) in [torch.Tensor] for t in tensors + ) + + +def custom_to_hf_state_dict( + state_dict: dict[str, Any] | Iterator[tuple[str, torch.Tensor]], + reverse_param_names_mapping: dict[str, tuple[str, int, int]], +) -> dict[str, Any]: + """Convert fastvideo's custom model format to diffusers format using + reverse_param_names_mapping. + + Args: + state_dict: State dict in fastvideo's custom format + reverse_param_names_mapping: Reverse mapping from fastvideo's custom format to diffusers format + + Returns: + State dict in diffusers format + """ + assert ( + len(reverse_param_names_mapping) > 0 + ), "reverse_param_names_mapping is empty" + if isinstance(state_dict, Iterator): + state_dict = dict(state_dict) + new_state_dict = {} + # Group parameters that need to be split (merged parameters) + merge_groups: dict[str, list[tuple[str, int, int]]] = {} + + # First pass: collect all merge groups + for training_key, ( + diffusers_key, + merge_index, + num_params_to_merge, + ) in reverse_param_names_mapping.items(): + if merge_index is not None: + # This is a merged parameter that needs to be split + if training_key not in merge_groups: + merge_groups[training_key] = [] + merge_groups[training_key].append( + (diffusers_key, merge_index, num_params_to_merge) + ) + + # Second pass: handle merged parameters by splitting them + used_keys = set() + for training_key, splits in merge_groups.items(): + if training_key in state_dict: + v = state_dict[training_key] + # Sort by merge_index to ensure correct order + splits.sort(key=lambda x: x[1]) + total = splits[0][2] + split_size = v.shape[0] // total + split_tensors = torch.split(v, split_size, dim=0) + + for diffusers_key, split_index, _ in splits: + new_state_dict[diffusers_key] = split_tensors[split_index] + used_keys.add(training_key) + + # Third pass: handle regular parameters (direct mappings) + for training_key, v in state_dict.items(): + if training_key in used_keys: + continue + + if training_key in reverse_param_names_mapping: + diffusers_key, merge_index, _ = reverse_param_names_mapping[ + training_key + ] + if merge_index is None: + # Direct mapping + new_state_dict[diffusers_key] = v + else: + # No mapping found, keep as is + new_state_dict[training_key] = v + + return new_state_dict + + +def pred_noise_to_pred_video( + pred_noise: torch.Tensor, + noise_input_latent: torch.Tensor, + timestep: torch.Tensor, + scheduler: Any, +) -> torch.Tensor: + """Convert predicted noise to clean latent.""" + timestep = timestep.expand(noise_input_latent.shape[0]) + dtype = pred_noise.dtype + device = pred_noise.device + pred_noise = pred_noise.float().to(device) + noise_input_latent = noise_input_latent.float().to(device) + sigmas = scheduler.sigmas.float().to(device) + timesteps = scheduler.timesteps.float().to(device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 + ) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + pred_video = noise_input_latent - sigma_t * pred_noise + return pred_video.to(dtype) + + +def shift_timestep( + timestep: torch.Tensor, shift: float, num_train_timestep: float +) -> torch.Tensor: + if shift == 1: + return timestep + t = timestep / num_train_timestep + denominator = 1 + (shift - 1) * t + return num_train_timestep * (shift * t / denominator) + + +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Modified from diffusers/src/diffusers/optimization.py +"""PyTorch optimization for diffusion models.""" + + +class SchedulerType(Enum): + LINEAR = "linear" + COSINE = "cosine" + COSINE_WITH_RESTARTS = "cosine_with_restarts" + COSINE_WITH_MIN_LR = "cosine_with_min_lr" + POLYNOMIAL = "polynomial" + CONSTANT = "constant" + CONSTANT_WITH_WARMUP = "constant_with_warmup" + PIECEWISE_CONSTANT = "piecewise_constant" + + +def get_constant_schedule( + optimizer: Optimizer, last_epoch: int = -1 +) -> LambdaLR: + """Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch) + + +def get_constant_schedule_with_warmup( + optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1 +) -> LambdaLR: + """Create a schedule with a constant learning rate preceded by a warmup period during which the + learning rate increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1.0, num_warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def get_piecewise_constant_schedule( + optimizer: Optimizer, step_rules: str, last_epoch: int = -1 +) -> LambdaLR: + """Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + step_rules (`string`): + The rules for the learning rate. ex: rule_steps="1:10,0.1:20,0.01:30,0.005" it means that the learning rate + if multiple 1 for the first 10 steps, multiple 0.1 for the next 20 steps, multiple 0.01 for the next 30 + steps and multiple 0.005 for the other steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + rules_dict = {} + rule_list = step_rules.split(",") + for rule_str in rule_list[:-1]: + value_str, steps_str = rule_str.split(":") + steps = int(steps_str) + value = float(value_str) + rules_dict[steps] = value + last_lr_multiple = float(rule_list[-1]) + + def create_rules_function( + rules_dict: dict, last_lr_multiple: float + ) -> Callable[[int], float]: + + def rule_func(steps: int) -> float: + for step_threshold, lr_multiple in sorted(rules_dict.items()): + if steps < step_threshold: + return lr_multiple + return last_lr_multiple + + return rule_func + + rules_func = create_rules_function(rules_dict, last_lr_multiple) + + return LambdaLR(optimizer, rules_func, last_epoch=last_epoch) + + +def get_linear_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + last_epoch: int = -1, +) -> LambdaLR: + """Create a schedule with a learning rate that decreases linearly from the initial lr set in the + optimizer to 0, after a warmup period during which it increases linearly from 0 to the initial + lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return max( + 0.0, + float(num_training_steps - current_step) + / float(max(1, num_training_steps - num_warmup_steps)), + ) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float = 0.5, + last_epoch: int = -1, +) -> LambdaLR: + """Create a schedule with a learning rate that decreases following the values of the cosine + function between the initial lr set in the optimizer to 0, after a warmup period during which it + increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_periods (`float`, *optional*, defaults to 0.5): + The number of periods of the cosine function in a schedule (the default is to just decrease from the max + value to 0 following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps) + ) + return max( + 0.0, + 0.5 + * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)), + ) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_schedule_with_min_lr( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + min_lr_ratio: float = 0.1, + num_cycles: float = 0.5, + last_epoch: int = -1, +) -> LambdaLR: + """Create a schedule with a learning rate that decreases following the values of the cosine + function between the initial lr set in the optimizer to a minimum lr (min_lr_ratio * + initial_lr), after a warmup period during which it increases linearly between 0 and the initial + lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + min_lr_ratio (`float`, *optional*, defaults to 0.1): + The ratio of minimum learning rate to initial learning rate. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of periods of the cosine function in a schedule. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps) + ) + # Cosine decay from 1.0 to min_lr_ratio over num_cycles periods + # Use the same formula as standard cosine but ensure minimum is min_lr_ratio instead of 0 + cosine_value = 0.5 * ( + 1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress) + ) + # Ensure the value doesn't go below min_lr_ratio + return max(min_lr_ratio, cosine_value) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_cosine_with_hard_restarts_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: int = 1, + last_epoch: int = -1, +) -> LambdaLR: + """Create a schedule with a learning rate that decreases following the values of the cosine + function between the initial lr set in the optimizer to 0, with several hard restarts, after a + warmup period during which it increases linearly between 0 and the initial lr set in the + optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`int`, *optional*, defaults to 1): + The number of hard restarts to use. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps) + ) + if progress >= 1.0: + return 0.0 + return max( + 0.0, + 0.5 + * ( + 1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0)) + ), + ) + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def get_polynomial_decay_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + lr_end: float = 1e-7, + power: float = 1.0, + last_epoch: int = -1, +) -> LambdaLR: + """Create a schedule with a learning rate that decreases as a polynomial decay from the initial + lr set in the optimizer to end lr defined by *lr_end*, after a warmup period during which it + increases linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + lr_end (`float`, *optional*, defaults to 1e-7): + The end LR. + power (`float`, *optional*, defaults to 1.0): + Power factor. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT + implementation at + https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37 + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_init = optimizer.defaults["lr"] + if not (lr_init > lr_end): + raise ValueError( + f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})" + ) + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + elif current_step > num_training_steps: + return lr_end / lr_init # as LambdaLR multiplies by lr_init + else: + lr_range = lr_init - lr_end + decay_steps = num_training_steps - num_warmup_steps + pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps + decay = lr_range * pct_remaining**power + lr_end + return decay / lr_init # as LambdaLR multiplies by lr_init + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +# Type alias for scheduler functions +SchedulerFunction = Callable[..., LambdaLR] + +TYPE_TO_SCHEDULER_FUNCTION: dict[SchedulerType, SchedulerFunction] = { + SchedulerType.LINEAR: get_linear_schedule_with_warmup, + SchedulerType.COSINE: get_cosine_schedule_with_warmup, + SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup, + SchedulerType.COSINE_WITH_MIN_LR: get_cosine_schedule_with_min_lr, + SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup, + SchedulerType.CONSTANT: get_constant_schedule, + SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup, + SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule, +} + + +def get_scheduler( + name: str | SchedulerType, + optimizer: Optimizer, + step_rules: str | None = None, + num_warmup_steps: int | None = None, + num_training_steps: int | None = None, + num_cycles: int = 1, + power: float = 1.0, + min_lr_ratio: float = 0.1, + last_epoch: int = -1, +) -> LambdaLR: + """Unified API to get any scheduler from its name. + + Args: + name (`str` or `SchedulerType`): + The name of the scheduler to use. + optimizer (`torch.optim.Optimizer`): + The optimizer that will be used during training. + step_rules (`str`, *optional*): + A string representing the step rules to use. This is only used by the `PIECEWISE_CONSTANT` scheduler. + num_warmup_steps (`int`, *optional*): + The number of warmup steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_training_steps (`int``, *optional*): + The number of training steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_cycles (`int`, *optional*): + The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler. + power (`float`, *optional*, defaults to 1.0): + Power factor. See `POLYNOMIAL` scheduler + min_lr_ratio (`float`, *optional*, defaults to 0.1): + The ratio of minimum learning rate to initial learning rate. Used in `COSINE_WITH_MIN_LR` scheduler. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + """ + name = SchedulerType(name) + schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] + if name == SchedulerType.CONSTANT: + return schedule_func(optimizer, last_epoch=last_epoch) + + if name == SchedulerType.PIECEWISE_CONSTANT: + return schedule_func( + optimizer, step_rules=step_rules, last_epoch=last_epoch + ) + + # All other schedulers require `num_warmup_steps` + if num_warmup_steps is None: + raise ValueError( + f"{name} requires `num_warmup_steps`, please provide that argument." + ) + + if name == SchedulerType.CONSTANT_WITH_WARMUP: + return schedule_func( + optimizer, num_warmup_steps=num_warmup_steps, last_epoch=last_epoch + ) + + # All other schedulers require `num_training_steps` + if num_training_steps is None: + raise ValueError( + f"{name} requires `num_training_steps`, please provide that argument." + ) + + if name == SchedulerType.COSINE_WITH_RESTARTS: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + last_epoch=last_epoch, + ) + + if name == SchedulerType.POLYNOMIAL: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + power=power, + last_epoch=last_epoch, + ) + + if name == SchedulerType.COSINE_WITH_MIN_LR: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + min_lr_ratio=min_lr_ratio, + last_epoch=last_epoch, + ) + + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + last_epoch=last_epoch, + ) + + +class EMA_FSDP_schedule: + """ + FSDP2-friendly EMA with two modes: + - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank. + Provides a context manager to temporarily swap EMA weights into the live model for teacher forward. + - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only + using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward. + + Usage (local_shard for CM teacher): + ema = EMA_FSDP(model, decay=0.999, mode="local_shard") + for step in ...: + ema.update(model) + with ema.apply_to_model(model): + with torch.no_grad(): + y_teacher = model(...) + + Usage (rank0_full for export): + ema = EMA_FSDP(model, decay=0.999, mode="rank0_full") + ema.update(model) + ema.state_dict() # on rank 0 + """ + + def __init__( + self, + module, + min_decay: float = 0.5, + max_decay: float = 0.9, + step_decay: float = 0.01, + ckpt_decay: float = 0.9, + mode: str = "local_shard", + ): + self.min_decay = min_decay + self.max_decay = max_decay + self.step_decay = step_decay + self.ckpt_decay = ckpt_decay + self.mode = mode + self.ckpt_shadow: dict[str, torch.Tensor] = {} + self.policy_shadow: dict[str, torch.Tensor] = {} + self.rank = dist.get_rank() if dist.is_initialized() else 0 + if self.mode not in {"local_shard", "rank0_full"}: + raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}") + self._init_shadow(module) + + @staticmethod + def _to_local_tensor(t: torch.Tensor) -> torch.Tensor: + # DTensor-aware to_local fetch; fall back to raw tensor + try: + from torch.distributed.tensor import DTensor # type: ignore + + if isinstance(t, DTensor): + return t.to_local() + except Exception: + pass + return t + + @torch.no_grad() + def _init_shadow(self, module): + # local_shard: maintain EMA of local shards for requires_grad params + self.ckpt_shadow = {} + self.policy_shadow = {} + for name, p in module.named_parameters(): + if not p.requires_grad: + continue + local = self._to_local_tensor(p.detach()) + self.ckpt_shadow[name] = local.clone().float().cpu() + self.policy_shadow[name] = local.clone().float().cpu() + + @torch.no_grad() + def update_policy_shadow(self, module, step: int): + # local_shard: update local shard EMA on every rank + decay = min(self.max_decay, self.min_decay + self.step_decay * step) + for name, p in module.named_parameters(): + if not p.requires_grad: + continue + local = self._to_local_tensor(p.detach()) + v_cpu = local.float().cpu() + if name not in self.policy_shadow: + self.policy_shadow[name] = v_cpu.clone() + else: + self.policy_shadow[name].mul_(decay).add_( + v_cpu, alpha=1.0 - decay + ) + print(f"Updated policy shadow with decay {decay}") + + @torch.no_grad() + def update_ckpt_shadow(self, module): + d = self.ckpt_decay + # local_shard: update local shard EMA on every rank + for name, p in module.named_parameters(): + if not p.requires_grad: + continue + local = self._to_local_tensor(p.detach()) + v_cpu = local.float().cpu() + if name not in self.ckpt_shadow: + self.ckpt_shadow[name] = v_cpu.clone() + else: + self.ckpt_shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d) + print(f"Updated ckpt shadow with decay {d}") + + def state_dict(self) -> dict[str, torch.Tensor]: + if self.mode == "rank0_full": + return ( + {k: v.clone() for k, v in self.ckpt_shadow.items()} + if self.rank == 0 + else {} + ) + return {k: v.clone() for k, v in self.ckpt_shadow.items()} + + def load_state_dict(self, sd: dict[str, torch.Tensor]): + self.ckpt_shadow = {k: v.clone() for k, v in sd.items()} + + @torch.no_grad() + def copy_to_unwrapped(self, module) -> None: + """Copy EMA weights into a non-sharded (unwrapped) module. + + Intended for export/eval. For mode="rank0_full", only rank 0 has the full EMA state. + """ + if self.mode == "rank0_full" and self.rank != 0: + return + name_to_param = dict(module.named_parameters()) + for n, w in self.ckpt_shadow.items(): + if n in name_to_param: + p = name_to_param[n] + p.data.copy_(w.to(dtype=p.dtype, device=p.device)) + + def replace_parameters_with_ema(self, module): + with torch.no_grad(): + for name, p in module.named_parameters(): + if not p.requires_grad: + continue + # Save local shard + p_local = self._to_local_tensor(p.detach()) + if p_local.numel() == 0: + # Nothing to swap on this rank for this param + continue + if name in self.ckpt_shadow: + ema_cpu = self.ckpt_shadow[name] + if ema_cpu.numel() != p_local.numel(): + # Shard shape mismatch (e.g., empty shard here), skip + print( + f"Shard shape mismatch for parameter {name}, ema_cpu.numel(): {ema_cpu.numel()}, p_local.numel(): {p_local.numel()}" + ) + continue + # Copy EMA shard into local param shard + p_local.copy_( + ema_cpu.to(dtype=p_local.dtype, device=p_local.device) + ) + return module + + class _ApplyEMACtx: + def __init__(self, ema: "EMA_FSDP_schedule", module): + self.ema = ema + self.module = module + self.saved: dict[str, torch.Tensor] = {} + + def __enter__(self): + if self.ema.mode != "local_shard": + raise RuntimeError( + "EMA apply_to_model is only supported for mode='local_shard'" + ) + # print(f'Using ckpt shadow!!!!') + with torch.no_grad(): + for name, p in self.module.named_parameters(): + if not p.requires_grad: + continue + # Save local shard + p_local = EMA_FSDP_schedule._to_local_tensor(p.detach()) + if p_local.numel() == 0: + # Nothing to swap on this rank for this param + continue + self.saved[name] = p_local.clone().to( + device=p_local.device, dtype=p_local.dtype + ) + if name in self.ema.ckpt_shadow: + ema_cpu = self.ema.ckpt_shadow[name] + if ema_cpu.numel() != p_local.numel(): + # Shard shape mismatch (e.g., empty shard here), skip + continue + # Copy EMA shard into local param shard + p_local.copy_( + ema_cpu.to( + dtype=p_local.dtype, device=p_local.device + ) + ) + return self.module + + def __exit__(self, exc_type, exc, tb): + with torch.no_grad(): + for name, p in self.module.named_parameters(): + if name in self.saved: + p_local = EMA_FSDP_schedule._to_local_tensor(p.detach()) + if p_local.numel() == 0: + continue + saved_local = self.saved[name] + if saved_local.numel() != p_local.numel(): + continue + p_local.copy_(saved_local) + self.saved.clear() + return False + + class _ApplyPolicyEMACtx: + def __init__(self, ema: "EMA_FSDP_schedule", module): + self.ema = ema + self.module = module + self.saved: dict[str, torch.Tensor] = {} + + def __enter__(self): + if self.ema.mode != "local_shard": + raise RuntimeError( + "EMA apply_to_model is only supported for mode='local_shard'" + ) + with torch.no_grad(): + # print(f'Using policy shadow!!!!') + for name, p in self.module.named_parameters(): + if not p.requires_grad: + continue + # Save local shard + p_local = EMA_FSDP_schedule._to_local_tensor(p.detach()) + if p_local.numel() == 0: + # Nothing to swap on this rank for this param + continue + self.saved[name] = p_local.clone().to( + device=p_local.device, dtype=p_local.dtype + ) + if name in self.ema.policy_shadow: + ema_cpu = self.ema.policy_shadow[name] + if ema_cpu.numel() != p_local.numel(): + # Shard shape mismatch (e.g., empty shard here), skip + continue + # Copy EMA shard into local param shard + p_local.copy_( + ema_cpu.to( + dtype=p_local.dtype, device=p_local.device + ) + ) + return self.module + + def __exit__(self, exc_type, exc, tb): + with torch.no_grad(): + for name, p in self.module.named_parameters(): + if name in self.saved: + p_local = EMA_FSDP_schedule._to_local_tensor(p.detach()) + if p_local.numel() == 0: + continue + saved_local = self.saved[name] + if saved_local.numel() != p_local.numel(): + continue + p_local.copy_(saved_local) + self.saved.clear() + return False + + def apply_ckpt_shadow_to_model(self, module): + return EMA_FSDP_schedule._ApplyEMACtx(self, module) + + def apply_policy_shadow_to_model(self, module): + return EMA_FSDP_schedule._ApplyPolicyEMACtx(self, module) diff --git a/worldcompass/fastvideo/training/world_compass_in_train_pipeline.py b/worldcompass/fastvideo/training/world_compass_in_train_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..6e58b48a8880f8ec9f433d45a338fcad156b72b9 --- /dev/null +++ b/worldcompass/fastvideo/training/world_compass_in_train_pipeline.py @@ -0,0 +1,2117 @@ +# SPDX-License-Identifier: Apache-2.0 +import dataclasses +import math +import os +from re import U +import time +from abc import ABC, abstractmethod +from collections import deque +from collections.abc import Iterator +from typing import Any + +import imageio +import numpy as np +import torch +import torch.distributed as dist +import torchvision +from diffusers import FlowMatchEulerDiscreteScheduler +from einops import rearrange +from torch.utils.data import DataLoader +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm.auto import tqdm +import re +import json + +from fastvideo.models.utils import ( + generate_points_in_sphere, + select_aligned_memory_frames_context_per_chunk_w_latent_sink_fov_refine_hunyuan, +) + +from fastvideo.distributed.parallel_state import get_sp_parallel_rank + +import fastvideo.envs as envs +from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionMetadataBuilder, +) +from fastvideo.dataset import build_hy_camera_dataloader +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.validation_dataset import ValidationDataset +from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + get_local_torch_device, + get_sp_group, + get_world_group, + get_gpu_group, +) +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines import ( + ComposedPipelineBase, + ForwardBatch, + LoRAPipeline, + TrainingBatch, +) +from fastvideo.training.activation_checkpoint import apply_activation_checkpointing +from fastvideo.training.training_utils import ( + EMA_FSDP_schedule, + clip_grad_norm_while_handling_failing_dtensor_cases, + compute_density_for_timestep_sampling, + get_scheduler, + load_checkpoint, + normalize_dit_input, + save_checkpoint, + shard_latents_dim_across_sp, +) +from fastvideo.utils import is_vsa_available, set_random_seed, shallow_asdict +from fastvideo.training.muon import get_muon_optimizer +from fastvideo.models.hyvideo.models.autoencoders import hunyuanvideo_15_vae_w_cache + +from reward_function.reward_function import CompassReward + +import wandb # isort: skip +import random +import copy + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +def merge_tensor_by_mask(tensor_1, tensor_2, mask, dim): + assert tensor_1.shape == tensor_2.shape + # Mask is a 0/1 vector. Choose tensor_2 when the value is 1; otherwise, tensor_1 + masked_indices = torch.nonzero(mask).squeeze(1) + tmp = tensor_1.clone() + if dim == 0: + tmp[masked_indices] = tensor_2[masked_indices] + elif dim == 1: + tmp[:, masked_indices] = tensor_2[:, masked_indices] + elif dim == 2: + tmp[:, :, masked_indices] = tensor_2[:, :, masked_indices] + return tmp + + +def _get_trainable_params(model: torch.nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +class WorldCompassInTrainPipeline(LoRAPipeline, ABC): + """A pipeline for training a model. + + All training pipelines should inherit from this class. All reusable components and code should + be implemented in this class. + """ + + _required_config_modules = ["scheduler", "transformer"] + validation_pipeline: ComposedPipelineBase + train_dataloader: StatefulDataLoader + train_loader_iter: Iterator[dict[str, Any]] + current_epoch: int = 0 + + def __init__( + self, + model_path: str, + fastvideo_args: TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> None: + self.lora_training = fastvideo_args.lora_training + if self.lora_training and fastvideo_args.lora_rank is None: + raise ValueError("lora rank must be set when using lora training") + + set_random_seed(fastvideo_args.seed) # for lora param init + super().__init__( + model_path, fastvideo_args, required_config_modules, loaded_modules + ) # type: ignore + + # Initialize KV cache storage for GRPO sampling + self.points_local = generate_points_in_sphere(50000, 8.0) + self.sample_kv_caches = {} # Dict to store KV caches for all samples and chunks + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + raise RuntimeError( + "create_pipeline_stages should not be called for training pipeline" + ) + + def set_schemas(self) -> None: + self.train_dataset_schema = pyarrow_schema_t2v + + def initialize_training_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing training pipeline...") + self.device = get_local_torch_device() + self.training_args = training_args + world_group = get_world_group() + self.world_size = world_group.world_size + self.global_rank = world_group.rank + self.sp_group = get_sp_group() + self.gpu_group = get_gpu_group() + self.rank_in_sp_group = self.sp_group.rank_in_group + self.sp_world_size = self.sp_group.world_size + self.local_rank = world_group.local_rank + self.transformer = self.get_module("transformer") + # self.vae = self.get_module("vae") + + # Load VAE from config path + vae_path = training_args.vae_path + if not vae_path: + raise ValueError("vae_path must be provided in training_args") + + self.vae = hunyuanvideo_15_vae_w_cache.AutoencoderKLConv3D.from_pretrained( + vae_path, torch_dtype=torch.float32 + ).to("cpu") + self.vae = self.vae.to(torch.float32) + + # Initialize reward model with configurable camera estimator + camera_estimator = getattr(training_args, "camera_estimator", "dav3") + cache_dir = getattr(training_args, "cache_dir", None) or None + logger.info( + f"Initializing reward model with camera_estimator={camera_estimator}" + ) + self.reward_model = CompassReward( + device=self.device, camera_estimator=camera_estimator, cache_dir=cache_dir + ) + + self.seed = training_args.seed + self.set_schemas() + self.action = training_args.action + + assert self.seed is not None, "seed must be set" + set_random_seed(self.seed) + self.transformer.train() + if training_args.enable_gradient_checkpointing_type is not None: + self.transformer = apply_activation_checkpointing( + self.transformer, + checkpointing_type=training_args.enable_gradient_checkpointing_type, + ) + + self.set_trainable() + params_to_optimize = self.transformer.parameters() + params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize)) + + self.vae.requires_grad = False + + if self.training_args.eval_only: + self.transformer.requires_grad = False + self.transformer = self.transformer.to(torch.bfloat16) + + self.optimizer = get_muon_optimizer( + model=self.transformer, + lr=training_args.learning_rate, # Learning rate + weight_decay=training_args.weight_decay, # Weight decay + adamw_betas=(0.9, 0.999), # AdamW betas for 1D parameters + adamw_eps=1e-8, # AdamW epsilon + ) + + self.init_steps = 0 + logger.info("optimizer: %s", self.optimizer) + + self.lr_scheduler = get_scheduler( + training_args.lr_scheduler, + optimizer=self.optimizer, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) + + if self.action: + self.train_dataset, self.train_dataloader = build_hy_camera_dataloader( + json_path=training_args.json_path, + causal=training_args.causal, + window_frames=training_args.window_frames, + batch_size=training_args.train_batch_size, + num_data_workers=training_args.dataloader_num_workers, + drop_last=True, + drop_first_row=False, + seed=self.seed, + cfg_rate=training_args.training_cfg_rate, + i2v_rate=training_args.i2v_rate, + random_pose_path=training_args.random_pose_path, + neg_prompt_path=training_args.neg_prompt_path, + neg_byt5_prompt_path=training_args.neg_byt5_prompt_path, + ) + self.eval_dataset, self.eval_dataloader = build_hy_camera_dataloader( + json_path=training_args.eval_json_path, + causal=training_args.causal, + window_frames=training_args.window_frames, + batch_size=training_args.train_batch_size, + num_data_workers=training_args.dataloader_num_workers, + drop_last=True, + drop_first_row=False, + seed=self.seed, + cfg_rate=training_args.training_cfg_rate, + i2v_rate=training_args.i2v_rate, + random_pose_path=training_args.random_pose_path, + neg_prompt_path=training_args.neg_prompt_path, + neg_byt5_prompt_path=training_args.neg_byt5_prompt_path, + ) + + self.num_update_steps_per_epoch = math.ceil( + len(self.train_dataloader) + / training_args.gradient_accumulation_steps + * training_args.sp_size + / training_args.train_sp_batch_size + ) + self.num_train_epochs = math.ceil( + training_args.max_train_steps / self.num_update_steps_per_epoch + ) + + self.current_epoch = 0 + + if self.global_rank == 0: + project = training_args.tracker_project_name or "fastvideo" + wandb_config = dataclasses.asdict(training_args) + wandb.login(key=training_args.wandb_key) + wandb.init( + config=wandb_config, + name=training_args.output_dir.split("/")[-1], + entity=training_args.wandb_entity, + project=project, + ) + + @abstractmethod + def initialize_validation_pipeline(self, training_args: TrainingArgs): + raise NotImplementedError("Training pipelines must implement this method") + + def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: + self.vae.eval() + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + training_batch.samples_loss = {} + training_batch.samples_grad_norm = {} + return training_batch + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + extra_kwargs = { + "byt5_text_states": batch["byt5_text_states"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ), + "byt5_text_mask": batch["byt5_text_mask"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ), + } + multitask_mask = self.get_task_mask("i2v", batch["latent"].shape[2]) + cond_latents = self._prepare_cond_latents( + "i2v", batch["image_cond"], batch["latent"], multitask_mask + ) + + training_batch.latents = batch["latent"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.cond_latents = cond_latents.to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.latents_concat = torch.concat( + [batch["latent"], cond_latents], dim=1 + ).to(get_local_torch_device(), dtype=torch.bfloat16) + training_batch.prompt = batch["prompt"] + training_batch.prompt_embed = batch["prompt_embed"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.prompt_mask = batch["prompt_mask"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.vision_states = batch["vision_states"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.extra_kwargs = extra_kwargs + training_batch.w2c = batch["w2c"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.intrinsic = batch["intrinsic"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.action = batch["action"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + + return training_batch + + def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch: + latents_shape = training_batch.raw_latent_shape + patch_size = self.training_args.pipeline_config.dit_config.patch_size + current_vsa_sparsity = training_batch.current_vsa_sparsity + assert latents_shape is not None + assert training_batch.timesteps is not None + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder( # type: ignore + ).build( # type: ignore + raw_latent_shape=latents_shape[2:5], + current_timestep=training_batch.timesteps, + patch_size=patch_size, + VSA_sparsity=current_vsa_sparsity, + device=get_local_torch_device(), + ) + else: + training_batch.attn_metadata = None + + return training_batch + + def _build_rope_idx(self, training_batch: TrainingBatch) -> TrainingBatch: + rank_in_sp_group = get_sp_parallel_rank() + per_sp_seq_length = ( + training_batch.latents.shape[2] * training_batch.per_seq_length + ) + + training_batch.current_start = per_sp_seq_length * rank_in_sp_group + training_batch.current_end = ( + per_sp_seq_length * rank_in_sp_group + per_sp_seq_length + ) + return training_batch + + def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: + max_grad_norm = self.training_args.max_grad_norm + + # TODO(will): perhaps move this into transformer api so that we can do + # the following: + # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + if max_grad_norm is not None: + model_parts = [self.transformer] + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + [p for m in model_parts for p in m.parameters()], + max_grad_norm, + foreach=None, + ) + assert grad_norm is not float("nan") or grad_norm is not float("inf") + grad_norm = grad_norm.item() if grad_norm is not None else 0.0 + else: + grad_norm = 0.0 + training_batch.grad_norm = grad_norm + return training_batch + + def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: + max_grad_norm = self.training_args.max_grad_norm + + # TODO(will): perhaps move this into transformer api so that we can do + # the following: + # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + if max_grad_norm is not None: + model_parts = [self.transformer] + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + [p for m in model_parts for p in m.parameters()], + max_grad_norm, + foreach=None, + ) + assert grad_norm is not float("nan") or grad_norm is not float("inf") + grad_norm = grad_norm.item() if grad_norm is not None else 0.0 + else: + grad_norm = 0.0 + training_batch.grad_norm = grad_norm + return training_batch + + def _create_sample_kv_cache(self): + _kv_cache = [] + _kv_cache_neg = [] + transformer_num_layers = len(self.transformer.double_blocks) + for i in range(transformer_num_layers): + _kv_cache.append( + {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None} + ) + _kv_cache_neg.append( + {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None} + ) + + return {"positive": _kv_cache, "negative": _kv_cache_neg} + + def _build_kv_cache_from_previous_chunks( + self, + current_kv_cache, + latents_curr, + training_batch, + generate_latent_num, + update_latent_num, + stabilization_level, + negative=False, + ): + """Build KV cache from previous chunks to accelerate current chunk sampling.""" + transformer_dtype = self.transformer.dtype + device = latents_curr.device + # Calculate the range for KV cache building (previous chunks) + kv_cache_latent_num = generate_latent_num - update_latent_num + # kv_cache_latent_num = 4 + + if generate_latent_num >= 20: + selected_frame_indices = select_aligned_memory_frames_context_per_chunk_w_latent_sink_fov_refine_hunyuan( + training_batch.w2c[:, :generate_latent_num, :, :][0] + .to(torch.float32) + .cpu() + .detach() + .numpy(), + generate_latent_num - 4, + memory_frames=20, + temporal_context_size=12, + pred_latent_size=4, + points_local=self.points_local.to(device), + device=device, + ) + else: + selected_frame_indices = list(range(0, kv_cache_latent_num)) + + if kv_cache_latent_num <= 0: + return current_kv_cache # No previous chunks to cache + + # Extract latents for KV cache building (previous chunks only) + latents_for_cache = latents_curr[:, :, selected_frame_indices, :, :] + cond_latents_for_cache = training_batch.cond_latents[ + :, :, selected_frame_indices, :, : + ] + latents_input_for_cache = torch.cat( + [latents_for_cache, cond_latents_for_cache], dim=1 + ) + + action_for_cache = ( + training_batch.action[:, selected_frame_indices] if self.action else None + ) + + # Create timestep input for cache building (use stabilization level for previous chunks) + timestep_for_cache = torch.full( + (len(selected_frame_indices),), + stabilization_level - 1, + device=latents_curr.device, + dtype=torch.long, + ) + + viewmats_for_cache = training_batch.w2c[:, selected_frame_indices].to(device) + Ks_for_cache = training_batch.intrinsic[:, selected_frame_indices].to(device) + action_for_cache = training_batch.action[:, selected_frame_indices].to(device) + + for kv_cache_layer in current_kv_cache["positive"]: + kv_cache_layer["k_vision"] = None + kv_cache_layer["v_vision"] = None + + # Prepare inputs for KV cache building + cache_kwargs = { + "hidden_states": latents_input_for_cache, + "timestep": timestep_for_cache, + "timestep_r": None, + "return_dict": False, + "mask_type": "i2v", + "action": action_for_cache, + "viewmats": viewmats_for_cache, + "Ks": Ks_for_cache, + "kv_cache": current_kv_cache["positive"], + "cache_vision": True, + "rope_temporal_size": latents_for_cache.shape[2], + "start_rope_start_idx": 0, + } + + with ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True), + torch.no_grad(), + ): + current_kv_cache["positive"] = self.transformer( + txt_branch=False, + input_dict=cache_kwargs, + ) + + return current_kv_cache, selected_frame_indices + + def _split_kv_cache(self, kv_cache, idx): + new_kv_cache = [] + + for layer_idx in range(len(kv_cache)): + new_kv_cache.append( + { + "k": kv_cache[layer_idx]["k"][idx].unsqueeze(0).detach().clone(), + "v": kv_cache[layer_idx]["v"][idx].unsqueeze(0).detach().clone(), + } + ) + + return new_kv_cache + + def flux_step( + self, + model_output: torch.Tensor, + latents: torch.Tensor, + eta: float, + sigmas: torch.Tensor, + index: int, + prev_sample: torch.Tensor, + grpo: bool, + sde_solver: bool, + ): + sigma = sigmas[index] + dsigma = sigmas[index + 1] - sigma + prev_sample_mean = latents + dsigma * model_output + + pred_original_sample = latents - sigma * model_output + + delta_t = sigma - sigmas[index + 1] + std_dev_t = eta * math.sqrt(delta_t) + + if sde_solver: + score_estimate = -(latents - pred_original_sample * (1 - sigma)) / sigma**2 + log_term = -0.5 * eta**2 * score_estimate + prev_sample_mean = prev_sample_mean + log_term * dsigma + + prev_sample_mean = ( + prev_sample_mean + torch.randn_like(prev_sample_mean) * std_dev_t + ) + + return prev_sample_mean, pred_original_sample + + def _shard_latents_dim_across_sp(self, latents, timesteps, action): + latents = shard_latents_dim_across_sp(latents, total_dim=latents.dim()) + timesteps = shard_latents_dim_across_sp(timesteps, total_dim=timesteps.dim()) + if self.action: + action = shard_latents_dim_across_sp(action, total_dim=action.dim()) + return latents, timesteps, action + + @torch.no_grad() + def _sample_model_ode( + self, + training_batch: TrainingBatch, + selected_chunk_id, + noise_latents, + sampling_steps=None, + ) -> TrainingBatch: + self.transformer.eval() + + chunk_latent_num = ( + self.training_args.single_chunk_size + ) # hard code right now, can be changed later. -> done + stabilization_level = 15 + + # 配置sample schedule + if sampling_steps is None: + sampling_steps = self.training_args.sampling_steps + + sampling_batch_size = self.training_args.sampling_batch_size + + kv_cache = self._create_sample_kv_cache() + + t_expand_txt = ( + torch.tensor([0]).to(get_local_torch_device()).to(noise_latents.dtype) + ) + input_dict = { + "timestep_txt": t_expand_txt, + "text_states": training_batch.prompt_embed, + "encoder_attention_mask": training_batch.prompt_mask, + "vision_states": training_batch.vision_states, + "mask_type": "i2v", + "extra_kwargs": training_batch.extra_kwargs, + "kv_cache": kv_cache["positive"], + "cache_txt": True, + } + with ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True), + torch.no_grad(), + ): + kv_cache["positive"] = self.transformer( + txt_branch=True, + input_dict=input_dict, + ) + + sigma_schedule = torch.linspace(1, 0, sampling_steps + 1) + shift = 5 + sigma_schedule = (shift * sigma_schedule) / (1 + (shift - 1) * sigma_schedule) + + for chunk_i in tqdm( + range(0, selected_chunk_id - 1), desc="ODE Chunk", leave=True + ): + torch.cuda.empty_cache() + + generate_latent_num = (chunk_i + 1) * chunk_latent_num + + latents_curr = noise_latents[:, :, :generate_latent_num, :, :] + cond_latent_curr = training_batch.cond_latents[ + :, :, :generate_latent_num, :, : + ] + + w2c_curr = training_batch.w2c[:, :generate_latent_num, :, :] + intrinsic_curr = training_batch.intrinsic[:, :generate_latent_num, :, :] + action_curr = training_batch.action[:, :generate_latent_num] + + update_latent_num = chunk_latent_num + + selected_frame_indices = [] + if chunk_i > 0: + kv_cache, selected_frame_indices = ( + self._build_kv_cache_from_previous_chunks( + kv_cache, + latents_curr, + training_batch, + generate_latent_num, + update_latent_num, + stabilization_level, + ) + ) + + for i in range(0, sampling_steps): + sigma = sigma_schedule[i] + timestep_value = int(sigma * 1000) + + timestep_input = torch.full( + ( + 1, + update_latent_num, + ), + timestep_value, + device=latents_curr.device, + dtype=torch.long, + ) + + latent_concat_curr = torch.cat([latents_curr, cond_latent_curr], dim=1) + + input_dict = { + "hidden_states": latent_concat_curr[ + :, :, -update_latent_num:, :, : + ], + "timestep": timestep_input, + "timestep_r": None, + "return_dict": False, + "mask_type": "i2v", + "action": action_curr[:, -update_latent_num:], + "viewmats": w2c_curr[:, -update_latent_num:, :, :], + "Ks": intrinsic_curr[:, -update_latent_num:, :, :], + "kv_cache": kv_cache["positive"], + "cache_vision": False, + "rope_temporal_size": len(selected_frame_indices) + + update_latent_num, + "start_rope_start_idx": len(selected_frame_indices), + } + + with ( + torch.autocast( + device_type="cuda", dtype=torch.bfloat16, enabled=True + ), + torch.no_grad(), + ): + model_pred = self.transformer( + txt_branch=False, + input_dict=input_dict, + )[0] + + z, pred_original = self.flux_step( + model_pred[:, :, -update_latent_num:, :, :].to(torch.float32), + latents_curr[:, :, -update_latent_num:, :, :].to(torch.float32), + eta=0.0, + sigmas=sigma_schedule, + index=i, + prev_sample=None, + grpo=False, + sde_solver=False, + ) + latents_curr[:, :, -update_latent_num:, :, :] = z + + latents_curr[:, :, -update_latent_num:, :, :] = pred_original + + return latents_curr + + @torch.no_grad() + def _sample_reference_model(self, training_batch: TrainingBatch) -> TrainingBatch: + + # if training_batch.current_timestep % 8 == 0: + # self.transformer = self.ema_generator.replace_parameters_with_ema(self.transformer) + + with self.ema_generator.apply_policy_shadow_to_model(self.transformer): + self.transformer.eval() + + chunk_latent_num = ( + self.training_args.single_chunk_size + ) # hard code right now, can be changed later. -> done + stabilization_level = 15 + + # 配置sample schedule + sampling_steps = self.training_args.sampling_steps + sampling_batch_size = self.training_args.sampling_batch_size + + sigma_schedule = torch.linspace(1, 0, self.training_args.sampling_steps + 1) + shift = 5 + sigma_schedule = (shift * sigma_schedule) / ( + 1 + (shift - 1) * sigma_schedule + ) + + latents = training_batch.latents + bsz, latent_channels, _, latent_h, latent_w = latents.shape + + latent_t = self.training_args.window_frames + # 随机选择一个chunk作为优化目标 + # 计算当前步数对应的chunk_id，从1逐渐增加到7 + # MIN to MAX + + # Get chunk selection parameters from training_args + min_chunk_id = self.training_args.min_chunk_id + max_chunk_id = self.training_args.max_chunk_id + chunk_strategy = self.training_args.chunk_selection_strategy + + if chunk_strategy == "min2max": + selected_chunk_id = min_chunk_id + ( + (training_batch.current_timestep - 1) % max_chunk_id + ) + elif chunk_strategy == "max2min": + selected_chunk_id = max_chunk_id - ( + (training_batch.current_timestep - 1) + % (max_chunk_id - min_chunk_id + 1) + ) + else: + raise ValueError( + f'Invalid chunk_selection_strategy: {chunk_strategy}, must be one of "min2max", "max2min", or "random".' + ) + + training_batch.selected_chunk_id = selected_chunk_id + + noise = torch.randn( + (1, latent_channels, latent_t, latent_h, latent_w), # （c,t,h,w) + device=latents.device, + dtype=latents.dtype, + ) + noise = self.gpu_group.all_gather(noise, dim=0) + + context_latents = None + context_num = 0 + if selected_chunk_id > 1: + context_latents = self._sample_model_ode( + training_batch, selected_chunk_id, noise + ) + context_num = context_latents.shape[2] + + training_batch.sample_kwargs = {} + + if selected_chunk_id > 1: + if ( + hasattr(self.vae.config, "shift_factor") + and self.vae.config.shift_factor + ): + context_latents = ( + context_latents / self.vae.config.scaling_factor + + self.vae.config.shift_factor + ) + else: + context_latents = context_latents / self.vae.config.scaling_factor + + self.vae.to(self.device) + context_video_frames = self.vae.decode( + context_latents.to(torch.float32), return_dict=False + )[0] + self.vae.to("cpu") + context_video_frames = ( + (context_video_frames / 2 + 0.5).clamp(0, 1).cpu().float() + ) + context_video_frames = np.transpose( + context_video_frames[0], (1, 2, 3, 0) + ) + context_video_frames = ( + (context_video_frames * 255).numpy().astype(np.uint8) + ) + else: + context_video_frames = None + + kv_cache = self._create_sample_kv_cache() + + t_expand_txt = ( + torch.tensor([0]).to(get_local_torch_device()).to(noise.dtype) + ) + input_dict = { + "timestep_txt": t_expand_txt, + "text_states": training_batch.prompt_embed, + "encoder_attention_mask": training_batch.prompt_mask, + "vision_states": training_batch.vision_states, + "mask_type": "i2v", + "extra_kwargs": training_batch.extra_kwargs, + "kv_cache": kv_cache["positive"], + "cache_txt": True, + } + with ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True), + torch.no_grad(), + ): + kv_cache["positive"] = self.transformer( + txt_branch=True, + input_dict=input_dict, + ) + + if selected_chunk_id > 1: + noise_latents = torch.cat( + [context_latents, noise[:, :, context_num:, :, :]], dim=2 + ) + start_chunk_id = selected_chunk_id - 1 + generate_latent_num = selected_chunk_id * chunk_latent_num + kv_cache, selected_frame_indices = ( + self._build_kv_cache_from_previous_chunks( + kv_cache, + noise_latents[:, :, :generate_latent_num, :, :], + training_batch, + generate_latent_num, + chunk_latent_num, + stabilization_level, + ) + ) + training_batch.kv_cache = kv_cache + else: + start_chunk_id = 0 + selected_frame_indices = [] + training_batch.kv_cache = kv_cache + + for sample_idx in tqdm( + range(0, self.training_args.grpo_generation_num), + desc="GRPO Batch Sampling", + leave=False, + ): + new_noise = torch.randn_like(noise) + + if context_latents is not None: + noise_latents = torch.cat( + [context_latents, new_noise[:, :, context_num:, :, :]], dim=2 + ) + else: + noise_latents = new_noise + + video_num_in_each_gpu = ( + sampling_batch_size // self.training_args.gpu_para + ) + rank_in_gpu_group = self.gpu_group.rank_in_group + + for chunk_i in range(start_chunk_id, selected_chunk_id): + generate_latent_num = (chunk_i + 1) * chunk_latent_num + + latents_curr = noise_latents[:, :, :generate_latent_num, :, :] + cond_latent_curr = training_batch.cond_latents[ + :, :, :generate_latent_num, :, : + ] + + w2c_curr = training_batch.w2c[:, :generate_latent_num, :, :] + intrinsic_curr = training_batch.intrinsic[ + :, :generate_latent_num, :, : + ] + action_curr = training_batch.action[:, :generate_latent_num] + + update_latent_num = chunk_latent_num + + for i in range(0, sampling_steps): + sigma = sigma_schedule[i] + timestep_value = int(sigma * 1000) + + latent_concat_curr = torch.cat( + [latents_curr, cond_latent_curr], dim=1 + ) + timestep_input = torch.full( + ( + 1, + update_latent_num, + ), + timestep_value, + device=latents_curr.device, + dtype=torch.long, + ) + + if i == 0: + tmp_kwargs = { + "hidden_states": latent_concat_curr[ + :, :, -update_latent_num:, :, : + ], + "timestep": timestep_input, + "timestep_r": None, + "return_dict": False, + "mask_type": "i2v", + "action": action_curr[:, -update_latent_num:], + "viewmats": w2c_curr[:, -update_latent_num:, :, :], + "Ks": intrinsic_curr[:, -update_latent_num:, :, :], + "kv_cache": kv_cache["positive"], + "cache_vision": False, + "rope_temporal_size": len(selected_frame_indices) + + update_latent_num, + "start_rope_start_idx": len(selected_frame_indices), + } + + input_dict = { + "hidden_states": latent_concat_curr[ + :, :, -update_latent_num:, :, : + ], + "timestep": timestep_input, + "timestep_r": None, + "return_dict": False, + "mask_type": "i2v", + "action": action_curr[:, -update_latent_num:], + "viewmats": w2c_curr[:, -update_latent_num:, :, :], + "Ks": intrinsic_curr[:, -update_latent_num:, :, :], + "kv_cache": kv_cache["positive"], + "cache_vision": False, + "rope_temporal_size": len(selected_frame_indices) + + update_latent_num, + "start_rope_start_idx": len(selected_frame_indices), + } + + with ( + torch.autocast( + device_type="cuda", dtype=torch.bfloat16, enabled=True + ), + torch.no_grad(), + ): + model_pred = self.transformer( + txt_branch=False, + input_dict=input_dict, + )[0] + + z, pred_original = self.flux_step( + model_pred[:, :, -update_latent_num:, :, :].to( + torch.float32 + ), + latents_curr[:, :, -update_latent_num:, :, :].to( + torch.float32 + ), + eta=0.0, + sigmas=sigma_schedule, + index=i, + prev_sample=None, + grpo=False, + sde_solver=False, + ) + + latents_curr[:, :, -update_latent_num:, :, :] = z + + latents_curr[:, :, -update_latent_num:, :, :] = pred_original + tmp_kwargs["pred_latents"] = ( + latents_curr[:, :, -chunk_latent_num:, :, :].detach().clone() + ) + + with torch.no_grad(): + if latents_curr.shape[2] > chunk_latent_num: + latents_curr = torch.cat( + [ + latents_curr[:, :, :1, :, :], + latents_curr[:, :, -(chunk_latent_num + 2) :, :, :], + ], + dim=2, + ) + + if ( + hasattr(self.vae.config, "shift_factor") + and self.vae.config.shift_factor + ): + latents = ( + latents_curr / self.vae.config.scaling_factor + + self.vae.config.shift_factor + ) + else: + latents = latents_curr / self.vae.config.scaling_factor + + self.vae.to(self.device) + video_frames = self.vae.decode( + latents.to(torch.float32), return_dict=False + )[0] + self.vae.to("cpu") + video_frames = (video_frames / 2 + 0.5).clamp(0, 1).cpu().float() + video_frames = np.transpose(video_frames[0], (1, 2, 3, 0)) + video_frames = (video_frames * 255).numpy().astype(np.uint8) + + if context_video_frames is not None: + video_frames = np.concatenate( + [context_video_frames, video_frames[-16:]], axis=0 + ) + + # Use generated_videos_dir from training_args + generated_videos_base = self.training_args.generated_videos_dir + if not generated_videos_base: + generated_videos_base = os.path.join( + self.training_args.output_dir, "generated_videos" + ) + video_path = os.path.join( + generated_videos_base, + self.training_args.output_dir.split("/")[-1], + f"step_{training_batch.current_timestep}", + ) + + os.makedirs(video_path, exist_ok=True) + + for local_video_idx in range(video_num_in_each_gpu): + # global_video_idx = self.rank_in_sp_group * video_num_in_each_gpu + sample_idx * sampling_batch_size + local_video_idx + global_video_idx = ( + rank_in_gpu_group * video_num_in_each_gpu + + sample_idx * sampling_batch_size + + local_video_idx + ) + + save_video_path = os.path.join( + video_path, + f"{self.global_rank // self.training_args.gpu_para}_{global_video_idx}.mp4", + ) + imageio.mimsave(save_video_path, video_frames) + + absolute_path = os.path.abspath(save_video_path) + + reward_info = self.reward_model.reward( + absolute_path, + gt_camera_pose=w2c_curr, + gt_action=action_curr[:, -update_latent_num:], + caption=training_batch.prompt, + interval=1, + update_latent_num=chunk_latent_num, + ) + action_acc = reward_info["action_acc"] + fine_action_acc = reward_info["fine_action_acc"] + hpsv3_acc = reward_info["hpsv3_acc"] + hpsv3_quality_acc = reward_info["hpsv3_quality_acc"] + hpsv3_quality_drift_score = reward_info[ + "hpsv3_quality_drift_score" + ] + + action_reward = torch.tensor(action_acc).to(latents_curr) + fine_action_reward = torch.tensor(fine_action_acc).to( + latents_curr + ) + hpsv3_reward = torch.tensor(hpsv3_acc).to(latents_curr) + hpsv3_quality_reward = torch.tensor(hpsv3_quality_acc).to( + latents_curr + ) + hpsv3_quality_drift_reward = torch.tensor( + hpsv3_quality_drift_score + ).to(latents_curr) + + new_filename = f"{self.global_rank // self.training_args.gpu_para}_{global_video_idx}_chunk_{selected_chunk_id}_action_{round(action_acc, 1)}_Faction_{round(fine_action_acc, 1)}_hpsv3_{round(hpsv3_acc, 1)}_quality_{round(hpsv3_quality_acc, 1)}_drift_{round(hpsv3_quality_drift_score, 1)}.mp4" + new_absolute_path = os.path.join( + os.path.dirname(absolute_path), new_filename + ) + if os.path.exists(absolute_path): + os.rename(absolute_path, new_absolute_path) + + training_batch.sample_kwargs[f"sample_{global_video_idx}"] = { + "pred_latents": tmp_kwargs["pred_latents"][local_video_idx] + .unsqueeze(0) + .detach() + .clone(), + "action": tmp_kwargs["action"] if self.action else None, + "viewmats": tmp_kwargs["viewmats"] if self.action else None, + "Ks": tmp_kwargs["Ks"] if self.action else None, + "rope_temporal_size": tmp_kwargs["rope_temporal_size"], + "start_rope_start_idx": tmp_kwargs["start_rope_start_idx"], + "cache_vision": tmp_kwargs["cache_vision"], + "return_dict": False, + "action_reward": action_reward, + "fine_action_reward": fine_action_reward, + "hpsv3_reward": hpsv3_reward, + "hpsv3_quality_reward": hpsv3_quality_reward, + "hpsv3_quality_drift_reward": hpsv3_quality_drift_reward, + "selected_chunk_id": selected_chunk_id, + } + + torch.cuda.empty_cache() + training_batch.sigma_schedule = sigma_schedule + training_batch.selected_chunk_id = selected_chunk_id + + return training_batch + + def _prepare_grpo_inputs(self, training_batch: TrainingBatch) -> TrainingBatch: + + action_rewards = [] + fine_action_rewards = [] + hpsv3_rewards = [] + hpsv3_quality_rewards = [] + hpsv3_quality_drift_rewards = [] + + for sample_key in training_batch.sample_kwargs.keys(): + action_rewards.append( + training_batch.sample_kwargs[sample_key]["action_reward"] + ) + fine_action_rewards.append( + training_batch.sample_kwargs[sample_key]["fine_action_reward"] + ) + hpsv3_rewards.append( + training_batch.sample_kwargs[sample_key]["hpsv3_reward"] + ) + hpsv3_quality_rewards.append( + training_batch.sample_kwargs[sample_key]["hpsv3_quality_reward"] + ) + hpsv3_quality_drift_rewards.append( + training_batch.sample_kwargs[sample_key]["hpsv3_quality_drift_reward"] + ) + + action_rewards = torch.tensor(action_rewards).to(self.device) + fine_action_rewards = torch.tensor(fine_action_rewards).to(self.device) + hpsv3_rewards = torch.tensor(hpsv3_rewards).to(self.device) + hpsv3_quality_rewards = torch.tensor(hpsv3_quality_rewards).to(self.device) + hpsv3_quality_drift_rewards = torch.tensor(hpsv3_quality_drift_rewards).to( + self.device + ) + + world_group = get_world_group() + + group_action_rewards = self.gpu_group.all_gather(action_rewards, dim=0) + group_action_rewards_mean = group_action_rewards.mean() + group_action_rewards_std = group_action_rewards.std() + 1e-8 + world_action_rewards = world_group.all_gather(group_action_rewards, dim=0) + world_action_rewards_std = world_group.all_gather( + group_action_rewards_std.unsqueeze(0), dim=0 + ) + + group_fine_action_rewards = self.gpu_group.all_gather( + fine_action_rewards, dim=0 + ) + group_fine_action_rewards_mean = group_fine_action_rewards.mean() + group_fine_action_rewards_std = group_fine_action_rewards.std() + 1e-8 + world_fine_action_rewards = world_group.all_gather( + group_fine_action_rewards, dim=0 + ) + world_fine_action_rewards_std = world_group.all_gather( + group_fine_action_rewards_std.unsqueeze(0), dim=0 + ) + + group_hpsv3_rewards = self.gpu_group.all_gather(hpsv3_rewards, dim=0) + group_hpsv3_rewards_mean = group_hpsv3_rewards.mean() + group_hpsv3_rewards_std = group_hpsv3_rewards.std() + 1e-8 + world_hpsv3_rewards = world_group.all_gather(group_hpsv3_rewards, dim=0) + world_hpsv3_rewards_std = world_group.all_gather( + group_hpsv3_rewards_std.unsqueeze(0), dim=0 + ) + + group_hpsv3_quality_rewards = self.gpu_group.all_gather( + hpsv3_quality_rewards, dim=0 + ) + group_hpsv3_quality_rewards_mean = group_hpsv3_quality_rewards.mean() + group_hpsv3_quality_rewards_std = group_hpsv3_quality_rewards.std() + 1e-8 + world_hpsv3_quality_rewards = world_group.all_gather( + group_hpsv3_quality_rewards, dim=0 + ) + world_hpsv3_quality_rewards_std = world_group.all_gather( + group_hpsv3_quality_rewards_std.unsqueeze(0), dim=0 + ) + + group_hpsv3_quality_drift_rewards = self.gpu_group.all_gather( + hpsv3_quality_drift_rewards, dim=0 + ) + group_hpsv3_quality_drift_rewards_mean = ( + group_hpsv3_quality_drift_rewards.mean() + ) + group_hpsv3_quality_drift_rewards_std = ( + group_hpsv3_quality_drift_rewards.std() + 1e-8 + ) + world_hpsv3_quality_drift_rewards = world_group.all_gather( + group_hpsv3_quality_drift_rewards, dim=0 + ) + world_hpsv3_quality_drift_rewards_std = world_group.all_gather( + group_hpsv3_quality_drift_rewards_std.unsqueeze(0), dim=0 + ) + + if self.training_args.std_type == "sample_max": + group_action_rewards_std = world_action_rewards_std.max() + group_fine_action_rewards_std = world_fine_action_rewards_std.max() + group_hpsv3_rewards_std = world_hpsv3_rewards_std.max() + group_hpsv3_quality_rewards_std = world_hpsv3_quality_rewards_std.max() + group_hpsv3_quality_drift_rewards_std = ( + world_hpsv3_quality_drift_rewards_std.max() + ) + elif self.training_args.std_type == "global": + group_action_rewards_std = world_action_rewards.std() + 1e-8 + group_fine_action_rewards_std = world_fine_action_rewards.std() + 1e-8 + group_hpsv3_rewards_std = world_hpsv3_rewards.std() + 1e-8 + group_hpsv3_quality_rewards_std = world_hpsv3_quality_rewards.std() + 1e-8 + group_hpsv3_quality_drift_rewards_std = ( + world_hpsv3_quality_drift_rewards.std() + 1e-8 + ) + elif self.training_args.std_type == "sample": + group_action_rewards_std = group_action_rewards_std + group_fine_action_rewards_std = group_fine_action_rewards_std + group_hpsv3_rewards_std = group_hpsv3_rewards_std + group_hpsv3_quality_rewards_std = group_hpsv3_quality_rewards_std + group_hpsv3_quality_drift_rewards_std = ( + group_hpsv3_quality_drift_rewards_std + ) + + all_action_rewards = world_group.all_gather(action_rewards, dim=0) + all_fine_action_rewards = world_group.all_gather(fine_action_rewards, dim=0) + all_hpsv3_rewards = world_group.all_gather(hpsv3_rewards, dim=0) + all_hpsv3_quality_rewards = world_group.all_gather(hpsv3_quality_rewards, dim=0) + all_hpsv3_quality_drift_rewards = world_group.all_gather( + hpsv3_quality_drift_rewards, dim=0 + ) + + training_batch.action_reward_mean = all_action_rewards.mean() + training_batch.fine_action_reward_mean = all_fine_action_rewards.mean() + training_batch.hpsv3_reward_mean = all_hpsv3_rewards.mean() + training_batch.hpsv3_quality_reward_mean = all_hpsv3_quality_rewards.mean() + training_batch.hpsv3_quality_drift_reward_mean = ( + all_hpsv3_quality_drift_rewards.mean() + ) + + if self.training_args.action_reward_type == "fine_action": + action_advantages = [] + for sample_key in training_batch.sample_kwargs.keys(): + fine_action_reward = training_batch.sample_kwargs[sample_key][ + "fine_action_reward" + ] + action_advantages.append( + (fine_action_reward - group_fine_action_rewards_mean) + / group_fine_action_rewards_std + ) + training_batch.sample_kwargs[sample_key]["action_advantages"] = ( + fine_action_reward - group_fine_action_rewards_mean + ) / group_fine_action_rewards_std + action_advantages = torch.tensor(action_advantages) + elif self.training_args.action_reward_type == "action": + action_advantages = [] + for sample_key in training_batch.sample_kwargs.keys(): + action_reward = training_batch.sample_kwargs[sample_key][ + "action_reward" + ] + action_advantages.append( + (action_reward - group_action_rewards_mean) + / group_action_rewards_std + ) + training_batch.sample_kwargs[sample_key]["action_advantages"] = ( + action_reward - group_action_rewards_mean + ) / group_action_rewards_std + action_advantages = torch.tensor(action_advantages) + + hpsv3_advantages = [] + for sample_key in training_batch.sample_kwargs.keys(): + hpsv3_reward = training_batch.sample_kwargs[sample_key]["hpsv3_reward"] + hpsv3_advantages.append( + (hpsv3_reward - group_hpsv3_rewards_mean) / group_hpsv3_rewards_std + ) + training_batch.sample_kwargs[sample_key]["hpsv3_advantages"] = ( + hpsv3_reward - group_hpsv3_rewards_mean + ) / group_hpsv3_rewards_std + hpsv3_advantages = torch.tensor(hpsv3_advantages) + + hpsv3_quality_advantages = [] + for sample_key in training_batch.sample_kwargs.keys(): + hpsv3_quality_reward = training_batch.sample_kwargs[sample_key][ + "hpsv3_quality_reward" + ] + hpsv3_quality_advantages.append( + (hpsv3_quality_reward - group_hpsv3_quality_rewards_mean) + / group_hpsv3_quality_rewards_std + ) + training_batch.sample_kwargs[sample_key]["hpsv3_quality_advantages"] = ( + hpsv3_quality_reward - group_hpsv3_quality_rewards_mean + ) / group_hpsv3_quality_rewards_std + hpsv3_quality_advantages = torch.tensor(hpsv3_quality_advantages) + + hpsv3_quality_drift_advantages = [] + for sample_key in training_batch.sample_kwargs.keys(): + hpsv3_quality_drift_reward = training_batch.sample_kwargs[sample_key][ + "hpsv3_quality_drift_reward" + ] + hpsv3_quality_drift_advantages.append( + (hpsv3_quality_drift_reward - group_hpsv3_quality_drift_rewards_mean) + / group_hpsv3_quality_drift_rewards_std + ) + training_batch.sample_kwargs[sample_key][ + "hpsv3_quality_drift_advantages" + ] = ( + hpsv3_quality_drift_reward - group_hpsv3_quality_drift_rewards_mean + ) / group_hpsv3_quality_drift_rewards_std + hpsv3_quality_drift_advantages = torch.tensor(hpsv3_quality_drift_advantages) + + overall_reward = ( + self.training_args.action_reward_weight * action_advantages + + self.training_args.hpsv3_reward_weight * hpsv3_advantages + + self.training_args.hpsv3_quality_reward_weight * hpsv3_quality_advantages + + self.training_args.hpsv3_quality_drift_reward_weight * hpsv3_quality_drift_advantages + ) + + sorted_indices = torch.argsort(overall_reward) + + top_indices = sorted_indices[ + -(self.training_args.bestofn // 2) // self.gpu_group.world_size : + ] + bottom_indices = sorted_indices[ + : (self.training_args.bestofn // 2) // self.gpu_group.world_size + ] + selected_indices = torch.cat([top_indices, bottom_indices]) + + shuffled_order = torch.randperm( + len(selected_indices), device=selected_indices.device + ) + selected_indices = selected_indices[shuffled_order] + training_batch.sample_kwargs["shuffled_sample_keys"] = [ + list(training_batch.sample_kwargs.keys())[idx] for idx in selected_indices + ] + + return training_batch + + def _nft_forward_and_compute_loss( + self, training_batch: TrainingBatch + ) -> TrainingBatch: + self.transformer.train() + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + assert training_batch.attn_metadata is not None + else: + assert training_batch.attn_metadata is None + + with set_forward_context( + current_timestep=training_batch.current_timestep, + attn_metadata=training_batch.attn_metadata, + ): + sigma_schedule = torch.linspace(1, 0, self.training_args.sampling_steps + 1) + shift = 5 + sigma_schedule = (shift * sigma_schedule) / ( + 1 + (shift - 1) * sigma_schedule + ) + training_num = int( + self.training_args.sampling_steps + * self.training_args.train_timestep_fraction + ) + sigma_schedule = sigma_schedule[:-1] + + print( + f"sample_for_training: {len(training_batch.sample_kwargs['shuffled_sample_keys'])}" + ) + for sample_idx, sample_key in tqdm( + list(enumerate(training_batch.sample_kwargs["shuffled_sample_keys"])), + desc="NFT Training Over Samples", + leave=False, + ): + train_sigma_schedule = random.sample( + list(sigma_schedule.cpu().numpy()), training_num + ) + random.shuffle(train_sigma_schedule) + + for timestep_idx in tqdm( + range(len(train_sigma_schedule)), + desc="NFT Training Over Timesteps", + leave=False, + ): + adv_clip_max = self.training_args.adv_clip_max + sigma = train_sigma_schedule[timestep_idx] + + x0 = training_batch.sample_kwargs[sample_key]["pred_latents"] + noise = torch.randn_like(x0) + + device = training_batch.sample_kwargs[sample_key]["viewmats"].device + dtype = training_batch.sample_kwargs[sample_key]["viewmats"].dtype + + update_latent_num = self.training_args.single_chunk_size + + latent_num = x0.shape[2] + stabilization_level = 15 + timestep_value = int(sigma * 1000) + timestep_input = torch.full( + (latent_num,), + stabilization_level - 1, + device=get_local_torch_device(), + dtype=torch.long, + ) + timestep_input[-update_latent_num:] = timestep_value + + noisy_latents = x0.detach().clone() + + noisy_latents = (1 - sigma) * x0 + sigma * noise + noisy_latents = noisy_latents.to(device).to(dtype) + + start_idx = training_batch.sample_kwargs[sample_key][ + "start_rope_start_idx" + ] + end_idx = training_batch.sample_kwargs[sample_key][ + "rope_temporal_size" + ] + + cond_latent_curr = training_batch.cond_latents[ + :, :, start_idx:end_idx, :, : + ] + + latents_for_model = torch.cat( + [noisy_latents, cond_latent_curr], dim=1 + ) + + input_dict = { + "hidden_states": latents_for_model, + "timestep": timestep_input, + "timestep_r": None, + "return_dict": False, + "mask_type": "i2v", + "action": training_batch.sample_kwargs[sample_key]["action"] + if self.action + else None, + "viewmats": training_batch.sample_kwargs[sample_key][ + "viewmats" + ][:, -update_latent_num:, :, :], + "Ks": training_batch.sample_kwargs[sample_key]["Ks"][ + :, -update_latent_num:, :, : + ], + "kv_cache": training_batch.kv_cache["positive"], + "cache_vision": False, + "rope_temporal_size": training_batch.sample_kwargs[sample_key][ + "rope_temporal_size" + ], + "start_rope_start_idx": training_batch.sample_kwargs[ + sample_key + ]["start_rope_start_idx"], + } + + with torch.no_grad(): + with self.ema_generator.apply_policy_shadow_to_model( + self.transformer + ): + model_pred_old = self.transformer( + txt_branch=False, + input_dict=input_dict, + )[0] + + model_pred = self.transformer( + txt_branch=False, + input_dict=input_dict, + )[0] + + action_weight = self.training_args.action_reward_weight + hpsv3_weight = self.training_args.hpsv3_reward_weight + hpsv3_quality_weight = ( + self.training_args.hpsv3_quality_reward_weight + ) + hpsv3_quality_drift_weight = ( + self.training_args.hpsv3_quality_drift_reward_weight + ) + + sum_weights = ( + action_weight + + hpsv3_weight + + hpsv3_quality_weight + + hpsv3_quality_drift_weight + ) + + action_weight_ = action_weight / sum_weights + hpsv3_weight_ = hpsv3_weight / sum_weights + hpsv3_quality_weight_ = hpsv3_quality_weight / sum_weights + hpsv3_quality_drift_weight_ = ( + hpsv3_quality_drift_weight / sum_weights + ) + + total_advantages = ( + action_weight_ + * training_batch.sample_kwargs[sample_key]["action_advantages"] + + hpsv3_weight_ + * training_batch.sample_kwargs[sample_key]["hpsv3_advantages"] + + hpsv3_quality_weight_ + * training_batch.sample_kwargs[sample_key][ + "hpsv3_quality_advantages" + ] + + hpsv3_quality_drift_weight_ + * training_batch.sample_kwargs[sample_key][ + "hpsv3_quality_drift_advantages" + ] + ) + + total_advantages = torch.clamp( + total_advantages, + -adv_clip_max, + adv_clip_max, + ) + normalized_advantages_clip = ( + total_advantages / adv_clip_max + ) / 2.0 + 0.5 + r = torch.clamp(normalized_advantages_clip, 0, 1) + + positive_prediction = model_pred + negative_prediction = 2 * model_pred_old.detach() - model_pred + + positive_x0 = ( + noisy_latents[:, :, -update_latent_num:, :, :] + - sigma * positive_prediction[:, :, -update_latent_num:, :, :] + ) + x0 = x0[:, :, -update_latent_num:, :, :] + with torch.no_grad(): + weight_factor = ( + torch.abs(positive_x0.double() - x0.double()) + .mean(dim=tuple(range(1, positive_x0.ndim)), keepdim=True) + .clip(min=0.00001) + ) + positive_loss = ((positive_x0 - x0) ** 2 / weight_factor).mean( + dim=tuple(range(1, positive_x0.ndim)) + ) + + negative_x0 = ( + noisy_latents[:, :, -update_latent_num:, :, :] + - sigma * negative_prediction[:, :, -update_latent_num:, :, :] + ) + with torch.no_grad(): + weight_factor = ( + torch.abs(negative_x0.double() - x0.double()) + .mean(dim=tuple(range(1, negative_x0.ndim)), keepdim=True) + .clip(min=0.00001) + ) + negative_loss = ((negative_x0 - x0) ** 2 / weight_factor).mean( + dim=tuple(range(1, negative_x0.ndim)) + ) + + policy_loss = r * positive_loss + (1.0 - r) * negative_loss + + final_loss = policy_loss / ( + self.training_args.gradient_accumulation_steps + ) + + # print({ + # 'sample_idx': sample_idx, + # 'timestep_idx': timestep_idx, + # 'action_advantages': round(training_batch.sample_kwargs[sample_key]['action_advantages'].item(), 2), + # 'hpsv3_advantages': round(training_batch.sample_kwargs[sample_key]['hpsv3_advantages'].item(), 2), + # 'total_advantages': round(total_advantages.item(), 2), + # 'r': round(r.item(), 2), + # 'policy_loss': round(policy_loss.item(), 2), + # 'final_loss': round(final_loss.item(), 2), + # }) + # print(f'output_dir: {self.training_args.output_dir}') + + final_loss.backward() + avg_loss = final_loss.detach().clone() + + world_group = get_world_group() + avg_loss = world_group.all_reduce(avg_loss, op=dist.ReduceOp.AVG) + training_batch.total_loss += avg_loss.item() + + if ( + sample_idx + 1 + ) % self.training_args.gradient_accumulation_steps == 0: + training_batch = self._clip_grad_norm(training_batch) + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad() + self.ema_generator.update_ckpt_shadow(self.transformer) + self.ema_generator.update_policy_shadow( + self.transformer, training_batch.current_timestep + ) + + training_batch.samples_grad_norm[sample_key] = ( + training_batch.grad_norm + ) + + return training_batch + + def _clear_cache(self, training_batch: TrainingBatch) -> TrainingBatch: + del training_batch + torch.cuda.empty_cache() + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + training_batch = self._prepare_training(training_batch) + + training_batch = self._get_next_batch(training_batch) + + training_batch = self._sample_reference_model(training_batch) + + training_batch = self._prepare_grpo_inputs(training_batch) + + training_batch = self._nft_forward_and_compute_loss(training_batch) + + self._clear_cache(training_batch) + + training_batch.total_loss = training_batch.total_loss + training_batch.grad_norm = training_batch.grad_norm + return training_batch + + def _resume_from_checkpoint(self) -> None: + logger.info( + "Loading checkpoint from %s", self.training_args.resume_from_checkpoint + ) + resumed_step = load_checkpoint( + self.transformer, + self.global_rank, + self.training_args.resume_from_checkpoint, + self.optimizer, + self.train_dataloader, + self.lr_scheduler, + self.noise_random_generator, + ) + if resumed_step > 0: + self.init_steps = resumed_step + logger.info("Successfully resumed from step %s", resumed_step) + else: + logger.warning("Failed to load checkpoint, starting from step 0") + self.init_steps = 0 + + def get_task_mask(self, task_type, latent_target_length): + if task_type == "t2v": + mask = torch.zeros(latent_target_length) + elif task_type == "i2v": + mask = torch.zeros(latent_target_length) + mask[0] = 1.0 + else: + raise ValueError(f"{task_type} is not supported !") + return mask + + def _prepare_cond_latents(self, task_type, cond_latents, latents, multitask_mask): + """Prepare conditional latents and mask for multitask training. + + Args: + task_type: Type of task ("i2v" or "t2v"). + cond_latents: Conditional latents tensor. + latents: Main latents tensor. + multitask_mask: Multitask mask tensor. + + Returns: + tuple[torch.Tensor, torch.Tensor]: + - latents_concat: Concatenated conditional latents. + - mask_concat: Concatenated mask tensor. + """ + latents_concat = None + mask_concat = None + + if cond_latents is not None and task_type == "i2v": + latents_concat = cond_latents.repeat(1, 1, latents.shape[2], 1, 1) + latents_concat[:, :, 1:, :, :] = 0.0 + else: + latents_concat = torch.zeros( + latents.shape[0], + latents.shape[1], + latents.shape[2], + latents.shape[3], + latents.shape[4], + ).to(latents.device) + + mask_zeros = torch.zeros( + latents.shape[0], 1, latents.shape[2], latents.shape[3], latents.shape[4] + ) + mask_ones = torch.ones( + latents.shape[0], 1, latents.shape[2], latents.shape[3], latents.shape[4] + ) + mask_concat = merge_tensor_by_mask( + mask_zeros.cpu(), mask_ones.cpu(), mask=multitask_mask.cpu(), dim=2 + ).to(device=latents.device) + + cond_latents = torch.concat([latents_concat, mask_concat], dim=1) + + return cond_latents + + def _eval(self, step: int) -> None: + with torch.no_grad(): + action_acc = [] + hpsv3_acc = [] + hpsv3_quality_acc = [] + hpsv3_drift_score = [] + + with self.ema_generator.apply_ckpt_shadow_to_model(self.transformer): + self.vae.eval() + self.transformer.eval() + + eval_batch = TrainingBatch() + + sampling_steps = 20 + latent_t = 32 + + chunk_latent_num = self.training_args.single_chunk_size + video_chunk_num = latent_t // chunk_latent_num + + for batch_idx, batch in enumerate(self.eval_dataloader): + latents = batch["latent"] + prompt = batch["prompt"] + w2c = batch["w2c"] + action = batch["action"] + + bsz, latent_channels, _, latent_h, latent_w = latents.shape + noise = torch.randn( + ( + 1, + latent_channels, + latent_t, + latent_h, + latent_w, + ), # （c,t,h,w) + device=get_local_torch_device(), + dtype=latents.dtype, + ) + + extra_kwargs = { + "byt5_text_states": batch["byt5_text_states"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ), + "byt5_text_mask": batch["byt5_text_mask"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ), + } + multitask_mask = self.get_task_mask("i2v", batch["latent"].shape[2]) + cond_latents = self._prepare_cond_latents( + "i2v", batch["image_cond"], batch["latent"], multitask_mask + ) + + eval_batch.cond_latents = cond_latents.to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.latents_concat = torch.concat( + [batch["latent"], cond_latents], dim=1 + ).to(get_local_torch_device(), dtype=torch.bfloat16) + eval_batch.prompt_embed = batch["prompt_embed"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.prompt_mask = batch["prompt_mask"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.vision_states = batch["vision_states"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.extra_kwargs = extra_kwargs + eval_batch.w2c = batch["w2c"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.intrinsic = batch["intrinsic"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + eval_batch.action = batch["action"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + + latents = self._sample_model_ode( + eval_batch, + video_chunk_num + 1, + noise, + sampling_steps=sampling_steps, + ) + + if ( + hasattr(self.vae.config, "shift_factor") + and self.vae.config.shift_factor + ): + latents = ( + latents / self.vae.config.scaling_factor + + self.vae.config.shift_factor + ) + else: + latents = latents / self.vae.config.scaling_factor + + self.vae.to(self.device) + video_frames = self.vae.decode( + latents.to(torch.float32), return_dict=False + )[0] + self.vae.to("cpu") + video_frames = (video_frames / 2 + 0.5).clamp(0, 1).cpu().float() + video_frames = np.transpose(video_frames[0], (1, 2, 3, 0)) + video_frames = (video_frames * 255).numpy().astype(np.uint8) + + # Use generated_videos_dir from training_args + generated_videos_base = self.training_args.generated_videos_dir + if not generated_videos_base: + generated_videos_base = os.path.join( + self.training_args.output_dir, "generated_videos" + ) + video_path = os.path.join( + generated_videos_base, + self.training_args.output_dir.split("/")[-1], + "000_eval", + f"sample_{self.global_rank + (batch_idx * self.world_size)}", + ) + os.makedirs(video_path, exist_ok=True) + save_video_path = os.path.join(video_path, f"step_{step}.mp4") + print(f"video_frames shape: {video_frames.shape}") + + imageio.mimsave(save_video_path, video_frames) + absolute_path = os.path.abspath(save_video_path) + + score = self.reward_model.score_video( + absolute_path, + caption=prompt, + gt_camera_pose=w2c, + gt_action=action, + interval=1, + latent_num=latent_t, + ) + + new_filename = ( + f"step_{step}" + f"_action_[{torch.tensor(score['action_acc']).mean().item():.1f}]" + f"_hpsv3_[{torch.tensor(score['hps_acc']).mean().item():.1f}]" + f"_quality_[{torch.tensor(score['hps_quality_acc']).mean().item():.1f}]" + f"_drift_[{torch.tensor(score['hps_drift_score']).mean().item():.1f}]" + ".mp4" + ) + + new_absolute_path = os.path.join( + os.path.dirname(absolute_path), new_filename + ) + if os.path.exists(absolute_path): + os.rename(absolute_path, new_absolute_path) + + action_acc += score["action_acc"] + hpsv3_acc += score["hps_acc"] + hpsv3_quality_acc += score["hps_quality_acc"] + hpsv3_drift_score += score["hps_drift_score"] + + world_group = get_world_group() + all_action_acc = world_group.all_gather( + torch.tensor(action_acc).to(get_local_torch_device()).unsqueeze(0), + dim=0, + ) + all_hpsv3_acc = world_group.all_gather( + torch.tensor(hpsv3_acc).to(get_local_torch_device()).unsqueeze(0), + dim=0, + ) + all_hpsv3_quality_acc = world_group.all_gather( + torch.tensor(hpsv3_quality_acc) + .to(get_local_torch_device()) + .unsqueeze(0), + dim=0, + ) + all_hpsv3_drift_score = world_group.all_gather( + torch.tensor(hpsv3_drift_score) + .to(get_local_torch_device()) + .unsqueeze(0), + dim=0, + ) + + ave_action_acc = ( + torch.tensor(action_acc).mean().to(get_local_torch_device()) + ) + ave_hpsv3_acc = ( + torch.tensor(hpsv3_acc).mean().to(get_local_torch_device()) + ) + ave_hpsv3_quality_acc = ( + torch.tensor(hpsv3_quality_acc).mean().to(get_local_torch_device()) + ) + ave_hpsv3_drift_score = ( + torch.tensor(hpsv3_drift_score).mean().to(get_local_torch_device()) + ) + + ave_action_acc = world_group.all_reduce( + ave_action_acc, op=dist.ReduceOp.AVG + ) + ave_hpsv3_acc = world_group.all_reduce( + ave_hpsv3_acc, op=dist.ReduceOp.AVG + ) + ave_hpsv3_quality_acc = world_group.all_reduce( + ave_hpsv3_quality_acc, op=dist.ReduceOp.AVG + ) + ave_hpsv3_drift_score = world_group.all_reduce( + ave_hpsv3_drift_score, op=dist.ReduceOp.AVG + ) + + return { + "ave_action_acc": ave_action_acc.cpu().item(), + "ave_hpsv3_acc": ave_hpsv3_acc.cpu().item(), + "ave_hpsv3_quality_acc": ave_hpsv3_quality_acc.cpu().item(), + "ave_hpsv3_drift_score": ave_hpsv3_drift_score.cpu().item(), + "all_action_acc": all_action_acc.cpu().tolist(), + "all_hpsv3_acc": all_hpsv3_acc.cpu().tolist(), + "all_hpsv3_quality_acc": all_hpsv3_quality_acc.cpu().tolist(), + "all_hpsv3_drift_score": all_hpsv3_drift_score.cpu().tolist(), + } + + def train(self) -> None: + assert self.seed is not None, "seed must be set" + set_random_seed(self.seed + self.global_rank) + logger.info( + "rank: %s: start training", self.global_rank, local_main_process_only=False + ) + + if not self.post_init_called: + self.post_init() + num_trainable_params = _get_trainable_params(self.transformer) + logger.info( + "Starting training with %s B trainable parameters", + round(num_trainable_params / 1e9, 3), + ) + + import subprocess + + kill_cmds = [ + "ps -ef | grep run.py | awk '{print $2}' | head -n -1 | xargs kill -9" + ] + for cmd in kill_cmds: + try: + subprocess.run(cmd, shell=True, check=False) + except Exception as e: + logger.warning(f"Error running kill command: {cmd} -- {e}") + + # Set random seeds for deterministic training + self.noise_random_generator = torch.Generator(device="cpu").manual_seed( + self.seed + ) + self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(self.seed) + self.validation_random_generator = torch.Generator(device="cpu").manual_seed( + self.seed + ) + logger.info("Initialized random seeds with seed: %s", self.seed) + + self.noise_scheduler = FlowMatchEulerDiscreteScheduler() + + if self.training_args.resume_from_checkpoint: + self._resume_from_checkpoint() + + self.train_loader_iter = iter(self.train_dataloader) + + step_times: deque[float] = deque(maxlen=100) + + self._log_training_info() + + self.ema_generator = EMA_FSDP_schedule( + self.transformer, + min_decay=self.training_args.ema_min_decay, + max_decay=self.training_args.ema_max_decay, + step_decay=self.training_args.ema_step_decay, + ckpt_decay=self.training_args.ema_ckpt_decay, + ) + + # Train! + progress_bar = tqdm( + range(0, self.training_args.max_train_steps), + initial=self.init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + disable=self.local_rank > 0, + ) + for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): + if step == 1: + eval_results = self._eval(step - 1) + if self.training_args.eval_only: + print(eval_results) + break + eval_results = self._eval(step - 1) + if self.global_rank == 0: + wandb.log( + { + f"eval_action_acc": eval_results["ave_action_acc"], + f"eval_hpsv3_acc": eval_results["ave_hpsv3_acc"], + f"eval_hpsv3_quality_acc": eval_results[ + "ave_hpsv3_quality_acc" + ], + f"ave_hpsv3_drift_score": eval_results[ + "ave_hpsv3_drift_score" + ], + }, + step=step, + ) + + start_time = time.perf_counter() + if vsa_available: + vsa_sparsity = self.training_args.VSA_sparsity + vsa_decay_rate = self.training_args.VSA_decay_rate + vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps + current_decay_times = min( + step // vsa_decay_interval_steps, vsa_sparsity // vsa_decay_rate + ) + current_vsa_sparsity = current_decay_times * vsa_decay_rate + else: + current_vsa_sparsity = 0.0 + + training_batch = TrainingBatch() + training_batch.current_timestep = step + training_batch.current_vsa_sparsity = current_vsa_sparsity + training_batch = self.train_one_step(training_batch) + + samples_loss = training_batch.samples_loss + samples_grad_norm = training_batch.samples_grad_norm + + step_time = time.perf_counter() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + total_loss = training_batch.total_loss / len(samples_grad_norm) + total_grad_norm = sum(samples_grad_norm.values()) / len(samples_grad_norm) + + progress_bar.set_postfix( + { + "loss": total_loss, + "step_time": f"{step_time:.2f}s", + "grad_norm": total_grad_norm, + } + ) + progress_bar.update(1) + + if self.global_rank == 0: + wandb.log( + { + f"ave_loss": training_batch.total_loss / len(samples_grad_norm), + f"ave_grad_norm": sum(samples_grad_norm.values()) + / len(samples_grad_norm), + }, + step=step, + ) + + wandb.log( + { + "learning_rate": self.lr_scheduler.get_last_lr()[0], + "step_time": step_time, + "avg_step_time": avg_step_time, + "action_reward": training_batch.action_reward_mean, + "fine_action_reward": training_batch.fine_action_reward_mean, + "hpsv3_reward": training_batch.hpsv3_reward_mean, + "hpsv3_quality_reward": training_batch.hpsv3_quality_reward_mean, + }, + step=step, + ) + + training_batch = TrainingBatch() + + if step % self.training_args.checkpointing_steps == 0: + eval_results = self._eval(step - 1) + if self.global_rank == 0: + wandb.log( + { + f"eval_action_acc": eval_results["ave_action_acc"], + f"eval_hpsv3_acc": eval_results["ave_hpsv3_acc"], + f"eval_hpsv3_quality_acc": eval_results[ + "ave_hpsv3_quality_acc" + ], + f"ave_hpsv3_drift_score": eval_results[ + "ave_hpsv3_drift_score" + ], + }, + step=step, + ) + + with self.ema_generator.apply_ckpt_shadow_to_model(self.transformer): + save_checkpoint( + self.transformer, + self.global_rank, + self.training_args.output_dir, + step, + self.optimizer, + self.train_dataloader, + self.lr_scheduler, + self.noise_random_generator, + ) + self.transformer.train() + self.sp_group.barrier() + + wandb.finish() + save_checkpoint( + self.transformer, + self.global_rank, + self.training_args.output_dir, + self.training_args.max_train_steps, + self.optimizer, + self.train_dataloader, + self.lr_scheduler, + self.noise_random_generator, + ) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + def _log_training_info(self) -> None: + total_batch_size = ( + self.world_size + * self.training_args.gradient_accumulation_steps + / self.training_args.sp_size + * self.training_args.train_sp_batch_size + ) + logger.info("***** Running training *****") + logger.info(" Num examples = %s", len(self.train_dataset)) + logger.info(" Dataloader size = %s", len(self.train_dataloader)) + logger.info(" Num Epochs = %s", self.num_train_epochs) + logger.info(" Resume training from step %s", self.init_steps) # type: ignore + logger.info( + " Instantaneous batch size per device = %s", + self.training_args.train_batch_size, + ) + logger.info( + " Total train batch size (w. data & sequence parallel, accumulation) = %s", + total_batch_size, + ) + logger.info( + " Gradient Accumulation steps = %s", + self.training_args.gradient_accumulation_steps, + ) + logger.info( + " Total optimization steps = %s", self.training_args.max_train_steps + ) + logger.info( + " Total training parameters per FSDP shard = %s B", + round(_get_trainable_params(self.transformer) / 1e9, 3), + ) + # print dtype + logger.info( + " Master weight dtype: %s", self.transformer.parameters().__next__().dtype + ) + + gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + logger.info("GPU memory usage before train_one_step: %s MB", gpu_memory_usage) + logger.info("VSA validation sparsity: %s", self.training_args.VSA_sparsity) diff --git a/worldcompass/fastvideo/training/world_compass_train_pipeline.py b/worldcompass/fastvideo/training/world_compass_train_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..74658c179ac755c187a1ca82788c01b5a6a572c7 --- /dev/null +++ b/worldcompass/fastvideo/training/world_compass_train_pipeline.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +from copy import deepcopy +import os +import sys + +sys.path.append(os.path.abspath(".")) + +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.logger import init_logger +from fastvideo.training.world_compass_in_train_pipeline import ( + WorldCompassInTrainPipeline, +) +from fastvideo.utils import is_vsa_available + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class WorldCompassTrainPipeline(WorldCompassInTrainPipeline): + _required_config_modules = ["transformer"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + pass + + def create_training_stages(self, training_args: TrainingArgs): + pass + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + pass + + +def main(args) -> None: + logger.info("Starting training pipeline...") + + pipeline = WorldCompassTrainPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args + ) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args) diff --git a/worldcompass/fastvideo/utils.py b/worldcompass/fastvideo/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8117e7129d7a049766f51daeb45ec92ecc416310 --- /dev/null +++ b/worldcompass/fastvideo/utils.py @@ -0,0 +1,855 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py + +import argparse +import ctypes +import hashlib +import importlib +import importlib.util +import inspect +import json +import math +import os +import signal +import socket +import sys +import tempfile +import threading +import traceback +from collections.abc import Callable +from dataclasses import dataclass, fields, is_dataclass +from functools import lru_cache, partial, wraps +from typing import Any, TypeVar, cast + +import cloudpickle +import filelock +import torch +import yaml +from diffusers.loaders.lora_base import ( + _best_guess_weight_name, +) # watch out for potetential removal from diffusers +from huggingface_hub import snapshot_download +from remote_pdb import RemotePdb +from torch.distributed.fsdp import MixedPrecision + +import fastvideo.envs as envs +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +T = TypeVar("T") + +# TODO(will): used to convert fastvideo_args.precision to torch.dtype. Find a +# cleaner way to do this. +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +STR_BACKEND_ENV_VAR: str = "FASTVIDEO_ATTENTION_BACKEND" +STR_ATTN_CONFIG_ENV_VAR: str = "FASTVIDEO_ATTENTION_CONFIG" + + +def find_nccl_library() -> str: + """We either use the library file specified by the `VLLM_NCCL_SO_PATH` environment variable, or + we find the library file brought by PyTorch. + + After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be + found by `ctypes` automatically. + """ + so_file = envs.FASTVIDEO_NCCL_SO_PATH + + # manually load the nccl library + if so_file: + logger.info( + "Found nccl from environment variable FASTVIDEO_NCCL_SO_PATH=%s", + so_file, + ) + else: + if torch.version.cuda is not None: + so_file = "libnccl.so.2" + elif torch.version.hip is not None: + so_file = "librccl.so.1" + else: + raise ValueError("NCCL only supports CUDA and ROCm backends.") + logger.info("Found nccl from library %s", so_file) + return str(so_file) + + +prev_set_stream = torch.cuda.set_stream + +_current_stream = None + + +def _patched_set_stream(stream: torch.cuda.Stream | None) -> None: + global _current_stream + _current_stream = stream + if stream is not None: + prev_set_stream(stream) + + +torch.cuda.set_stream = _patched_set_stream + + +def current_stream() -> torch.cuda.Stream | None: + """Replace `torch.cuda.current_stream()` with `fastvideo.utils.current_stream()`. it turns out + that `torch.cuda.current_stream()` is quite expensive, as it will construct a new stream object + at each call. here we patch `torch.cuda.set_stream` to keep track of the current stream + directly, so that we can avoid calling `torch.cuda.current_stream()`. + + the underlying hypothesis is that we do not call `torch._C._cuda_setStream` + from C/C++ code. + """ + from fastvideo.platforms import current_platform + + # For non-CUDA platforms, return None + if not current_platform.is_cuda_alike(): + return None + + global _current_stream + if _current_stream is None: + # when this function is called before any stream is set, + # we return the default stream. + # On ROCm using the default 0 stream in combination with RCCL + # is hurting performance. Therefore creating a dedicated stream + # per process + _current_stream = ( + torch.cuda.Stream() + if current_platform.is_rocm() + else torch.cuda.current_stream() + ) + return _current_stream + + +class StoreBoolean(argparse.Action): + def __init__( + self, option_strings, dest, default=False, required=False, help=None + ): + super().__init__( + option_strings=option_strings, + dest=dest, + nargs="?", + const=True, + default=default, + required=required, + help=help, + ) + + def __call__(self, parser, namespace, values, option_string=None): + if values is None: + setattr(namespace, self.dest, True) + elif isinstance(values, str): + if values.lower() == "true": + setattr(namespace, self.dest, True) + elif values.lower() == "false": + setattr(namespace, self.dest, False) + else: + raise ValueError( + f"Invalid boolean value: {values}. Expected 'true' or 'false'." + ) + else: + setattr(namespace, self.dest, bool(values)) + + +class SortedHelpFormatter(argparse.HelpFormatter): + """SortedHelpFormatter that sorts arguments by their option strings.""" + + def add_arguments(self, actions): + actions = sorted(actions, key=lambda x: x.option_strings) + super().add_arguments(actions) + + +class FlexibleArgumentParser(argparse.ArgumentParser): + """ArgumentParser that allows both underscore and dash in names.""" + + def __init__(self, *args, **kwargs) -> None: + # Set the default 'formatter_class' to SortedHelpFormatter + if "formatter_class" not in kwargs: + kwargs["formatter_class"] = SortedHelpFormatter + super().__init__(*args, **kwargs) + + def parse_args( # type: ignore[override] + self, args=None, namespace=None + ) -> argparse.Namespace: + if args is None: + args = sys.argv[1:] + + if "--config" in args: + args = self._pull_args_from_config(args) + + # Convert underscores to dashes and vice versa in argument names + processed_args = [] + for arg in args: + if arg.startswith("--"): + if "=" in arg: + key, value = arg.split("=", 1) + key = "--" + key[len("--") :].replace("_", "-") + processed_args.append(f"{key}={value}") + else: + processed_args.append( + "--" + arg[len("--") :].replace("_", "-") + ) + elif arg.startswith("-O") and arg != "-O" and len(arg) == 2: + # allow -O flag to be used without space, e.g. -O3 + processed_args.append("-O") + processed_args.append(arg[2:]) + else: + processed_args.append(arg) + + namespace = super().parse_args(processed_args, namespace) + + # Track which arguments were explicitly provided + namespace._provided = set() + + i = 0 + while i < len(args): + arg = args[i] + if arg.startswith("--"): + # Handle --key=value format + if "=" in arg: + key = arg.split("=")[0][2:].replace("-", "_") + namespace._provided.add(key) + i += 1 + # Handle --key value format + else: + key = arg[2:].replace("-", "_") + namespace._provided.add(key) + # Skip the value if there is one + if i + 1 < len(args) and not args[i + 1].startswith("-"): + i += 2 + else: + i += 1 + else: + i += 1 + + return namespace # type: ignore[no-any-return] + + def _pull_args_from_config(self, args: list[str]) -> list[str]: + """Method to pull arguments specified in the config file into the command-line args + variable. + + The arguments in config file will be inserted between + the argument list. + + example: + ```yaml + port: 12323 + tensor-parallel-size: 4 + ``` + ```python + $: vllm {serve,chat,complete} "facebook/opt-12B" \ + --config config.yaml -tp 2 + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--config', 'config.yaml', + '-tp', '2' + ] + $: args = [ + "serve,chat,complete", + "facebook/opt-12B", + '--port', '12323', + '--tp-size', '4', + '-tp', '2' + ] + ``` + + Please note how the config args are inserted after the sub command. + this way the order of priorities is maintained when these are args + parsed by super(). + """ + assert ( + args.count("--config") <= 1 + ), "More than one config file specified!" + + index = args.index("--config") + if index == len(args) - 1: + raise ValueError("No config file specified! \ + Please check your command-line arguments.") + + file_path = args[index + 1] + + config_args = self._load_config_file(file_path) + + # 0th index is for {serve,chat,complete} + # followed by model_tag (only for serve) + # followed by config args + # followed by rest of cli args. + # maintaining this order will enforce the precedence + # of cli > config > defaults + if args[0] == "serve": + if index == 1: + raise ValueError( + "No model_tag specified! Please check your command-line arguments." + ) + args = ( + [args[0]] + + [args[1]] + + config_args + + args[2:index] + + args[index + 2 :] + ) + else: + args = [args[0]] + config_args + args[1:index] + args[index + 2 :] + + return args + + def _load_config_file(self, file_path: str) -> list[str]: + """Loads a yaml file and returns the key value pairs as a + flattened list with argparse like pattern + ```yaml + port: 12323 + tensor-parallel-size: 4 + vae_config: + load_encoder: false + load_decoder: true + ``` + returns: + processed_args: list[str] = [ + '--port': '12323', + '--tp-size': '4', + '--vae-config.load-encoder': 'false', + '--vae-config.load-decoder': 'true' + ] + """ + + extension: str = file_path.split(".")[-1] + if extension not in ("yaml", "yml", "json"): + raise ValueError( + "Config file must be of a yaml/yml/json type.\ + %s supplied", + extension, + ) + + processed_args: list[str] = [] + + config: dict[str, Any] = {} + try: + with open(file_path) as config_file: + config = yaml.safe_load(config_file) + except Exception as ex: + logger.error( + "Unable to read the config file at %s. \ + Make sure path is correct", + file_path, + ) + raise ex + + store_boolean_arguments = [ + action.dest + for action in self._actions + if isinstance(action, StoreBoolean) + ] + + def process_dict(prefix: str, d: dict[str, Any]): + for key, value in d.items(): + full_key = f"{prefix}.{key}" if prefix else key + + if ( + isinstance(value, bool) + and full_key not in store_boolean_arguments + ): + if value: + processed_args.append("--" + full_key) + else: + processed_args.append("--" + full_key) + processed_args.append("false") + elif isinstance(value, list): + processed_args.append("--" + full_key) + for item in value: + processed_args.append(str(item)) + elif isinstance(value, dict): + process_dict(full_key, value) + else: + processed_args.append("--" + full_key) + processed_args.append(str(value)) + + process_dict("", config) + + return processed_args + + +def get_lock(model_name_or_path: str): + lock_dir = tempfile.gettempdir() + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +def warn_for_unimplemented_methods(cls: type[T]) -> type[T]: + """A replacement for `abc.ABC`. + + When we use `abc.ABC`, subclasses will fail to instantiate + if they do not implement all abstract methods. + Here, we only require `raise NotImplementedError` in the + base class, and log a warning if the method is not implemented + in the subclass. + """ + + original_init = cls.__init__ + + def find_unimplemented_methods(self: object): + unimplemented_methods = [] + for attr_name in dir(self): + # bypass inner method + if attr_name.startswith("_"): + continue + + try: + attr = getattr(self, attr_name) + # get the func of callable method + if callable(attr): + attr_func = attr.__func__ + except AttributeError: + continue + src = inspect.getsource(attr_func) + if "NotImplementedError" in src: + unimplemented_methods.append(attr_name) + if unimplemented_methods: + method_names = ",".join(unimplemented_methods) + msg = f"Methods {method_names} not implemented in {self}" + logger.warning(msg) + + @wraps(original_init) + def wrapped_init(self, *args, **kwargs) -> None: + original_init(self, *args, **kwargs) + find_unimplemented_methods(self) + + type.__setattr__(cls, "__init__", wrapped_init) + return cls + + +def align_to(value: int, alignment: int) -> int: + """Align height, width according to alignment. + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) + + +def resolve_obj_by_qualname(qualname: str) -> Any: + """Resolve an object by its fully qualified name.""" + module_name, obj_name = qualname.rsplit(".", 1) + module = importlib.import_module(module_name) + return getattr(module, obj_name) + + +# From vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py +def import_pynvml(): + """ + Historical comments: + + libnvml.so is the library behind nvidia-smi, and + pynvml is a Python wrapper around it. We use it to get GPU + status without initializing CUDA context in the current process. + Historically, there are two packages that provide pynvml: + - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official + wrapper. It is a dependency of FastVideo, and is installed when users + install FastVideo. It provides a Python module named `pynvml`. + - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper. + Prior to version 12.0, it also provides a Python module `pynvml`, + and therefore conflicts with the official one which is a standalone Python file. + This causes errors when both of them are installed. + Starting from version 12.0, it migrates to a new module + named `pynvml_utils` to avoid the conflict. + It is so confusing that many packages in the community use the + unofficial one by mistake, and we have to handle this case. + For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial + one, and it will cause errors, see the issue + https://github.com/vllm-project/vllm/issues/12847 for example. + After all the troubles, we decide to copy the official `pynvml` + module to our codebase, and use it directly. + """ + import fastvideo.third_party.pynvml as pynvml + + return pynvml + + +def maybe_download_model( + model_name_or_path: str, local_dir: str | None = None, download: bool = True +) -> str: + """Check if the model path is a Hugging Face Hub model ID and download it if needed. + + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + # If the path exists locally, return it + if os.path.exists(model_name_or_path): + logger.info("Model already exists locally at %s", model_name_or_path) + return model_name_or_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info( + "Downloading model snapshot from HF Hub for %s...", + model_name_or_path, + ) + with get_lock(model_name_or_path): + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir, + ) + logger.info("Downloaded model to %s", local_path) + return str(local_path) + except Exception as e: + raise ValueError( + f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" + ) from e + + +def maybe_download_lora( + model_name_or_path: str, local_dir: str | None = None, download: bool = True +) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the model + download: Whether to download the model from Hugging Face Hub + + Returns: + Local path to the model + """ + + local_path = maybe_download_model(model_name_or_path, local_dir, download) + weight_name = _best_guess_weight_name( + model_name_or_path, file_extension=".safetensors" + ) + return os.path.join(local_path, weight_name) + + +def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: + """Verify that the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary + """ + + # Check for model_index.json which is required for diffusers models + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + raise ValueError( + f"Model directory {model_path} does not contain model_index.json. " + "Only Hugging Face diffusers format is supported." + ) + + # Check for transformer and vae directories + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + + if not os.path.exists(transformer_dir): + raise ValueError( + f"Model directory {model_path} does not contain a transformer/ directory." + ) + + if not os.path.exists(vae_dir): + raise ValueError( + f"Model directory {model_path} does not contain a vae/ directory." + ) + + # Load the config + with open(config_path) as f: + config = json.load(f) + + # Verify diffusers version exists + if "_diffusers_version" not in config: + raise ValueError("model_index.json does not contain _diffusers_version") + + logger.info("Diffusers version: %s", config["_diffusers_version"]) + return cast(dict[str, Any], config) + + +def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: + """Download and extract just the model_index.json for a Hugging Face model. + + Args: + model_name_or_path: Path or HF Hub model ID + + Returns: + The parsed model_index.json as a dictionary + """ + import tempfile + + from huggingface_hub import hf_hub_download + + # If it's a local path, verify it directly + if os.path.exists(model_name_or_path): + return verify_model_config_and_directory(model_name_or_path) + + # For remote models, download just the model_index.json + try: + with tempfile.TemporaryDirectory() as tmp_dir: + # Download just the model_index.json file + model_index_path = hf_hub_download( + repo_id=model_name_or_path, + filename="model_index.json", + local_dir=tmp_dir, + ) + + # Load the model_index.json + with open(model_index_path) as f: + config: dict[str, Any] = json.load(f) + + # Verify it has the required fields + if "_class_name" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _class_name field" + ) + + if "_diffusers_version" not in config: + raise ValueError( + f"model_index.json for {model_name_or_path} does not contain _diffusers_version field" + ) + + # Add the pipeline name for downstream use + config["pipeline_name"] = config["_class_name"] + + logger.info( + "Downloaded model_index.json for %s, pipeline: %s", + model_name_or_path, + config["_class_name"], + ) + return config + + except Exception as e: + raise ValueError( + f"Failed to download or parse model_index.json for {model_name_or_path}: {e}" + ) from e + + +def update_environment_variables(envs: dict[str, str]): + for k, v in envs.items(): + if k in os.environ and os.environ[k] != v: + logger.warning( + "Overwriting environment variable %s from '%s' to '%s'", + k, + os.environ[k], + v, + ) + os.environ[k] = v + + +def run_method( + obj: Any, + method: str | bytes | Callable, + args: tuple[Any], + kwargs: dict[str, Any], +) -> Any: + """Run a method of an object with the given arguments and keyword arguments. + + If the method is string, it will be converted to a method using getattr. If the method is + serialized bytes and will be deserialized using cloudpickle. If the method is a callable, it + will be called directly. + """ + if isinstance(method, bytes): + func = partial(cloudpickle.loads(method), obj) + elif isinstance(method, str): + try: + func = getattr(obj, method) + except AttributeError: + raise NotImplementedError( + f"Method {method!r} is not implemented." + ) from None + else: + func = partial(method, obj) # type: ignore + return func(*args, **kwargs) + + +def shallow_asdict(obj) -> dict[str, Any]: + if not is_dataclass(obj): + raise TypeError("Expected dataclass instance") + return {f.name: getattr(obj, f.name) for f in fields(obj)} + + +# TODO: validate that this is fine +def kill_itself_when_parent_died() -> None: + # if sys.platform == "linux": + # sigkill this process when parent worker manager dies + PR_SET_PDEATHSIG = 1 + import platform + + if platform.system() == "Linux": + libc = ctypes.CDLL("libc.so.6") + libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) + # elif platform.system() == "Darwin": + # libc = ctypes.CDLL("libc.dylib") + # logger.warning("kill_itself_when_parent_died is only supported in linux.") + else: + logger.warning( + "kill_itself_when_parent_died is only supported in linux." + ) + + +def get_exception_traceback() -> str: + etype, value, tb = sys.exc_info() + err_str = "".join(traceback.format_exception(etype, value, tb)) + return err_str + + +class TypeBasedDispatcher: + def __init__(self, mapping: list[tuple[type, Callable]]): + self._mapping = mapping + + def __call__(self, obj: Any): + for ty, fn in self._mapping: + if isinstance(obj, ty): + return fn(obj) + raise ValueError(f"Invalid object: {obj}") + + +# For non-torch.distributed debugging +def remote_breakpoint() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("localhost", 0)) # Let the OS pick an ephemeral port. + port = s.getsockname()[1] + RemotePdb(host="localhost", port=port).set_trace() + + +@dataclass +class MixedPrecisionState: + param_dtype: torch.dtype | None = None + reduce_dtype: torch.dtype | None = None + output_dtype: torch.dtype | None = None + compute_dtype: torch.dtype | None = None + mp_policy: MixedPrecision | None = None + + +# Thread-local storage for mixed precision state +_mixed_precision_state = threading.local() + + +def get_mixed_precision_state() -> MixedPrecisionState: + """Get the current mixed precision state.""" + if not hasattr(_mixed_precision_state, "state"): + raise ValueError("Mixed precision state not set") + return cast(MixedPrecisionState, _mixed_precision_state.state) + + +def set_mixed_precision_policy( + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + output_dtype: torch.dtype | None = None, + mp_policy: MixedPrecision | None = None, +): + """Set mixed precision policy globally. + + Args: + param_dtype: Parameter dtype used for training + reduce_dtype: Reduction dtype used for gradients + output_dtype: Optional output dtype + """ + state = MixedPrecisionState( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + _mixed_precision_state.state = state + + +def get_compute_dtype() -> torch.dtype: + """Get the current compute dtype from mixed precision policy. + + Returns: + torch.dtype: The compute dtype to use, defaults to get_default_dtype() if no policy set + """ + if not hasattr(_mixed_precision_state, "state"): + return torch.get_default_dtype() + else: + state = get_mixed_precision_state() + return state.param_dtype + + +def dict_to_3d_list( + mask_strategy: dict[str, Any] | None = None, + t_max: int | None = None, + l_max: int | None = None, + h_max: int | None = None, +) -> list[list[list[torch.Tensor | None]]]: + """Convert a dictionary of mask indices to a 3D list of tensors. + + Args: + mask_strategy: keys are "t_l_h", values are torch.Tensor masks. + t_max, l_max, h_max: if provided (all three), force the output shape to (t_max, l_max, h_max). + If all three are None, infer shape from the data. + """ + # Case 1: no data, but fixed shape requested + if mask_strategy is None: + assert ( + t_max is not None and l_max is not None and h_max is not None + ), "If mask_strategy is None, you must provide t_max, l_max, and h_max" + return [ + [[None for _ in range(h_max)] for _ in range(l_max)] + for _ in range(t_max) + ] + + # Parse all keys into integer tuples + indices = [tuple(map(int, key.split("_"))) for key in mask_strategy] + + # Decide on dimensions + if t_max is None and l_max is None and h_max is None: + # fully dynamic: infer from data + max_timesteps_idx = max(t for t, _, _ in indices) + 1 + max_layer_idx = max(l for _, l, _ in indices) + 1 # noqa: E741 + max_head_idx = max(h for _, _, h in indices) + 1 + else: + # require all three to be provided + assert t_max is not None and l_max is not None and h_max is not None, ( + "Either supply none of (t_max, l_max, h_max) to infer dimensions, " + "or supply all three to fix the shape." + ) + max_timesteps_idx = t_max + max_layer_idx = l_max + max_head_idx = h_max + + # Preallocate + result = [ + [[None for _ in range(max_head_idx)] for _ in range(max_layer_idx)] + for _ in range(max_timesteps_idx) + ] + + # Fill in, skipping any out-of-bounds entries + for key, value in mask_strategy.items(): + t, l, h = map(int, key.split("_")) # noqa: E741 + if ( + 0 <= t < max_timesteps_idx + and 0 <= l < max_layer_idx + and 0 <= h < max_head_idx + ): + result[t][l][h] = value + # else: silently ignore any key that doesn't fit + + return result + + +def set_random_seed(seed: int) -> None: + from fastvideo.platforms import current_platform + + current_platform.seed_everything(seed) + + +@lru_cache(maxsize=1) +def is_vsa_available() -> bool: + return importlib.util.find_spec("vsa") is not None diff --git a/worldcompass/fastvideo/version.py b/worldcompass/fastvideo/version.py new file mode 100644 index 0000000000000000000000000000000000000000..1276d0254ffc2b5e4fcf48cd868123238d4ad06f --- /dev/null +++ b/worldcompass/fastvideo/version.py @@ -0,0 +1 @@ +__version__ = "0.1.5" diff --git a/worldcompass/fastvideo/worker/__init__.py b/worldcompass/fastvideo/worker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..535f558a4d38999a117bfcfef32d5291be8b3ccb --- /dev/null +++ b/worldcompass/fastvideo/worker/__init__.py @@ -0,0 +1,5 @@ +from .executor import Executor +from .gpu_worker import run_worker_process +from .multiproc_executor import MultiprocExecutor + +__all__ = ["Executor", "run_worker_process", "MultiprocExecutor"] diff --git a/worldcompass/fastvideo/worker/executor.py b/worldcompass/fastvideo/worker/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..ceea9776a9fc02cc0e567d2398c393bd5cbc6175 --- /dev/null +++ b/worldcompass/fastvideo/worker/executor.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, TypeVar, cast + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines import ForwardBatch +from fastvideo.utils import init_logger + +logger = init_logger(__name__) + +_R = TypeVar("_R") + + +class Executor(ABC): + def __init__(self, fastvideo_args: FastVideoArgs): + self.fastvideo_args = fastvideo_args + + self._init_executor() + + @abstractmethod + def _init_executor(self) -> None: + raise NotImplementedError + + @classmethod + def get_class(cls, fastvideo_args: FastVideoArgs) -> type["Executor"]: + if fastvideo_args.distributed_executor_backend == "mp": + from fastvideo.worker.multiproc_executor import MultiprocExecutor + + return cast(type["Executor"], MultiprocExecutor) + else: + raise ValueError( + f"Unsupported distributed executor backend: {fastvideo_args.distributed_executor_backend}" + ) + + def execute_forward( + self, + forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + outputs: list[dict[str, Any]] = self.collective_rpc( + "execute_forward", + kwargs={ + "forward_batch": forward_batch, + "fastvideo_args": fastvideo_args, + }, + ) + return cast(ForwardBatch, outputs[0]["output_batch"]) + + @abstractmethod + def set_lora_adapter( + self, lora_nickname: str, lora_path: str | None = None + ) -> None: + """Set the LoRA adapter for the workers.""" + raise NotImplementedError + + @abstractmethod + def collective_rpc( + self, + method: str | Callable[..., _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ) -> list[_R]: + """Execute an RPC call on all workers. + + Args: + method: Name of the worker method to execute, or a callable that + is serialized and sent to all workers to execute. + + If the method is a callable, it should accept an additional + `self` argument, in addition to the arguments passed in `args` + and `kwargs`. The `self` argument will be the worker object. + timeout: Maximum time in seconds to wait for execution. Raises a + :exc:`TimeoutError` on timeout. `None` means wait indefinitely. + args: Positional arguments to pass to the worker method. + kwargs: Keyword arguments to pass to the worker method. + + Returns: + A list containing the results from each worker. + + Note: + It is recommended to use this API to only pass control messages, + and set up data-plane communication to pass data. + """ + raise NotImplementedError + + @abstractmethod + def shutdown(self) -> None: + """Shutdown the executor.""" + raise NotImplementedError diff --git a/worldcompass/fastvideo/worker/gpu_worker.py b/worldcompass/fastvideo/worker/gpu_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..fbbedaa2590bb30191e3aa080d4df5e30de97ee3 --- /dev/null +++ b/worldcompass/fastvideo/worker/gpu_worker.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +import contextlib +import faulthandler +import multiprocessing as mp +import os +import signal +import sys +from multiprocessing.connection import Connection +from typing import Any, TextIO, cast + +import psutil +import torch + +import fastvideo.envs as envs +from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, +) +from fastvideo.distributed.parallel_state import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ForwardBatch, build_pipeline +from fastvideo.platforms import current_platform +from fastvideo.utils import ( + get_exception_traceback, + kill_itself_when_parent_died, +) + +logger = init_logger(__name__) + +# ANSI color codes +CYAN = "\033[1;36m" +RESET = "\033[0;0m" + + +class Worker: + def __init__( + self, + fastvideo_args: FastVideoArgs, + local_rank: int, + rank: int, + pipe: Connection, + master_port: int, + ): + self.fastvideo_args = fastvideo_args + self.local_rank = local_rank + self.rank = rank + # TODO(will): don't hardcode this + self.distributed_init_method = "env://" + self.pipe = pipe + self.master_port = master_port + self.init_device() + + # Init request dispatcher + # TODO(will): add request dispatcher: use TypeBasedDispatcher from + # utils.py + # self._request_dispatcher = TypeBasedDispatcher( + # [ + # (RpcReqInput, self.handle_rpc_request), + # (GenerateRequest, self.handle_generate_request), + # (ExpertDistributionReq, self.expert_distribution_handle), + # ] + # ) + + def init_device(self) -> None: + """Initialize the device for the worker.""" + + # torch.distributed.all_reduce does not free the input tensor until + # the synchronization point. This causes the memory usage to grow + # as the number of all_reduce calls increases. This env var disables + # this behavior. + # Related issue: + # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + # This env var set by Ray causes exceptions with graph building. + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + + # Platform-agnostic device initialization + self.device = get_local_torch_device() + + # _check_if_gpu_supports_dtype(self.model_config.dtype) + if current_platform.is_cuda_alike(): + self.init_gpu_memory = torch.cuda.mem_get_info()[0] + else: + # For MPS, we can't get memory info the same way + self.init_gpu_memory = 0 + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(self.master_port) + os.environ["LOCAL_RANK"] = str(self.local_rank) + os.environ["RANK"] = str(self.rank) + os.environ["WORLD_SIZE"] = str(self.fastvideo_args.num_gpus) + + # Initialize the distributed environment. + maybe_init_distributed_environment_and_model_parallel( + self.fastvideo_args.tp_size, self.fastvideo_args.sp_size + ) + + self.pipeline = build_pipeline(self.fastvideo_args) + + def execute_forward( + self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs + ) -> ForwardBatch: + output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) + return cast(ForwardBatch, output_batch) + + def set_lora_adapter( + self, lora_nickname: str, lora_path: str | None = None + ) -> None: + self.pipeline.set_lora_adapter(lora_nickname, lora_path) + + def shutdown(self) -> dict[str, Any]: + """Gracefully shut down the worker process.""" + logger.info( + "Worker %d shutting down...", + self.rank, + local_main_process_only=False, + ) + # Clean up resources + if hasattr(self, "pipeline") and self.pipeline is not None: + # Clean up pipeline resources if needed + pass + + # Destroy the distributed environment + cleanup_dist_env_and_memory(shutdown_ray=False) + + logger.info( + "Worker %d shutdown complete", + self.rank, + local_main_process_only=False, + ) + return {"status": "shutdown_complete"} + + def event_loop(self) -> None: + """Event loop for the worker.""" + logger.info( + "Worker %d starting event loop...", + self.rank, + local_main_process_only=False, + ) + while True: + try: + recv_rpc = self.pipe.recv() + method_name = recv_rpc.get("method") + + # Handle shutdown request + if method_name == "shutdown": + response = self.shutdown() + with contextlib.suppress(Exception): + self.pipe.send(response) + break # Exit the loop + + # Handle regular RPC calls + if method_name == "execute_forward": + forward_batch = recv_rpc["kwargs"]["forward_batch"] + fastvideo_args = recv_rpc["kwargs"]["fastvideo_args"] + output_batch = self.execute_forward( + forward_batch, fastvideo_args + ) + logging_info = None + if envs.FASTVIDEO_STAGE_LOGGING: + logging_info = output_batch.logging_info + self.pipe.send( + { + "output_batch": output_batch.output.cpu(), + "logging_info": logging_info, + } + ) + elif method_name == "set_lora_adapter": + lora_nickname = recv_rpc["kwargs"]["lora_nickname"] + lora_path = recv_rpc["kwargs"]["lora_path"] + self.set_lora_adapter(lora_nickname, lora_path) + logger.info( + "Worker %d set LoRA adapter %s with path %s", + self.rank, + lora_nickname, + lora_path, + ) + self.pipe.send({"status": "lora_adapter_set"}) + else: + # Handle other methods dynamically if needed + args = recv_rpc.get("args", ()) + kwargs = recv_rpc.get("kwargs", {}) + if hasattr(self, method_name): + method = getattr(self, method_name) + result = method(*args, **kwargs) + self.pipe.send(result) + else: + self.pipe.send( + {"error": f"Unknown method: {method_name}"} + ) + except KeyboardInterrupt: + logger.error( + "Worker %d in loop received KeyboardInterrupt, aborting forward pass", + self.rank, + ) + try: + self.pipe.send( + {"error": "Operation aborted by KeyboardInterrupt"} + ) + logger.info( + "Worker %d sent error response after interrupt", + self.rank, + ) + except Exception as e: + logger.error( + "Worker %d failed to send error response: %s", + self.rank, + str(e), + ) + continue + + +def run_worker_process( + fastvideo_args: FastVideoArgs, + local_rank: int, + rank: int, + pipe: Connection, + master_port: int, +): + # Add process-specific prefix to stdout and stderr + process_name = mp.current_process().name + pid = os.getpid() + _add_prefix(sys.stdout, process_name, pid) + _add_prefix(sys.stderr, process_name, pid) + + # Config the process + kill_itself_when_parent_died() + faulthandler.enable() + parent_process = psutil.Process().parent() + + logger.info( + "Worker %d initializing...", rank, local_main_process_only=False + ) + + try: + worker = Worker(fastvideo_args, local_rank, rank, pipe, master_port) + logger.info("Worker %d sending ready", rank) + pipe.send( + { + "status": "ready", + "local_rank": local_rank, + } + ) + worker.event_loop() + except Exception: + traceback = get_exception_traceback() + logger.error("Worker %d hit an exception: %s", rank, traceback) + parent_process.send_signal(signal.SIGQUIT) + + +def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None: + """Prepend each output line with process-specific prefix.""" + + prefix = f"{CYAN}({worker_name} pid={pid}){RESET} " + file_write = file.write + + def write_with_prefix(s: str): + if not s: + return + if file.start_new_line: # type: ignore[attr-defined] + file_write(prefix) + idx = 0 + while (next_idx := s.find("\n", idx)) != -1: + next_idx += 1 + file_write(s[idx:next_idx]) + if next_idx == len(s): + file.start_new_line = True # type: ignore[attr-defined] + return + file_write(prefix) + idx = next_idx + file_write(s[idx:]) + file.start_new_line = False # type: ignore[attr-defined] + + file.start_new_line = True # type: ignore[attr-defined] + file.write = write_with_prefix # type: ignore[method-assign] diff --git a/worldcompass/fastvideo/worker/multiproc_executor.py b/worldcompass/fastvideo/worker/multiproc_executor.py new file mode 100644 index 0000000000000000000000000000000000000000..ba2f872355f834039f1feeb8a29ec5ec816d9f5d --- /dev/null +++ b/worldcompass/fastvideo/worker/multiproc_executor.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +import atexit +import contextlib +import multiprocessing as mp +import os +import signal +import socket +import time +from collections.abc import Callable +from multiprocessing.process import BaseProcess +from typing import Any + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.worker.executor import Executor +from fastvideo.worker.gpu_worker import run_worker_process + +logger = init_logger(__name__) + + +class MultiprocExecutor(Executor): + def _init_executor(self) -> None: + self.world_size = self.fastvideo_args.num_gpus + self.shutting_down = False + + # this will force the use of the `spawn` multiprocessing start if cuda + # is initialized + mp.set_start_method("spawn", force=True) + + self.workers: list[BaseProcess] = [] + self.worker_pipes = [] + + # Check if master_port is provided in fastvideo_args + if ( + hasattr(self.fastvideo_args, "master_port") + and self.fastvideo_args.master_port is not None + ): + self.master_port = self.fastvideo_args.master_port + logger.info("Using provided master port: %s", self.master_port) + else: + # Auto-find available port + import random + + for port in range(29503 + random.randint(0, 10000), 65535): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + if s.connect_ex(("localhost", port)) != 0: + self.master_port = port + break + else: + raise ValueError("No unused port found to use as master port") + logger.info("Auto-selected master port: %s", self.master_port) + + # Create pipes and start workers + for rank in range(self.world_size): + executor_pipe, worker_pipe = mp.Pipe(duplex=True) + self.worker_pipes.append(executor_pipe) + worker = mp.Process( + target=run_worker_process, + name=f"FVWorkerProc-{rank}", + kwargs=dict( + fastvideo_args=self.fastvideo_args, + local_rank=rank, + rank=rank, + pipe=worker_pipe, + master_port=self.master_port, + ), + ) + worker.start() + self.workers.append(worker) + + # Wait for all workers to be ready + for idx, pipe in enumerate(self.worker_pipes): + data = pipe.recv() + if data["status"] != "ready" or data["local_rank"] != idx: + raise RuntimeError(f"Worker {idx} failed to start") + logger.info("%d workers ready", self.world_size) + + # Register shutdown on exit + atexit.register(self.shutdown) + + def execute_forward( + self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs + ) -> ForwardBatch: + responses = self.collective_rpc( + "execute_forward", + kwargs={ + "forward_batch": forward_batch, + "fastvideo_args": fastvideo_args, + }, + ) + output = responses[0]["output_batch"] + + logging_info = None + if envs.FASTVIDEO_STAGE_LOGGING: + logging_info = responses[0]["logging_info"] + else: + logging_info = None + + result_batch = ForwardBatch( + data_type=forward_batch.data_type, + output=output, + logging_info=logging_info, + ) + + return result_batch + + def set_lora_adapter( + self, lora_nickname: str, lora_path: str | None = None + ) -> None: + responses = self.collective_rpc( + "set_lora_adapter", + kwargs={"lora_nickname": lora_nickname, "lora_path": lora_path}, + ) + for i, response in enumerate(responses): + if response["status"] != "lora_adapter_set": + raise RuntimeError( + f"Worker {i} failed to set LoRA adapter to {lora_path}" + ) + + def collective_rpc( + self, + method: str | Callable, + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None, + ) -> list[Any]: + kwargs = kwargs or {} + + try: + for pipe in self.worker_pipes: + pipe.send({"method": method, "args": args, "kwargs": kwargs}) + + responses = [] + for pipe in self.worker_pipes: + response = pipe.recv() + responses.append(response) + return responses + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e + except KeyboardInterrupt as e: + # if we catch a KeyboardInterrupt, user wants to stop the execution. + # we need to send a signal to all workers to stop. + logger.info( + "Received KeyboardInterrupt, sending SIGINT to all workers" + ) + for worker in self.workers: + if worker.pid is not None: + os.kill(worker.pid, signal.SIGINT) + raise e + except Exception as e: + raise e + + def shutdown(self) -> None: + """Properly shut down the executor and its workers.""" + if hasattr(self, "shutting_down") and self.shutting_down: + return # Prevent multiple shutdown calls + + logger.info("Shutting down MultiprocExecutor...") + self.shutting_down = True + + # First try gentle termination + try: + # Send termination message to all workers + for pipe in self.worker_pipes: + with contextlib.suppress(Exception): + pipe.send({"method": "shutdown", "args": (), "kwargs": {}}) + + # Give workers some time to exit gracefully + start_time = time.perf_counter() + while time.perf_counter() - start_time < 5.0: # 5 seconds timeout + if all(not worker.is_alive() for worker in self.workers): + break + time.sleep(0.1) + + # Force terminate any remaining workers + for worker in self.workers: + if worker.is_alive(): + worker.terminate() + + # Final timeout for terminate + start_time = time.perf_counter() + while time.perf_counter() - start_time < 2.0: # 2 seconds timeout + if all(not worker.is_alive() for worker in self.workers): + break + time.sleep(0.1) + + # Kill if still alive + for worker in self.workers: + if worker.is_alive(): + worker.kill() + worker.join(timeout=1.0) + + except Exception as e: + logger.error("Error during shutdown: %s", e) + # Last resort, try to kill all workers + for worker in self.workers: + with contextlib.suppress(Exception): + if worker.is_alive(): + worker.kill() + + # Clean up pipes + for pipe in self.worker_pipes: + with contextlib.suppress(Exception): + pipe.close() + + self.workers = [] + self.worker_pipes = [] + logger.info("MultiprocExecutor shutdown complete") + + def __del__(self): + """Ensure cleanup on garbage collection.""" + self.shutdown() + + def __enter__(self): + """Support for context manager protocol.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Ensure cleanup when exiting context.""" + self.shutdown() diff --git a/worldcompass/prepare_dataset/prepare_custom_action.py b/worldcompass/prepare_dataset/prepare_custom_action.py new file mode 100644 index 0000000000000000000000000000000000000000..40c9a5e668e577324a054df44ea3c79bbb733bbd --- /dev/null +++ b/worldcompass/prepare_dataset/prepare_custom_action.py @@ -0,0 +1,298 @@ +"""Action to Pose Conversion Module Converts camera motion actions (forward, backward, left, right, +yaw, pitch) into camera poses represented as 4x4 transformation matrices for video generation +training.""" + +import numpy as np +import json +import torch +from scipy.spatial.transform import Rotation as R +import random +from tqdm import tqdm + + +def rot_x(theta): + """Generate rotation matrix around X-axis (pitch). + + Args: + theta: Rotation angle in radians + + Returns: + 3x3 rotation matrix for X-axis rotation + """ + c, s = np.cos(theta), np.sin(theta) + return np.array([[1, 0, 0], [0, c, -s], [0, s, c]]) + + +def rot_y(theta): + """Generate rotation matrix around Y-axis (yaw). + + Args: + theta: Rotation angle in radians + + Returns: + 3x3 rotation matrix for Y-axis rotation + """ + c, s = np.cos(theta), np.sin(theta) + return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]) + + +def rot_z(theta): + """Generate rotation matrix around Z-axis (roll). + + Args: + theta: Rotation angle in radians + + Returns: + 3x3 rotation matrix for Z-axis rotation + """ + c, s = np.cos(theta), np.sin(theta) + return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + + +def action_to_pose(motions, existing_pose=None): + """Convert a sequence of camera motion actions into camera poses. + + Args: + motions: List of motion dictionaries, each containing: + - "forward": Move along camera's positive Z-axis + - "backward": Move along camera's negative Z-axis + - "right": Move along camera's positive X-axis + - "left": Move along camera's negative X-axis + - "yaw": Rotate around camera's Y-axis (OpenCV convention) + - "pitch": Rotate around camera's X-axis + - "roll": Rotate around camera's Z-axis + Example: [{"forward": 1.0}, {"yaw": np.pi/2}, {"pitch": np.pi/6}] + existing_pose: Optional list of existing poses. If provided, the transformation + continues from the last pose in the list. + + Returns: + List of 4x4 transformation matrices representing camera poses for each action. + Each matrix is in world coordinates. + """ + # Initialize transformation matrix + if existing_pose is not None: + T = existing_pose[-1].copy() + else: + T = np.eye(4) # Start with identity matrix (origin) + + poses = [] + + for move in motions: + # Apply rotations first (in camera's local coordinate system) + if "yaw" in move: + R = rot_y(move["yaw"]) + T[:3, :3] = T[:3, :3] @ R + if "pitch" in move: + R = rot_x(move["pitch"]) + T[:3, :3] = T[:3, :3] @ R + + # Apply translations (in camera's local coordinate system, then transform to world) + forward = move.get("forward", 0.0) + right = move.get("right", 0.0) + backward = move.get("backward", 0.0) + left = move.get("left", 0.0) + + # Forward: positive Z in camera space + if forward != 0: + local_t = np.array([0, 0, forward]) + world_t = T[:3, :3] @ local_t + T[:3, 3] += world_t + + # Right: positive X in camera space + if right != 0: + local_t = np.array([right, 0, 0]) + world_t = T[:3, :3] @ local_t + T[:3, 3] += world_t + + # Backward: negative Z in camera space + if backward != 0: + local_t = np.array([0, 0, -backward]) + world_t = T[:3, :3] @ local_t + T[:3, 3] += world_t + + # Left: negative X in camera space + if left != 0: + local_t = np.array([-left, 0, 0]) + world_t = T[:3, :3] @ local_t + T[:3, 3] += world_t + + poses.append(T.copy()) + + return poses + + +def one_hot_to_one_dimension(one_hot): + """Convert one-hot encoded action vectors to single dimension labels. + + Args: + one_hot: Tensor of one-hot encoded vectors with 4 dimensions representing + [forward, backward, left, right] and their combinations + + Returns: + Tensor of integer labels (0-8) representing action classes: + 0: no action, 1: forward, 2: backward, 3: left, 4: right, + 5: forward+left, 6: forward+right, 7: backward+left, 8: backward+right + """ + mapping = { + (0, 0, 0, 0): 0, # No action + (1, 0, 0, 0): 1, # Forward only + (0, 1, 0, 0): 2, # Backward only + (0, 0, 1, 0): 3, # Left only + (0, 0, 0, 1): 4, # Right only + (1, 0, 1, 0): 5, # Forward + Left + (1, 0, 0, 1): 6, # Forward + Right + (0, 1, 1, 0): 7, # Backward + Left + (0, 1, 0, 1): 8, # Backward + Right + } + y = torch.tensor([mapping[tuple(row.tolist())] for row in one_hot]) + return y + + +def generate_revisit_action(num): + """ + Generate a revisit trajectory pattern: turn left for first half, then turn right. + Useful for training camera return-to-origin behaviors. + + Args: + num: Total number of actions to generate + + Returns: + List of action dictionaries with yaw rotations + """ + trajectory = [] + ANGLE = np.deg2rad(2) # 2-degree rotation per step + + # First half: turn left (negative yaw) + for _ in range(num // 2): + trajectory.append({"yaw": -ANGLE}) + + # Second half: turn right (positive yaw) + for _ in range(num - num // 2): + trajectory.append({"yaw": ANGLE}) + + return trajectory + + +def generate_random_action(num): + """Generates a single, unique basic trajectory of 50 steps with fixed parameters.""" + trajectory = [] + + ANGLE = np.deg2rad(2) + STEP = 0.1 + + basic_chunks = [ + # 单一平移，权重2 + {"forward": STEP}, + {"backward": STEP}, + {"left": STEP}, + {"right": STEP}, + # 单一旋转，权重2 + {"yaw": ANGLE}, + {"yaw": -ANGLE}, + {"pitch": ANGLE}, + {"pitch": -ANGLE}, + # 纯平移组合，权重4 + {"forward": STEP, "right": STEP}, + {"forward": STEP, "left": STEP}, + {"backward": STEP, "right": STEP}, + {"backward": STEP, "left": STEP}, + # 纯旋转组合，权重4 + {"yaw": ANGLE, "pitch": ANGLE}, + {"yaw": ANGLE, "pitch": -ANGLE}, + {"yaw": -ANGLE, "pitch": ANGLE}, + {"yaw": -ANGLE, "pitch": -ANGLE}, + # 平移与偏航组合，权重4 + {"forward": STEP, "yaw": ANGLE}, + {"forward": STEP, "yaw": -ANGLE}, + {"backward": STEP, "yaw": ANGLE}, + {"backward": STEP, "yaw": -ANGLE}, + {"left": STEP, "yaw": ANGLE}, + {"left": STEP, "yaw": -ANGLE}, + {"right": STEP, "yaw": ANGLE}, + {"right": STEP, "yaw": -ANGLE}, + # 平移与俯仰组合，权重4 + {"forward": STEP, "pitch": ANGLE}, + {"forward": STEP, "pitch": -ANGLE}, + {"backward": STEP, "pitch": ANGLE}, + {"backward": STEP, "pitch": -ANGLE}, + {"left": STEP, "pitch": ANGLE}, + {"left": STEP, "pitch": -ANGLE}, + {"right": STEP, "pitch": ANGLE}, + {"right": STEP, "pitch": -ANGLE}, + ] + + weights = [ + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 6, + 6, + 6, + 6, + 2, + 2, + 2, + 2, + 6, + 6, + 6, + 6, + 6, + 6, + 6, + 6, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + ] + + trajectory = [] + total_steps = 0 + while total_steps < num: + chunk_to_add = random.choices(basic_chunks, weights=weights, k=1)[0] + + num_steps = random.randint(2, 16) + + num_steps = min(num_steps, num - total_steps) + + for _ in range(num_steps): + trajectory.append(chunk_to_add.copy()) + + total_steps += num_steps + + return trajectory + + +intrinsic = [ + [969.6969696969696, 0.0, 960.0], + [0.0, 969.6969696969696, 540.0], + [0.0, 0.0, 1.0], +] + +if __name__ == "__main__": + all_rand_poses = [] + for i in tqdm(range(1000)): + random_action = generate_random_action(128) + pose = action_to_pose(random_action) + custom_w2c = {} + for i, p in enumerate(pose): + custom_w2c[str(i)] = {"extrinsic": p.tolist(), "K": intrinsic} + + all_rand_poses.append(custom_w2c) + + json.dump( + all_rand_poses, + open(f"./dataset/harder_random_poses.json", "w"), + indent=4, + ensure_ascii=False, + ) diff --git a/worldcompass/prepare_dataset/prepare_image_text_latent_simple.py b/worldcompass/prepare_dataset/prepare_image_text_latent_simple.py new file mode 100644 index 0000000000000000000000000000000000000000..01940294e097bc955f06e6f40e6f58a285d0bc6d --- /dev/null +++ b/worldcompass/prepare_dataset/prepare_image_text_latent_simple.py @@ -0,0 +1,425 @@ +"""Prepare image and text latents for training. + +This script extracts image and text features using HunyuanVideo encoders +and saves them as .pt files for training. + +Usage: + # Single GPU + python dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/input.json \ + --output_dir /path/to/output \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 + + # Multi-GPU + torchrun --nproc_per_node=8 dataset/prepare_image_text_latent_simple.py \ + --input_json /path/to/input.json \ + --output_dir /path/to/output \ + --hunyuan_checkpoint_path /path/to/hunyuanvideo_1_5 + +Input JSON format: + [ + {"image_path": "/path/to/image1.jpg", "caption": "A sunset"}, + {"image_path": "/path/to/image2.png", "caption": "A cat"} + ] + +Output: + - {output_dir}/latents/{item_id}.pt + - {output_dir}/latents.json (for --json_path in training) +""" + +import os +import sys + +sys.path.append(os.path.abspath(".")) + +import argparse +import json +import re +from tqdm import tqdm +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from fastvideo.models.hyvideo.models.autoencoders import ( + hunyuanvideo_15_vae_w_cache, +) +from fastvideo.models.hyvideo.models.vision_encoder import VisionEncoder +from fastvideo.models.hyvideo.models.text_encoders import ( + PROMPT_TEMPLATE, + TextEncoder, +) +from fastvideo.models.hyvideo.models.text_encoders.byT5 import ( + load_glyph_byT5_v2, +) +from fastvideo.models.hyvideo.models.text_encoders.byT5.format_prompt import ( + MultilingualPromptFormat, +) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_json", + type=str, + required=True, + help="Input JSON file with image_path and caption fields", + ) + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Output directory for latent files", + ) + parser.add_argument( + "--hunyuan_checkpoint_path", + type=str, + required=True, + help="Path to HunyuanVideo checkpoint (contains vae/, text_encoder/, vision_encoder/)", + ) + parser.add_argument("--target_height", type=int, default=480) + parser.add_argument("--target_width", type=int, default=832) + parser.add_argument( + "--skip_existing", + action="store_true", + help="Skip if output file already exists", + ) + return parser.parse_args() + + +class LatentExtractor: + """Extract latents from images and text using HunyuanVideo encoders. + + This class loads and manages all encoder models: + - VAE: Encodes images to latent space + - Vision Encoder (SigLIP): Extracts visual features for conditioning + - Text Encoder (LLM): Encodes text prompts + - byT5: Encodes glyph/text features + """ + + def __init__(self, checkpoint_path, device, target_size=(480, 832)): + self.device = device + self.target_size = target_size # (height, width) + self._load_models(checkpoint_path) + + def _load_models(self, checkpoint_path): + """Load all encoder models from checkpoint directory.""" + print(f"Loading models from {checkpoint_path}...") + + # ===== VAE: Variational Autoencoder for image-to-latent encoding ===== + self.vae = ( + hunyuanvideo_15_vae_w_cache.AutoencoderKLConv3D.from_pretrained( + os.path.join(checkpoint_path, "vae"), torch_dtype=torch.float32 + ) + .to(self.device) + .eval() + ) + + # ===== Vision Encoder: SigLIP for extracting visual features ===== + self.vision_encoder = VisionEncoder( + vision_encoder_type="siglip", + vision_encoder_precision="fp16", + vision_encoder_path=os.path.join( + checkpoint_path, "vision_encoder/siglip" + ), + processor_type=None, + processor_path=None, + output_key=None, + logger=None, + device=self.device, + ) + + # ===== Text Encoder: LLM-based text encoder for prompt embeddings ===== + self.text_encoder = TextEncoder( + text_encoder_type="llm", + tokenizer_type="llm", + text_encoder_path=os.path.join(checkpoint_path, "text_encoder/llm"), + max_length=1000, + text_encoder_precision="fp16", + prompt_template=PROMPT_TEMPLATE["li-dit-encode-image-json"], + prompt_template_video=PROMPT_TEMPLATE["li-dit-encode-video-json"], + hidden_state_skip_layer=2, + apply_final_norm=False, + reproduce=False, + logger=None, + device=self.device, + ) + self.text_len = self.text_encoder.max_length + + # ===== byT5: Byte-level T5 for glyph/text encoding ===== + load_from = os.path.join(checkpoint_path, "text_encoder") + glyph_root = os.path.join(load_from, "Glyph-SDXL-v2") + byt5_args = dict( + byT5_google_path=os.path.join(load_from, "byt5-small"), + byT5_ckpt_path=os.path.join( + glyph_root, "checkpoints/byt5_model.pt" + ), + multilingual_prompt_format_color_path=os.path.join( + glyph_root, "assets/color_idx.json" + ), + multilingual_prompt_format_font_path=os.path.join( + glyph_root, "assets/multilingual_10-lang_idx.json" + ), + byt5_max_length=256, + ) + byt5_kwargs = load_glyph_byT5_v2( + byt5_args, + device=( + f"cuda:{self.device}" + if isinstance(self.device, int) + else str(self.device) + ), + ) + self.prompt_format = MultilingualPromptFormat( + font_path=byt5_args["multilingual_prompt_format_font_path"], + color_path=byt5_args["multilingual_prompt_format_color_path"], + ) + self.byt5_model = byt5_kwargs["byt5_model"] + self.byt5_tokenizer = byt5_kwargs["byt5_tokenizer"] + self.byt5_max_length = byt5_kwargs["byt5_max_length"] + print("All models loaded.") + + def _process_byt5_prompt(self, prompt_text): + """Process text prompt through byT5 encoder. Extracts quoted text (glyph text) from prompt + and encodes it. + + Args: + prompt_text: Input text prompt + + Returns: + byt5_embeddings: [1, max_length, 1472] tensor + byt5_mask: [1, max_length] attention mask + """ + byt5_embeddings = torch.zeros( + (1, self.byt5_max_length, 1472), device=self.device + ) + byt5_mask = torch.zeros( + (1, self.byt5_max_length), device=self.device, dtype=torch.int64 + ) + + # Extract quoted text (glyph text) from prompt using regex + pattern = r'\"(.*?)\"|"(.*?)"' + matches = re.findall(pattern, prompt_text) + glyph_texts = [m[0] or m[1] for m in matches] + glyph_texts = ( + list(dict.fromkeys(glyph_texts)) + if len(glyph_texts) > 1 + else glyph_texts + ) + + if glyph_texts: + # Format and encode glyph text + text_styles = [ + {"color": None, "font-family": None} for _ in glyph_texts + ] + formatted_text = self.prompt_format.format_prompt( + glyph_texts, text_styles + ) + inputs = self.byt5_tokenizer( + formatted_text, + padding="max_length", + max_length=self.byt5_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_ids = inputs.input_ids.to(self.device) + text_mask = inputs.attention_mask.to(self.device) + byt5_embeddings = self.byt5_model( + text_ids, attention_mask=text_mask.float() + )[0] + byt5_mask = text_mask + + return byt5_embeddings, byt5_mask + + @torch.no_grad() + def extract(self, image_path, caption): + """Extract all latents from an image and its caption. + + Args: + image_path: Path to input image + caption: Text caption/prompt for the image + + Returns: + dict containing: + - latent: VAE-encoded image latent [1, C, T, H, W] + - image_cond: Condition latent (first frame) [1, C, 1, H, W] + - prompt_embeds: Text embeddings from LLM encoder + - prompt_mask: Attention mask for text + - vision_states: Visual features from SigLIP + - byt5_text_states: byT5 embeddings + - byt5_text_mask: byT5 attention mask + """ + # ===== Step 1: Load and preprocess image ===== + # Load image and convert to tensor [C, H, W], normalize to [-1, 1] + image = Image.open(image_path).convert("RGB") + image_tensor = ( + torch.from_numpy(np.array(image)).permute(2, 0, 1).float() + ) + image_tensor = ( + image_tensor / 255.0 - 0.5 + ) * 2.0 # Normalize to [-1, 1] + + # Resize to target resolution + image_tensor = F.interpolate( + image_tensor.unsqueeze(0), + size=self.target_size, + mode="bilinear", + align_corners=False, + ).squeeze(0) + image_tensor = image_tensor.to(self.device) + + # ===== Step 2: VAE encoding ===== + # Encode image as condition latent (single frame) + # Input shape: [B, C, T, H, W] = [1, 3, 1, H, W] + cond_latents = self.vae.encode( + image_tensor.unsqueeze(0).unsqueeze(2) + ).latent_dist.mode() + cond_latents.mul_(self.vae.config.scaling_factor) + + # Encode as video latent (for training) + # Input shape: [B, C, T, H, W] = [1, 3, 1, H, W] + latents = self.vae.encode( + image_tensor.unsqueeze(1).unsqueeze(0) + ).latent_dist.mode() + latents.mul_(self.vae.config.scaling_factor) + + # ===== Step 3: Vision encoder (SigLIP) ===== + # Convert back to [0, 255] uint8 for vision encoder + ref_image = (image_tensor + 1) * 127.5 + vision_states = self.vision_encoder.encode_images( + ref_image.cpu().numpy().astype(np.uint8) + ) + vision_states = vision_states.last_hidden_state.to( + device=self.device, dtype=torch.bfloat16 + ) + + # ===== Step 4: Text encoder (LLM) ===== + text_inputs = self.text_encoder.text2tokens( + caption, data_type="video", max_length=self.text_len + ) + prompt_outputs = self.text_encoder.encode( + text_inputs, data_type="video", device=self.device + ) + prompt_embeds = prompt_outputs.hidden_state.to( + dtype=self.text_encoder.dtype, device=self.device + ) + attention_mask = ( + prompt_outputs.attention_mask.to(self.device) + if prompt_outputs.attention_mask is not None + else None + ) + + # ===== Step 5: byT5 encoder ===== + byt5_embeddings, byt5_masks = self._process_byt5_prompt(caption) + + return { + "latent": latents.to(torch.bfloat16), + "prompt_embeds": prompt_embeds, + "image_cond": cond_latents.to(torch.bfloat16), + "vision_states": vision_states, + "prompt_mask": attention_mask, + "byt5_text_states": byt5_embeddings, + "byt5_text_mask": byt5_masks, + } + + +def main(): + args = parse_args() + + # ===== Distributed setup ===== + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if world_size > 1: + torch.cuda.set_device(local_rank) + device = torch.cuda.current_device() + + # Create output directories + latents_dir = os.path.join(args.output_dir, "latents") + os.makedirs(latents_dir, exist_ok=True) + + # Load input data + with open(args.input_json, "r") as f: + data_list = json.load(f) + + # ===== Distribute data across GPUs ===== + if world_size > 1: + per_gpu = (len(data_list) + world_size - 1) // world_size + data_list = data_list[local_rank * per_gpu : (local_rank + 1) * per_gpu] + + print(f"GPU {local_rank}: Processing {len(data_list)} items") + + # Initialize extractor + extractor = LatentExtractor( + args.hunyuan_checkpoint_path, + device, + (args.target_height, args.target_width), + ) + + # ===== Process each item ===== + output_items = [] + for idx, item in enumerate(tqdm(data_list, desc=f"GPU {local_rank}")): + image_path = item.get("image_path") + caption = item.get("caption", "") + + if not image_path or not os.path.exists(image_path): + continue + + item_id = f"{local_rank}_{idx:06d}" + output_path = os.path.join(latents_dir, f"{item_id}.pt") + + # Skip if already processed + if args.skip_existing and os.path.exists(output_path): + output_items.append( + { + "latent_path": output_path, + "image_path": image_path, + "caption": caption, + } + ) + continue + + try: + latents = extractor.extract(image_path, caption) + torch.save(latents, output_path) + output_items.append( + { + "latent_path": output_path, + "image_path": image_path, + "caption": caption, + } + ) + except Exception as e: + print(f"Error processing {image_path}: {e}") + + # ===== Save output JSON ===== + output_json = os.path.join(args.output_dir, f"latents_gpu{local_rank}.json") + with open(output_json, "w") as f: + json.dump(output_items, f, indent=2, ensure_ascii=False) + + # ===== Merge results on rank 0 ===== + if local_rank == 0 and world_size > 1: + import time + + time.sleep(3) # Wait for other GPUs to finish writing + all_items = [] + for gpu_id in range(world_size): + gpu_json = os.path.join( + args.output_dir, f"latents_gpu{gpu_id}.json" + ) + if os.path.exists(gpu_json): + with open(gpu_json, "r") as f: + all_items.extend(json.load(f)) + merged_json = os.path.join(args.output_dir, "latents.json") + with open(merged_json, "w") as f: + json.dump(all_items, f, indent=2, ensure_ascii=False) + print(f"Merged {len(all_items)} items to {merged_json}") + elif world_size == 1: + os.rename(output_json, os.path.join(args.output_dir, "latents.json")) + print( + f"Saved {len(output_items)} items to {args.output_dir}/latents.json" + ) + + +if __name__ == "__main__": + main() diff --git a/worldcompass/requirements.txt b/worldcompass/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..aad0ba5103c2dad934147760404d3a76546a7a09 --- /dev/null +++ b/worldcompass/requirements.txt @@ -0,0 +1,43 @@ +hpsv3==1.0.0 +torch==2.7.1 +torchvision==0.22.1 +torchdata==0.11.0 +triton==3.3.1 +tensorboard==2.20.0 +diffusers==0.35.1 +huggingface_hub==0.35.0 +safetensors==0.4.5 +peft==0.17.1 +accelerate==1.10.1 +numpy==1.26.4 +pandas==2.3.2 +scipy==1.15.1 +pyarrow==21.0.0 +Pillow==12.0.0 +pillow_heif==1.1.1 +imageio==2.37.0 +moviepy==1.0.3 +opencv_python==4.11.0.86 +trimesh==4.6.2 +plyfile==1.1 +e3nn==0.5.8 +omegaconf==2.3.0 +PyYAML==6.0.2 +pydantic==2.11.9 +addict==2.4.0 +matplotlib==3.9.0 +ray==2.43.0 +deepspeed==0.17.6 +loguru==0.7.3 +tqdm==4.66.5 +wandb==0.22.0 +uvicorn==0.38.0 +Requests==2.32.5 +fastapi==0.121.2 +remote_pdb==2.1.0 +psutil==6.0.0 +evo==1.34.0 +einops==0.8.1 +cloudpickle==3.1.1 +gsplat==1.5.3 +pycolmap==3.13.0 \ No newline at end of file diff --git a/worldcompass/reward_function/HunyuanWorldMirror/License.txt b/worldcompass/reward_function/HunyuanWorldMirror/License.txt new file mode 100644 index 0000000000000000000000000000000000000000..13c89039217923c9c21b29a66a5010c474d775bb --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/License.txt @@ -0,0 +1,82 @@ +TENCENT HUNYUANWORLD-MIRROR COMMUNITY LICENSE AGREEMENT +Tencent HunyuanWorld-Mirror Release Date: October 22, 2025 +THIS LICENSE AGREEMENT DOES NOT APPLY IN THE EUROPEAN UNION, UNITED KINGDOM AND SOUTH KOREA AND IS EXPRESSLY LIMITED TO THE TERRITORY, AS DEFINED BELOW. +By clicking to agree or by using, reproducing, modifying, distributing, performing or displaying any portion or element of the Tencent HunyuanWorld-Mirror Works, including via any Hosted Service, You will be deemed to have recognized and accepted the content of this Agreement, which is effective immediately. +1. DEFINITIONS. +a. “Acceptable Use Policy” shall mean the policy made available by Tencent as set forth in the Exhibit A. +b. “Agreement” shall mean the terms and conditions for use, reproduction, distribution, modification, performance and displaying of Tencent HunyuanWorld-Mirror Works or any portion or element thereof set forth herein. +c. “Documentation” shall mean the specifications, manuals and documentation for Tencent HunyuanWorld-Mirror made publicly available by Tencent. +d. “Hosted Service” shall mean a hosted service offered via an application programming interface (API), web access, or any other electronic or remote means. +e. “Licensee,” “You” or “Your” shall mean a natural person or legal entity exercising the rights granted by this Agreement and/or using the Tencent HunyuanWorld-Mirror Works for any purpose and in any field of use. +f. “Materials” shall mean, collectively, Tencent’s proprietary Tencent HunyuanWorld-Mirror and Documentation (and any portion thereof) as made available by Tencent under this Agreement. +g. “Model Derivatives” shall mean all: (i) modifications to Tencent HunyuanWorld-Mirror or any Model Derivative of Tencent HunyuanWorld-Mirror; (ii) works based on Tencent HunyuanWorld-Mirror or any Model Derivative of Tencent HunyuanWorld-Mirror; or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Tencent HunyuanWorld-Mirror or any Model Derivative of Tencent HunyuanWorld-Mirror, to that model in order to cause that model to perform similarly to Tencent HunyuanWorld-Mirror or a Model Derivative of Tencent HunyuanWorld-Mirror, including distillation methods, methods that use intermediate data representations, or methods based on the generation of synthetic data Outputs by Tencent HunyuanWorld-Mirror or a Model Derivative of Tencent HunyuanWorld-Mirror for training that model. For clarity, Outputs by themselves are not deemed Model Derivatives. +h. “Output” shall mean the information and/or content output of Tencent HunyuanWorld-Mirror or a Model Derivative that results from operating or otherwise using Tencent HunyuanWorld-Mirror or a Model Derivative, including via a Hosted Service. +i. “Tencent,” “We” or “Us” shall mean the applicable entity or entities in the Tencent corporate family that own(s) intellectual property or other rights embodied in or utilized by the Materials.. +j. “Tencent HunyuanWorld-Mirror” shall mean the 3D generation models and their software and algorithms, including trained model weights, parameters (including optimizer states), machine-learning model code, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements of the foregoing made publicly available by Us at [https://github.com/Tencent-Hunyuan/HunyuanWorld-Mirror]. +k. “Tencent HunyuanWorld-Mirror Works” shall mean: (i) the Materials; (ii) Model Derivatives; and (iii) all derivative works thereof. +l. “Territory” shall mean the worldwide territory, excluding the territory of the European Union, United Kingdom and South Korea. +m. “Third Party” or “Third Parties” shall mean individuals or legal entities that are not under common control with Us or You. +n. “including” shall mean including but not limited to. +2. GRANT OF RIGHTS. +We grant You, for the Territory only, a non-exclusive, non-transferable and royalty-free limited license under Tencent’s intellectual property or other rights owned by Us embodied in or utilized by the Materials to use, reproduce, distribute, create derivative works of (including Model Derivatives), and make modifications to the Materials, only in accordance with the terms of this Agreement and the Acceptable Use Policy, and You must not violate (or encourage or permit anyone else to violate) any term of this Agreement or the Acceptable Use Policy. +3. DISTRIBUTION. +You may, subject to Your compliance with this Agreement, distribute or make available to Third Parties the Tencent HunyuanWorld-Mirror Works, exclusively in the Territory, provided that You meet all of the following conditions: +a. You must provide all such Third Party recipients of the Tencent HunyuanWorld-Mirror Works or products or services using them a copy of this Agreement; +b. You must cause any modified files to carry prominent notices stating that You changed the files; +c. You are encouraged to: (i) publish at least one technology introduction blogpost or one public statement expressing Your experience of using the Tencent HunyuanWorld-Mirror Works; and (ii) mark the products or services developed by using the Tencent HunyuanWorld-Mirror Works to indicate that the product/service is “Powered by Tencent Hunyuan”; +d. All distributions to Third Parties (other than through a Hosted Service) must be accompanied by a “Notice” text file that contains the following notice: “Tencent HunyuanWorld-Mirror is licensed under the Tencent HunyuanWorld-Mirror Community License Agreement, Copyright © 2025 Tencent. All Rights Reserved. The trademark rights of “Tencent Hunyuan” are owned by Tencent or its affiliate”; +e. In the event that You use, integrate, implement, or otherwise deploy the Tencent Hunyuan Works, in whole or in part, to provide, enable, or support any service, product, or functionality to third parties, You shall clearly, accurately, and prominently disclose to all end users the full legal name and entity of the actual provider of such service, product, or functionality. You shall expressly and conspicuously state that Tencent is not affiliated with, associated with, sponsoring, or endorsing any such service, product, or functionality. You shall not use or display any name, logo, trademark, trade name, or other indicia of Tencent in any manner that could be construed as, or be likely to create, confusion, deception, or a false impression regarding any relationship, affiliation, sponsorship, or endorsement by Tencent. +You may add Your own copyright statement to Your modifications and, except as set forth in this Section and in Section 5, may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Model Derivatives as a whole, provided Your use, reproduction, modification, distribution, performance and display of the work otherwise complies with the terms and conditions of this Agreement (including as regards the Territory). If You receive Tencent HunyuanWorld-Mirror Works from a Licensee as part of an integrated end user product, then this Section 3 of this Agreement will not apply to You. +4. ADDITIONAL COMMERCIAL TERMS. +If, on the Tencent HunyuanWorld-Mirror version release date, the monthly active users of all products or services made available by or for Licensee is greater than 1 million monthly active users in the preceding calendar month, You must request a license from Tencent, which Tencent may grant to You in its sole discretion, and You are not authorized to exercise any of the rights under this Agreement unless or until Tencent otherwise expressly grants You such rights. +Subject to Tencent's written approval, you may request a license for the use of Tencent HunyuanWorld-Mirror by submitting the following information to hunyuan3d@tencent.com: +a. Your company’s name and associated business sector that plans to use Tencent HunyuanWorld-Mirror. +b. Your intended use case and the purpose of using Tencent HunyuanWorld-Mirror. +c. Your plans to modify Tencent HunyuanWorld-Mirror or create Model Derivatives. +5. RULES OF USE. +a. Your use of the Tencent HunyuanWorld-Mirror Works must comply with applicable laws and regulations (including trade compliance laws and regulations) and adhere to the Acceptable Use Policy for the Tencent HunyuanWorld-Mirror Works, which is hereby incorporated by reference into this Agreement. You must include the use restrictions referenced in these Sections 5(a) and 5(b) as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Tencent HunyuanWorld-Mirror Works and You must provide notice to subsequent users to whom You distribute that Tencent HunyuanWorld-Mirror Works are subject to the use restrictions in these Sections 5(a) and 5(b). +b. You must not use the Tencent HunyuanWorld-Mirror Works or any Output or results of the Tencent HunyuanWorld-Mirror Works to improve any other AI model (other than Tencent HunyuanWorld-Mirror or Model Derivatives thereof). +c. You must not use, reproduce, modify, distribute, or display the Tencent HunyuanWorld-Mirror Works, Output or results of the Tencent HunyuanWorld-Mirror Works outside the Territory. Any such use outside the Territory is unlicensed and unauthorized under this Agreement. +6. INTELLECTUAL PROPERTY. +a. Subject to Tencent’s ownership of Tencent HunyuanWorld-Mirror Works made by or for Tencent and intellectual property rights therein, conditioned upon Your compliance with the terms and conditions of this Agreement, as between You and Tencent, You will be the owner of any derivative works and modifications of the Materials and any Model Derivatives that are made by or for You. +b. No trademark licenses are granted under this Agreement, and in connection with the Tencent HunyuanWorld-Mirror Works, Licensee may not use any name or mark owned by or associated with Tencent or any of its affiliates, except as required for reasonable and customary use in describing and distributing the Tencent HunyuanWorld-Mirror Works. Tencent hereby grants You a license to use “Tencent Hunyuan” (the “Mark”) in the Territory solely as required to comply with the provisions of Section 3(c), provided that You comply with any applicable laws related to trademark protection. All goodwill arising out of Your use of the Mark will inure to the benefit of Tencent. +c. If You commence a lawsuit or other proceedings (including a cross-claim or counterclaim in a lawsuit) against Us or any person or entity alleging that the Materials or any Output, or any portion of any of the foregoing, infringe any intellectual property or other right owned or licensable by You, then all licenses granted to You under this Agreement shall terminate as of the date such lawsuit or other proceeding is filed. You will defend, indemnify and hold harmless Us from and against any claim by any Third Party arising out of or related to Your or the Third Party’s use or distribution of the Tencent HunyuanWorld-Mirror Works. +d. Tencent claims no rights in Outputs You generate. You and Your users are solely responsible for Outputs and their subsequent uses. +7. DISCLAIMERS OF WARRANTY AND LIMITATIONS OF LIABILITY. +a. We are not obligated to support, update, provide training for, or develop any further version of the Tencent HunyuanWorld-Mirror Works or to grant any license thereto. +b. UNLESS AND ONLY TO THE EXTENT REQUIRED BY APPLICABLE LAW, THE TENCENT HUNYUANWORLD-MIRROR WORKS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED “AS IS” WITHOUT ANY EXPRESS OR IMPLIED WARRANTIES OF ANY KIND INCLUDING ANY WARRANTIES OF TITLE, MERCHANTABILITY, NONINFRINGEMENT, COURSE OF DEALING, USAGE OF TRADE, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR DISTRIBUTING ANY OF THE TENCENT HUNYUANWORLD-MIRROR WORKS OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR OR A THIRD PARTY’S USE OR DISTRIBUTION OF ANY OF THE TENCENT HUNYUANWORLD-MIRROR WORKS OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT. +c. TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT SHALL TENCENT OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, FOR ANY DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO ANY OF THE TENCENT HUNYUANWORLD-MIRROR WORKS OR OUTPUTS, EVEN IF TENCENT OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING. +8. SURVIVAL AND TERMINATION. +a. The term of this Agreement shall commence upon Your acceptance of this Agreement or access to the Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. +b. We may terminate this Agreement if You breach any of the terms or conditions of this Agreement. Upon termination of this Agreement, You must promptly delete and cease use of the Tencent HunyuanWorld-Mirror Works. Sections 6(a), 6(c), 7 and 9 shall survive the termination of this Agreement. +9. GOVERNING LAW AND JURISDICTION. +a. This Agreement and any dispute arising out of or relating to it will be governed by the laws of the Hong Kong Special Administrative Region of the People’s Republic of China, without regard to conflict of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. +b. Exclusive jurisdiction and venue for any dispute arising out of or relating to this Agreement will be a court of competent jurisdiction in the Hong Kong Special Administrative Region of the People’s Republic of China, and Tencent and Licensee consent to the exclusive jurisdiction of such court with respect to any such dispute. +  +EXHIBIT A +ACCEPTABLE USE POLICY + +Tencent reserves the right to update this Acceptable Use Policy from time to time. +Last modified: November 5, 2024 + +Tencent endeavors to promote safe and fair use of its tools and features, including Tencent HunyuanWorld-Mirror. You agree not to use Tencent HunyuanWorld-Mirror or Model Derivatives: +1. Outside the Territory; +2. In any way that violates any applicable national, federal, state, local, international or any other law or regulation; +3. To harm Yourself or others; +4. To repurpose or distribute output from Tencent HunyuanWorld-Mirror or any Model Derivatives to harm Yourself or others; +5. To override or circumvent the safety guardrails and safeguards We have put in place; +6. For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; +7. To generate or disseminate verifiably false information and/or content with the purpose of harming others or influencing elections; +8. To generate or facilitate false online engagement, including fake reviews and other means of fake online engagement; +9. To intentionally defame, disparage or otherwise harass others; +10. To generate and/or disseminate malware (including ransomware) or any other content to be used for the purpose of harming electronic systems; +11. To generate or disseminate personal identifiable information with the purpose of harming others; +12. To generate or disseminate information (including images, code, posts, articles), and place the information in any public context (including –through the use of bot generated tweets), without expressly and conspicuously identifying that the information and/or content is machine generated; +13. To impersonate another individual without consent, authorization, or legal right; +14. To make high-stakes automated decisions in domains that affect an individual’s safety, rights or wellbeing (e.g., law enforcement, migration, medicine/health, management of critical infrastructure, safety components of products, essential services, credit, employment, housing, education, social scoring, or insurance); +15. In a manner that violates or disrespects the social ethics and moral standards of other countries or regions; +16. To perform, facilitate, threaten, incite, plan, promote or encourage violent extremism or terrorism; +17. For any use intended to discriminate against or harm individuals or groups based on protected characteristics or categories, online or offline social behavior or known or predicted personal or personality characteristics; +18. To intentionally exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm; +19. For military purposes; +20. To engage in the unauthorized or unlicensed practice of any profession including, but not limited to, financial, legal, medical/health, or other professional practices. \ No newline at end of file diff --git a/worldcompass/reward_function/HunyuanWorldMirror/Notice.txt b/worldcompass/reward_function/HunyuanWorldMirror/Notice.txt new file mode 100644 index 0000000000000000000000000000000000000000..bbaebba9e386016bdbde3a8e7f4fcd4de2ad0e2a --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/Notice.txt @@ -0,0 +1,76 @@ +Usage and Legal Notices: + +Tencent is pleased to support the open source community by making Tencent HunyuanWorld-Mirror . + +Copyright (C) 2025 Tencent. All rights reserved. The below software and/or models in this distribution may have been modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) Tencent. + +Tencent HunyuanWorld-Mirror is licensed under the TENCENT HUNYUANWORLD-MIRROR COMMUNITY LICENSE AGREEMENT except for the third-party components listed below, which is licensed under different terms. TTencent HunyuanWorld-Mirror does not impose any additional limitations beyond what is outlined in the respective licenses of these third-party components. Users must comply with all terms and conditions of original licenses of these third-party components and must ensure that the usage of the third party components adheres to all relevant laws and regulations. + +For avoidance of doubts, Tencent HunyuanWorld-Mirror means inference-enabling code, parameters, and/or weights of this Model, which are made publicly available by Tencent in accordance with TENCENT HUNYUANWORLD-MIRROR COMMUNITY LICENSE AGREEMENT. + + +Other dependencies and licenses: + + + +Open Source Software Licensed under the Apache-2.0: +-------------------------------------------------------------------- +1. gsplat + +Copyright 2025 Nerfstudio Team. + +You can access this component through: https://github.com/nerfstudio-project/gsplat + +Terms of the Apache-2.0: +-------------------------------------------------------------------- +Apache License +Version 2.0, January 2004 +http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + +"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. + +"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. + +"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. + +"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. + +"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. + +"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. + +"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). + +"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. + +"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." + +"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: + +You must give any other recipients of the Work or Derivative Works a copy of this License; and +You must cause any modified files to carry prominent notices stating that You changed the files; and +You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and +If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. +You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS \ No newline at end of file diff --git a/worldcompass/reward_function/HunyuanWorldMirror/README.md b/worldcompass/reward_function/HunyuanWorldMirror/README.md new file mode 100644 index 0000000000000000000000000000000000000000..946145c35839ae5b5f5f90f8626d7e310f6ea85f --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/README.md @@ -0,0 +1,74 @@ +# HunyuanWorldMirror - Inference Only + +This is a minimal inference-only version of HunyuanWorldMirror, stripped down +for use in reward functions. + +## Structure + +``` +HunyuanWorldMirror/ +├── __init__.py # Main package init +├── infer.py # Standalone inference script +├── src/ +│ ├── models/ +│ │ ├── heads/ # Camera and dense prediction heads +│ │ ├── layers/ # Transformer layers +│ │ ├── models/ # Core WorldMirror model +│ │ └── utils/ # Model utilities +│ └── utils/ # Inference utilities +├── License.txt +└── Notice.txt +``` + +## Usage + +### As a module: + +```python +from HunyuanWorldMirror import WorldMirror + +# Load model +model = WorldMirror.from_pretrained("tencent/HunyuanWorldMirror") +model.eval() + +# Prepare images +# ... (see infer.py for full example) + +# Run inference +with torch.no_grad(): + predictions = model(images) +``` + +### Standalone inference: + +```bash +python infer.py --input_path --output_dir +``` + +## Changes from Original + +- Removed all training code (`training/` directory) +- Removed demo applications (`app.py`, `demo/`) +- Removed examples and documentation +- Removed submodules +- All imports updated to use `HunyuanWorldMirror.src.*` prefix for proper module + resolution +- Cleaned up cache files and assets + +## Dependencies + +Core dependencies for inference: + +- torch +- numpy +- PIL +- cv2 +- gsplat (for Gaussian splatting) +- einops +- huggingface_hub + +Optional dependencies: + +- onnxruntime (for sky segmentation) +- pycolmap (for reconstruction) +- moviepy (for video rendering) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/__init__.py b/worldcompass/reward_function/HunyuanWorldMirror/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7c2fce445bdef7bc96a16f553ec00b5185135aa7 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/__init__.py @@ -0,0 +1,9 @@ +""" +HunyuanWorldMirror - Inference-only module for WorldMirror model +""" + +from reward_function.HunyuanWorldMirror.src.models.models.worldmirror import ( + WorldMirror, +) + +__all__ = ["WorldMirror"] diff --git a/worldcompass/reward_function/HunyuanWorldMirror/infer.py b/worldcompass/reward_function/HunyuanWorldMirror/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..a1d9357925679b778a663f1398abac144469b4ec --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/infer.py @@ -0,0 +1,686 @@ +import argparse +import glob +from pathlib import Path +import os +import time + +import cv2 +import numpy as np +import torch +from PIL import Image +import onnxruntime + +from reward_function.HunyuanWorldMirror.src.models.models.worldmirror import ( + WorldMirror, +) +from reward_function.HunyuanWorldMirror.src.utils.inference_utils import ( + prepare_images_to_tensor, +) +from reward_function.HunyuanWorldMirror.src.utils.video_utils import ( + video_to_image_frames, +) +from reward_function.HunyuanWorldMirror.src.models.utils.geometry import ( + depth_to_world_coords_points, +) +from reward_function.HunyuanWorldMirror.src.models.utils.geometry import ( + create_pixel_coordinate_grid, +) + +from reward_function.HunyuanWorldMirror.src.utils.save_utils import ( + save_depth_png, + save_depth_npy, + save_normal_png, +) +from reward_function.HunyuanWorldMirror.src.utils.save_utils import ( + save_scene_ply, + save_gs_ply, + save_points_ply, +) +from reward_function.HunyuanWorldMirror.src.utils.render_utils import ( + render_interpolated_video, +) + +from reward_function.HunyuanWorldMirror.src.utils.build_pycolmap_recon import ( + build_pycolmap_reconstruction, +) +from reward_function.HunyuanWorldMirror.src.models.utils.camera_utils import ( + vector_to_camera_matrices, +) + +# Import mask computation utilities +from reward_function.HunyuanWorldMirror.src.utils.geometry import ( + depth_edge, + normals_edge, +) +from reward_function.HunyuanWorldMirror.src.utils.visual_util import ( + segment_sky, + download_file_from_url, +) + + +def create_filter_mask( + pts3d_conf: np.ndarray, + depth_preds: np.ndarray, + normal_preds: np.ndarray, + sky_mask: np.ndarray, + confidence_percentile: float = 10.0, + edge_normal_threshold: float = 5.0, + edge_depth_threshold: float = 0.03, + apply_confidence_mask: bool = True, + apply_edge_mask: bool = True, + apply_sky_mask: bool = False, +) -> np.ndarray: + """Create comprehensive filter mask based on confidence, edges, and sky segmentation. This + follows the same logic as app.py for consistent mask computation. + + Args: + pts3d_conf: Point confidence scores [S, H, W] + depth_preds: Depth predictions [S, H, W, 1] + normal_preds: Normal predictions [S, H, W, 3] + sky_mask: Sky segmentation mask [S, H, W] + confidence_percentile: Percentile threshold for confidence filtering (0-100) + edge_normal_threshold: Normal angle threshold in degrees for edge detection + edge_depth_threshold: Relative depth threshold for edge detection + apply_confidence_mask: Whether to apply confidence-based filtering + apply_edge_mask: Whether to apply edge-based filtering + apply_sky_mask: Whether to apply sky mask filtering + + Returns: + final_mask: Boolean mask array [S, H, W] for filtering points + """ + S, H, W = pts3d_conf.shape[:3] + final_mask_list = [] + + for i in range(S): + final_mask = None + + if apply_confidence_mask: + # Compute confidence mask based on the pointmap confidence + confidences = pts3d_conf[i, :, :] # [H, W] + percentile_threshold = np.quantile( + confidences, confidence_percentile / 100.0 + ) + conf_mask = confidences >= percentile_threshold + if final_mask is None: + final_mask = conf_mask + else: + final_mask = final_mask & conf_mask + + if apply_edge_mask: + # Compute edge mask based on the normalmap + normal_pred = normal_preds[i] # [H, W, 3] + normal_edges = normals_edge( + normal_pred, tol=edge_normal_threshold, mask=final_mask + ) + # Compute depth mask based on the depthmap + depth_pred = depth_preds[i, :, :, 0] # [H, W] + depth_edges = depth_edge( + depth_pred, rtol=edge_depth_threshold, mask=final_mask + ) + edge_mask = ~(depth_edges & normal_edges) + if final_mask is None: + final_mask = edge_mask + else: + final_mask = final_mask & edge_mask + + if apply_sky_mask: + # Apply sky mask filtering (sky_mask is already inverted: True = non-sky) + sky_mask_frame = sky_mask[i] # [H, W] + if final_mask is None: + final_mask = sky_mask_frame + else: + final_mask = final_mask & sky_mask_frame + + final_mask_list.append(final_mask) + + # Stack all frame masks + if final_mask_list[0] is not None: + final_mask = np.stack(final_mask_list, axis=0) # [S, H, W] + else: + final_mask = np.ones(pts3d_conf.shape[:3], dtype=bool) # [S, H, W] + + return final_mask + + +def main(): + parser = argparse.ArgumentParser( + description="HunyuanWorld-Mirror inference" + ) + parser.add_argument( + "--input_path", + type=str, + default="examples/realistic/Ireland_Landscape", + help="Input can be: a directory of images; a single video file; or a directory containing multiple video files (.mp4/.avi/.mov/.webm/.gif). If directory has multiple videos, frames from all clips are extracted (using --fps) and merged in filename order.", + ) + parser.add_argument("--output_path", type=str, default="inference_output") + parser.add_argument( + "--fps", + type=int, + default=1, + help="Frames per second for video extraction", + ) + parser.add_argument( + "--target_size", + type=int, + default=518, + help="Target size for image resizing", + ) + parser.add_argument( + "--write_txt", + action="store_true", + help="Also write human-readable COLMAP txt (slow, huge)", + ) + # Mask filtering parameters + parser.add_argument( + "--confidence_percentile", + type=float, + default=10.0, + help="Confidence percentile threshold for filtering (0-100, filters bottom X percent)", + ) + parser.add_argument( + "--edge_normal_threshold", + type=float, + default=5.0, + help="Normal angle threshold in degrees for edge detection", + ) + parser.add_argument( + "--edge_depth_threshold", + type=float, + default=0.03, + help="Relative depth threshold for edge detection", + ) + parser.add_argument( + "--apply_confidence_mask", + action="store_true", + default=True, + help="Apply confidence-based filtering", + ) + parser.add_argument( + "--apply_edge_mask", + action="store_true", + default=True, + help="Apply edge-based filtering", + ) + parser.add_argument( + "--apply_sky_mask", + action="store_true", + default=False, + help="Apply sky mask filtering", + ) + # Save flags + parser.add_argument( + "--save_pointmap", + action="store_true", + default=True, + help="Save points PLY", + ) + parser.add_argument( + "--save_depth", action="store_true", default=True, help="Save depth PNG" + ) + parser.add_argument( + "--save_normal", + action="store_true", + default=True, + help="Save normal PNG", + ) + parser.add_argument( + "--save_gs", + action="store_true", + default=True, + help="Save Gaussians PLY", + ) + parser.add_argument( + "--save_rendered", + action="store_true", + default=True, + help="Save rendered video", + ) + parser.add_argument( + "--save_colmap", + action="store_true", + default=True, + help="Save COLMAP sparse", + ) + # Conditioning flags + parser.add_argument( + "--cond_pose", + action="store_true", + help="Use camera pose conditioning if available", + ) + parser.add_argument( + "--cond_intrinsics", + action="store_true", + help="Use intrinsics conditioning if available", + ) + parser.add_argument( + "--cond_depth", + action="store_true", + help="Use depth conditioning if available", + ) + args = parser.parse_args() + + # Print inference parameters + print(f"🔧 Configuration:") + print(f" - FPS: {args.fps}") + print(f" - Target size: {args.target_size}px") + print(f" - Mask Filtering:") + print( + f" - Confidence mask: {'✅' if args.apply_confidence_mask else '❌'} (percentile: {args.confidence_percentile}%)" + ) + print( + f" - Edge mask: {'✅' if args.apply_edge_mask else '❌'} (normal: {args.edge_normal_threshold}°, depth: {args.edge_depth_threshold})" + ) + print(f" - Sky mask: {'✅' if args.apply_sky_mask else '❌'}") + print(f" - Conditioning:") + print(f" - Pose: {'✅' if args.cond_pose else '❌'}") + print(f" - Intrinsics: {'✅' if args.cond_intrinsics else '❌'}") + print(f" - Depth: {'✅' if args.cond_depth else '❌'}") + + # 1) Init model - This requires internet access or the huggingface hub cache to be pre-downloaded + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = WorldMirror.from_pretrained("tencent/HunyuanWorld-Mirror").to( + device + ) + model.eval() + + input_path = Path(args.input_path) + + # Create output directory with filename-based subdirectory + if input_path.is_file(): + subdir_name = input_path.stem + elif input_path.is_dir(): + subdir_name = input_path.name + else: + raise ValueError( + f"❌ Invalid input path: {input_path} (must be directory or video file)" + ) + + outdir = Path(args.output_path) / subdir_name + outdir.mkdir(parents=True, exist_ok=True) + + # Determine input type and get image paths + video_exts = [".mp4", ".avi", ".mov", ".webm", ".gif"] + + if input_path.is_file() and input_path.suffix.lower() in video_exts: + # Case 1: Single video file - extract frames + print(f"📹 Processing video: {input_path}") + input_frames_dir = outdir / "input_frames" + input_frames_dir.mkdir(exist_ok=True) + + img_paths = video_to_image_frames( + str(input_path), str(input_frames_dir), fps=args.fps + ) + if not img_paths: + raise RuntimeError("❌ Failed to extract frames from video") + + img_paths = sorted(img_paths) + print(f"✅ Extracted {len(img_paths)} frames to {input_frames_dir}") + + elif input_path.is_dir(): + # Case 2: Directory of images + print(f"📁 Processing directory: {input_path}") + img_paths = [] + for ext in ["*.jpeg", "*.jpg", "*.png", "*.webp"]: + img_paths.extend(glob.glob(os.path.join(str(input_path), ext))) + if len(img_paths) == 0: + raise FileNotFoundError(f"❌ No image files found in {input_path}") + img_paths = sorted(img_paths) + print(f"✅ Loaded {len(img_paths)} images from {input_path}") + + else: + raise ValueError(f"❌ Invalid input path: {input_path}") + + # 3) Load and preprocess images + views = {} + imgs = prepare_images_to_tensor( + img_paths, target_size=args.target_size, resize_strategy="crop" + ).to( + device + ) # [1,S,3,H,W], in [0,1] + views["img"] = imgs + B, S, C, H, W = imgs.shape + cond_flags = [0, 0, 0] + print(f"📸 Loaded {S} images with shape {imgs.shape}") + + # 4) Inference + print("\n🚀 Starting inference pipeline...") + start_time = time.time() + use_amp = torch.cuda.is_available() and torch.cuda.is_bf16_supported() + if use_amp: + amp_dtype = torch.bfloat16 + else: + amp_dtype = torch.float32 + with torch.no_grad(): + with torch.amp.autocast("cuda", enabled=bool(use_amp), dtype=amp_dtype): + predictions = model( + views=views, cond_flags=cond_flags + ) # Multi-modal inference with priors + print(f"🕒 Inference time: {time.time() - start_time:.3f} seconds") + + # 4.5) Sky mask segmentation (if needed) + sky_mask = None + if args.apply_sky_mask: + print("\n🌤️ Computing sky masks...") + if not os.path.exists("skyseg.onnx"): + print("Downloading skyseg.onnx...") + download_file_from_url( + "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", + "skyseg.onnx", + ) + skyseg_session = onnxruntime.InferenceSession("skyseg.onnx") + sky_mask_list = [] + for i, img_path in enumerate(img_paths): + sky_mask_frame = segment_sky(img_path, skyseg_session) + # Resize mask to match H×W if needed + if sky_mask_frame.shape[0] != H or sky_mask_frame.shape[1] != W: + sky_mask_frame = cv2.resize(sky_mask_frame, (W, H)) + sky_mask_list.append(sky_mask_frame) + sky_mask = np.stack(sky_mask_list, axis=0) # [S, H, W] + sky_mask = sky_mask > 0 # Binary mask: True = non-sky, False = sky + print(f"✅ Sky masks computed for {S} frames") + else: + # Create dummy sky mask (all True = keep all points) + sky_mask = np.ones((S, H, W), dtype=bool) + + # 5) Save results + print("\n📤 Saving results...") + images_dir = outdir / "images" # original resolution images + images_dir.mkdir(exist_ok=True) + images_resized_dir = outdir / "images_resized" # resized images + images_resized_dir.mkdir(exist_ok=True) + if args.save_depth: + depth_dir = outdir / "depth" + depth_dir.mkdir(exist_ok=True) + if args.save_normal: + normal_dir = outdir / "normal" + normal_dir.mkdir(exist_ok=True) + if args.save_colmap: + sparse_dir = outdir / "sparse" / "0" + sparse_dir.mkdir(parents=True, exist_ok=True) + + # save images + processed_image_names = [] + for i in range(S): + im = ( + (imgs[0, i].permute(1, 2, 0).clamp(0, 1) * 255) + .to(torch.uint8) + .cpu() + .numpy() + ) + fname = f"image_{i + 1:04d}.png" + Image.fromarray(im).save(str(images_resized_dir / fname)) + pil_img = Image.open(img_paths[i]).convert("RGB") + processed_height, processed_width = ( + imgs[0, i].shape[1], + imgs[0, i].shape[2], + ) + processed_aspect_ratio = processed_width / processed_height + orig_width, orig_height = pil_img.size + new_height = int(orig_width / processed_aspect_ratio) + new_width = orig_width + pil_img = pil_img.resize( + (orig_width, new_height), Image.Resampling.BICUBIC + ) + pil_img.save(str(images_dir / fname)) + + processed_image_names.append(fname) + + # save pointmap with filtering + if "pts3d" in predictions and args.save_pointmap: + print("Computing filter mask for pointmap...") + + # Prepare data for mask computation + pts3d_conf_np = ( + predictions["pts3d_conf"][0].detach().cpu().numpy() + ) # [S, H, W] + depth_preds_np = ( + predictions["depth"][0].detach().cpu().numpy() + ) # [S, H, W, 1] + normal_preds_np = ( + predictions["normals"][0].detach().cpu().numpy() + ) # [S, H, W, 3] + + # Compute comprehensive filter mask + final_mask = create_filter_mask( + pts3d_conf=pts3d_conf_np, + depth_preds=depth_preds_np, + normal_preds=normal_preds_np, + sky_mask=sky_mask, + confidence_percentile=args.confidence_percentile, + edge_normal_threshold=args.edge_normal_threshold, + edge_depth_threshold=args.edge_depth_threshold, + apply_confidence_mask=args.apply_confidence_mask, + apply_edge_mask=args.apply_edge_mask, + apply_sky_mask=args.apply_sky_mask, + ) # [S, H, W] + + # Collect points and colors + pts_list = [] + pts_colors_list = [] + + for i in range(S): + pts = predictions["pts3d"][0, i] # [H,W,3] + img_colors = imgs[0, i].permute(1, 2, 0) # [H, W, 3] + img_colors = (img_colors * 255).to(torch.uint8) + + pts_list.append(pts.reshape(-1, 3)) + pts_colors_list.append(img_colors.reshape(-1, 3)) + + all_pts = torch.cat(pts_list, dim=0) + all_colors = torch.cat(pts_colors_list, dim=0) + + # Apply filter mask + final_mask_flat = final_mask.reshape(-1) # Flatten to [S*H*W] + final_mask_torch = torch.from_numpy(final_mask_flat).to(all_pts.device) + + filtered_pts = all_pts[final_mask_torch] + filtered_colors = all_colors[final_mask_torch] + + save_scene_ply( + outdir / "pts_from_pointmap.ply", filtered_pts, filtered_colors + ) + print( + f" - Saved {len(filtered_pts)} filtered points to {outdir / 'pts_from_pointmap.ply'}" + ) + + # save depthmap + if "depth" in predictions and args.save_depth: + for i in range(S): + # Save both PNG (for visualization) and NPY (for actual depth values) + save_depth_png( + depth_dir / f"depth_{i:04d}.png", + predictions["depth"][0, i, :, :, 0], + ) + save_depth_npy( + depth_dir / f"depth_{i:04d}.npy", + predictions["depth"][0, i, :, :, 0], + ) + print( + f" - Saved {S} depth maps to {depth_dir} (both PNG and NPY formats)" + ) + + # save normalmap + if "normals" in predictions and args.save_normal: + for i in range(S): + save_normal_png( + normal_dir / f"normal_{i:04d}.png", predictions["normals"][0, i] + ) + print(f" - Saved {S} normal maps to {normal_dir}") + + # Save Gaussians PLY and render video + if "splats" in predictions and args.save_gs: + # Get Gaussian parameters (already filtered by GaussianSplatRenderer) + means = predictions["splats"]["means"][0].reshape(-1, 3) + scales = predictions["splats"]["scales"][0].reshape(-1, 3) + quats = predictions["splats"]["quats"][0].reshape(-1, 4) + colors = ( + predictions["splats"]["sh"][0] + if "sh" in predictions["splats"] + else predictions["splats"]["colors"][0] + ).reshape(-1, 3) + opacities = predictions["splats"]["opacities"][0].reshape(-1) + + # Save Gaussian PLY + ply_path = outdir / "gaussians.ply" + save_gs_ply( + ply_path, + means, + scales, + quats, + colors, + opacities, + ) + + # Render video using the same filtered splats from predictions + num_views = S + if args.save_rendered: + e4x4 = predictions["camera_poses"] + k3x3 = predictions["camera_intrs"] + render_interpolated_video( + model.gs_renderer, + predictions["splats"], + e4x4, + k3x3, + (H, W), + outdir / "rendered", + interp_per_pair=15, + loop_reverse=num_views == 1, + ) + print(f" - Saved rendered.mp4 to {outdir}") + else: + print(f"⚠️ Not set --save_rendered flag, skipping video rendering") + + # Build and export COLMAP reconstruction (images + sparse) + if args.save_colmap: + print("Computing filter mask for COLMAP reconstruction...") + + final_width, final_height = new_width, new_height + print(f"colmap_width: {final_width}, colmap_height: {final_height}") + + # Prepare data for mask computation (reuse from pointmap if not already computed) + if not ("pts3d" in predictions and args.save_pointmap): + pts3d_conf_np = ( + predictions["pts3d_conf"][0].detach().cpu().numpy() + ) # [S, H, W] + depth_preds_np = ( + predictions["depth"][0].detach().cpu().numpy() + ) # [S, H, W, 1] + normal_preds_np = ( + predictions["normals"][0].detach().cpu().numpy() + ) # [S, H, W, 3] + + # Compute comprehensive filter mask + final_mask = create_filter_mask( + pts3d_conf=pts3d_conf_np, + depth_preds=depth_preds_np, + normal_preds=normal_preds_np, + sky_mask=sky_mask, + confidence_percentile=args.confidence_percentile, + edge_normal_threshold=args.edge_normal_threshold, + edge_depth_threshold=args.edge_depth_threshold, + apply_confidence_mask=args.apply_confidence_mask, + apply_edge_mask=args.apply_edge_mask, + apply_sky_mask=args.apply_sky_mask, + ) # [S, H, W] + + # Prepare extrinsics/intrinsics (camera-from-world) using resized image size + e3x4, intr = vector_to_camera_matrices( + predictions["camera_params"], image_hw=(final_height, final_width) + ) + _, intr_resize = vector_to_camera_matrices( + predictions["camera_params"], image_hw=(H, W) + ) + extrinsics = e3x4[0] # [S,3,4] + intrinsics = intr[0] # [S,3,3] + intrinsics_resize = intr_resize[0] # [S,3,3] + + points_list = [] + colors_list = [] + xyf_list = [] + + # Precompute pixel coordinate grid (XYF) like demo_colmap + xyf_grid = create_pixel_coordinate_grid( + num_frames=S, height=H, width=W + ) # [S,H,W,3] float32 + xyf_grid = xyf_grid.astype(np.int32) + + # Calculate scaling factors to map from processed to resized coordinates + scale_x = final_width / W + scale_y = final_height / H + + # Use the SAME coordinate transformation as GaussianSplatRenderer.prepare_splats + # to ensure consistency between Gaussian PLY and depth-based sparse points + for i in range(S): + d = predictions["depth"][0, i, :, :, 0] + w2c = extrinsics[i][:3, :4] # [3, 4] camera-to-world + w2c = torch.cat( + [w2c, torch.tensor([[0, 0, 0, 1]], device=w2c.device)], dim=0 + ) # [4,4] + c2w = torch.linalg.inv(w2c)[:3, :4] # [4,4] + K = intrinsics_resize[i] + pts_i, _, mask = depth_to_world_coords_points( + d[None], c2w[None], K[None] + ) + + img_colors = (imgs[0, i].permute(1, 2, 0) * 255).to(torch.uint8) + + # Apply filter mask from mask computation + filter_mask_frame = torch.from_numpy(final_mask[i]).to( + mask.device + ) # [H, W] + valid = ( + mask[0] & filter_mask_frame + ) # Combine depth validity with filter mask + + if valid.sum().item() == 0: + continue + xyf_np = xyf_grid[i][valid.cpu().numpy()] # [N,3] int32 + xyf_list.append(torch.from_numpy(xyf_np).to(valid.device)) + points_list.append(pts_i[0][valid]) + colors_list.append(img_colors[valid]) + + all_pts = torch.cat(points_list, dim=0) + all_cols = torch.cat(colors_list, dim=0) + all_xyf = torch.cat(xyf_list, dim=0) + + # Convert to numpy + extrinsics = extrinsics.detach().cpu().numpy() + intrinsics = intrinsics.detach().cpu().numpy() + f_pts = all_pts.detach().cpu().to(torch.float32).numpy() + f_cols = all_cols.detach().cpu().to(torch.uint8).numpy() + f_xyf = all_xyf.detach().cpu().to(torch.int32).numpy() + + # Scale 2D coordinates from processed image to resized image resolution (if still valid) + f_xyf[:, 0] = (f_xyf[:, 0] * scale_x).astype(np.int32) # x coordinates + f_xyf[:, 1] = (f_xyf[:, 1] * scale_y).astype(np.int32) # y coordinates + + # Build reconstruction using pycolmap (PINHOLE) with resized image size + # Standard COLMAP reconstruction with 2D-3D correspondences + image_size = np.array([final_width, final_height]) + reconstruction = build_pycolmap_reconstruction( + points=f_pts, + pixel_coords=f_xyf, + point_colors=f_cols, + poses=extrinsics, + intrinsics=intrinsics, + image_size=image_size, + shared_camera_model=False, + camera_model="SIMPLE_PINHOLE", + ) + + # Update image names to match saved files + for pyimageid in reconstruction.images: + reconstruction.images[pyimageid].name = processed_image_names[ + pyimageid - 1 + ] + + # Write BIN + reconstruction.write(str(sparse_dir)) + # Save points3D.ply + save_points_ply(sparse_dir / "points3D.ply", f_pts, f_cols) + + print(f" - Saved COLMAP BIN and points3D.ply to {sparse_dir}") + + +if __name__ == "__main__": + main() diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/__init__.py b/worldcompass/reward_function/HunyuanWorldMirror/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/__init__.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/camera_head.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/camera_head.py new file mode 100644 index 0000000000000000000000000000000000000000..f23aa4ebb83c15cf9bf46ec969bfa9c2f4e7bba2 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/camera_head.py @@ -0,0 +1,190 @@ +# inspired by https://github.com/facebookresearch/vggt/blob/main/src/models/heads/camera_head.py +import torch +import torch.nn as nn +import torch.nn.functional as F + +from reward_function.HunyuanWorldMirror.src.models.layers import Mlp +from reward_function.HunyuanWorldMirror.src.models.layers.block import Block + + +class CameraHead(nn.Module): + """ + Camera head module: predicts camera parameters from token representations using iterative refinement + + Processes dedicated camera tokens through a series of transformer blocks + """ + + def __init__( + self, + dim_in: int = 2048, + trunk_depth: int = 4, + num_heads: int = 16, + mlp_ratio: int = 4, + init_values: float = 0.01, + trans_act: str = "linear", + quat_act: str = "linear", + fl_act: str = "relu", + ): + super().__init__() + + self.out_dim = 9 + self.trans_act = trans_act + self.quat_act = quat_act + self.fl_act = fl_act + self.depth = trunk_depth + + # Build refinement network using transformer block sequence + self.refine_net = nn.Sequential( + *[ + Block( + dim=dim_in, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + init_values=init_values, + ) + for _ in range(trunk_depth) + ] + ) + + # Normalization for camera tokens and network output + self.token_norm = nn.LayerNorm(dim_in) + self.out_norm = nn.LayerNorm(dim_in) + + # Learnable initial camera parameter token + self.init_token = nn.Parameter(torch.zeros(1, 1, self.out_dim)) + self.param_embed = nn.Linear(self.out_dim, dim_in) + + # Generate adaptive normalization parameters: shift, scale, and gate + self.adapt_norm_gen = nn.Sequential( + nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True) + ) + + # Adaptive layer normalization (no learnable parameters) + self.adapt_norm = nn.LayerNorm( + dim_in, elementwise_affine=False, eps=1e-6 + ) + self.param_predictor = Mlp( + in_features=dim_in, + hidden_features=dim_in // 2, + out_features=self.out_dim, + drop=0, + ) + + def forward(self, feat_seq: list, steps: int = 4) -> list: + """Forward pass to predict camera parameters. + + Args: + feat_seq: List of token tensors from network, last one used for prediction + steps: Number of iterative refinement steps, default 4 + + Returns: + List of predicted camera encodings (post-activation) from each iteration + """ + # Use tokens from last block for camera prediction + latest_feat = feat_seq[-1] + + # Extract camera tokens + cam_tokens = latest_feat[:, :, 0] + cam_tokens = self.token_norm(cam_tokens) + + # Iteratively refine camera pose predictions + b, seq_len, feat_dim = cam_tokens.shape # seq_len expected to be 1 + curr_pred = None + pred_seq = [] + + for step in range(steps): + # Use learned initial token for first iteration + if curr_pred is None: + net_input = self.param_embed( + self.init_token.expand(b, seq_len, -1) + ) + else: + curr_pred = curr_pred.detach() + net_input = self.param_embed(curr_pred) + norm_shift, norm_scale, norm_gate = self.adapt_norm_gen( + net_input + ).chunk(3, dim=-1) + mod_cam_feat = norm_gate * self.apply_adaptive_modulation( + self.adapt_norm(cam_tokens), norm_shift, norm_scale + ) + mod_cam_feat = mod_cam_feat + cam_tokens + + proc_feat = self.refine_net(mod_cam_feat) + param_delta = self.param_predictor(self.out_norm(proc_feat)) + + if curr_pred is None: + curr_pred = param_delta + else: + curr_pred = curr_pred + param_delta + + # Apply final activation functions for translation, quaternion, and field-of-view + activated_params = self.apply_camera_parameter_activation(curr_pred) + pred_seq.append(activated_params) + + return pred_seq + + def apply_camera_parameter_activation( + self, params: torch.Tensor + ) -> torch.Tensor: + """Apply activation functions to camera parameter components. + + Args: + params: Tensor containing camera parameters [translation, quaternion, focal_length] + + Returns: + Activated camera parameters tensor + """ + trans_vec = params[..., :3] + quat_vec = params[..., 3:7] + fl_vec = params[..., 7:] # or field of view + + trans_vec = self.apply_parameter_activation(trans_vec, self.trans_act) + quat_vec = self.apply_parameter_activation(quat_vec, self.quat_act) + fl_vec = self.apply_parameter_activation(fl_vec, self.fl_act) + + activated_params = torch.cat([trans_vec, quat_vec, fl_vec], dim=-1) + return activated_params + + def apply_parameter_activation( + self, tensor: torch.Tensor, act_type: str + ) -> torch.Tensor: + """Apply specified activation function to parameter tensor. + + Args: + tensor: Tensor containing parameter values + act_type: Activation type ("linear", "inv_log", "exp", "relu") + + Returns: + Activated parameter tensor + """ + if act_type == "linear": + return tensor + elif act_type == "inv_log": + return self.apply_inverse_logarithm_transform(tensor) + elif act_type == "exp": + return torch.exp(tensor) + elif act_type == "relu": + return F.relu(tensor) + else: + raise ValueError(f"Unknown activation_type: {act_type}") + + def apply_inverse_logarithm_transform( + self, x: torch.Tensor + ) -> torch.Tensor: + """ + Apply inverse logarithm transform: sign(y) * (exp(|y|) - 1) + + Args: + x: Input tensor + + Returns: + Transformed tensor + """ + return torch.sign(x) * (torch.expm1(torch.abs(x))) + + def apply_adaptive_modulation( + self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor: + """Apply adaptive modulation to input tensor using scaling and shifting parameters.""" + # Modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19 + return x * (1 + scale) + shift diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/dense_head.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/dense_head.py new file mode 100644 index 0000000000000000000000000000000000000000..0c7e7aeed0589337433e103766783f33b7d7903c --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/heads/dense_head.py @@ -0,0 +1,734 @@ +# inspired by https://github.com/DepthAnything/Depth-Anything-V2 +from typing import List, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from reward_function.HunyuanWorldMirror.src.models.utils.grid import ( + create_uv_grid, + position_grid_to_embed, +) + + +class DPTHead(nn.Module): + """# DPT Head for dense prediction tasks. + + # This module implements the DPT (Dense Prediction Transformer) head as proposed in # "Vision + Transformers for Dense Prediction" ( + https://arxiv.org/abs/2103.13413). + # It takes features from a vision transformer backbone and generates dense (per-pixel) predictions + # by fusing multi-scale features through a series of projection, upsampling, and refinement blocks. + + # Args: + # dim_in (int): Number of input feature channels. + # patch_size (int, optional): Patch size used by the backbone, default is 14. + # output_dim (int, optional): Number of output channels, default is 4. + # activation (str, optional): Activation function type for the output head, default is "inv_log". + # conf_activation (str, optional): Activation function type for the confidence/output uncertainty head, default is "expp1". + # features (int, optional): Number of channels used in intermediate feature representations, default is 256. + # out_channels (List[int], optional): Number of channels for each intermediate multi-scale feature. + # intermediate_layer_idx (List[int], optional): Indices specifying which backbone layers to use for multi-scale fusion. + # pos_embed (bool, optional): Whether to add positional encoding to the features, default is True. + # feature_only (bool, optional): If True, only return intermediate features (skip final prediction and activations). + # down_ratio (int, optional): Downsampling ratio of the output predictions, default is 1 (no downsampling). + """ + + def __init__( + self, + dim_in: int, + patch_size: int = 14, + output_dim: int = 4, + activation: str = "inv_log+expp1", + features: int = 256, + out_channels: List[int] = [256, 512, 1024, 1024], + pos_embed: bool = True, + down_ratio: int = 1, + is_gsdpt: bool = False, + ) -> None: + super(DPTHead, self).__init__() + self.patch_size = patch_size + self.activation = activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + self.is_gsdpt = is_gsdpt + + self.norm = nn.LayerNorm(dim_in) + # Projection layers for each output channel from tokens. + self.projects = nn.ModuleList( + [ + nn.Conv2d( + in_channels=dim_in, + out_channels=oc, + kernel_size=1, + stride=1, + padding=0, + ) + for oc in out_channels + ] + ) + # Resize layers for upsampling feature maps. + self.resize_layers = nn.ModuleList( + [ + nn.ConvTranspose2d( + in_channels=out_channels[0], + out_channels=out_channels[0], + kernel_size=4, + stride=4, + padding=0, + ), + nn.ConvTranspose2d( + in_channels=out_channels[1], + out_channels=out_channels[1], + kernel_size=2, + stride=2, + padding=0, + ), + nn.Identity(), + nn.Conv2d( + in_channels=out_channels[3], + out_channels=out_channels[3], + kernel_size=3, + stride=2, + padding=1, + ), + ] + ) + self.scratch = _make_scratch(out_channels, features, expand=False) + + # Attach additional modules to scratch. + self.scratch.stem_transpose = None + + self.scratch.refinenet1 = _make_fusion_block(features) + self.scratch.refinenet2 = _make_fusion_block(features) + self.scratch.refinenet3 = _make_fusion_block(features) + self.scratch.refinenet4 = _make_fusion_block( + features, has_residual=False + ) + + head_features_1 = features + head_features_2 = 32 + + if self.is_gsdpt: + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, + head_features_1 // 2, + kernel_size=3, + stride=1, + padding=1, + ) + conv2_in_channels = head_features_1 // 2 + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d( + conv2_in_channels, + head_features_2, + kernel_size=3, + stride=1, + padding=1, + ), + nn.ReLU(inplace=True), + nn.Conv2d( + head_features_2, + output_dim, + kernel_size=1, + stride=1, + padding=0, + ), + ) + self.input_merger = nn.Sequential( + nn.Conv2d(3, conv2_in_channels, 7, 1, 3), nn.ReLU() + ) + else: + self.scratch.output_conv1 = nn.Conv2d( + head_features_1, + head_features_1 // 2, + kernel_size=3, + stride=1, + padding=1, + ) + conv2_in_channels = head_features_1 // 2 + self.scratch.output_conv2 = nn.Sequential( + nn.Conv2d( + conv2_in_channels, + head_features_2, + kernel_size=3, + stride=1, + padding=1, + ), + nn.ReLU(inplace=True), + nn.Conv2d( + head_features_2, + output_dim, + kernel_size=1, + stride=1, + padding=0, + ), + ) + + def forward( + self, + token_list: List[torch.Tensor], + images: torch.Tensor, + patch_start_idx: int, + frames_chunk_size: int = 8, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Forward pass with optional frame chunking for memory efficiency. + + Args: + token_list: List of token tensors from transformer, each [B, N, C] + images: Input images [B, S, 3, H, W], range [0, 1] + patch_start_idx: Starting index of patch tokens + frames_chunk_size: Number of frames per chunk. If None or >= S, process all at once + gradient_checkpoint: Whether to use gradient checkpointing + + Returns: + For is_gsdpt: predictions [B, S, ...] + Otherwise: (predictions, confidence), [B, S, X, H, W] and [B, S, 1, H, W] + """ + B, S, _, H, W = images.shape + + # Process all frames together if chunk size not specified or large enough + if frames_chunk_size is None or frames_chunk_size >= S: + return self._forward_impl(token_list, images, patch_start_idx) + + assert frames_chunk_size > 0 + + # Process frames in chunks + preds_chunks = [] + conf_chunks = [] + gs_chunks = [] + + for frame_start in range(0, S, frames_chunk_size): + frame_end = min(frame_start + frames_chunk_size, S) + + if self.is_gsdpt: + gs, preds, conf = self._forward_impl( + token_list, images, patch_start_idx, frame_start, frame_end + ) + gs_chunks.append(gs) + preds_chunks.append(preds) + conf_chunks.append(conf) + else: + preds, conf = self._forward_impl( + token_list, images, patch_start_idx, frame_start, frame_end + ) + preds_chunks.append(preds) + conf_chunks.append(conf) + + # Concatenate chunks along frame dimension + if self.is_gsdpt: + return ( + torch.cat(gs_chunks, dim=1), + torch.cat(preds_chunks, dim=1), + torch.cat(conf_chunks, dim=1), + ) + else: + return torch.cat(preds_chunks, dim=1), torch.cat(conf_chunks, dim=1) + + def _forward_impl( + self, + token_list: List[torch.Tensor], + images: torch.Tensor, + patch_start_idx: int, + frame_start: int = None, + frame_end: int = None, + ) -> torch.Tensor: + """Core forward implementation for DPT head. + + Args: + token_list: List of transformer tokens from each layer, [B, S, N, C] + images: Input images [B, S, 3, H, W] + patch_start_idx: Starting index of patch tokens + frame_start: Start index for frame chunking (optional) + frame_end: End index for frame chunking (optional) + + Returns: + If is_gsdpt: (features, preds, conf) + Else: (preds, conf) + """ + # Slice frames if chunking + if frame_start is not None and frame_end is not None: + images = images[:, frame_start:frame_end].contiguous() + + B, S, _, H, W = images.shape + ph = H // self.patch_size # patch height + pw = W // self.patch_size # patch width + + # Extract and project multi-level features + feats = [] + for proj, resize, tokens in zip( + self.projects, self.resize_layers, token_list + ): + # Extract patch tokens + patch_tokens = tokens[:, :, patch_start_idx:] + if frame_start is not None and frame_end is not None: + patch_tokens = patch_tokens[:, frame_start:frame_end] + + # Reshape to [B*S, N_patches, C] + patch_tokens = patch_tokens.reshape( + B * S, -1, patch_tokens.shape[-1] + ) + patch_tokens = self.norm(patch_tokens) + + # Convert to 2D feature map [B*S, C, ph, pw] + feat = patch_tokens.permute(0, 2, 1).reshape( + B * S, patch_tokens.shape[-1], ph, pw + ) + feat = proj(feat) + + if self.pos_embed: + feat = self._apply_pos_embed(feat, W, H) + feat = resize(feat) + feats.append(feat) + + # Fuse multi-level features + fused = self.scratch_forward(feats) + fused = custom_interpolate( + fused, + size=( + int(ph * self.patch_size / self.down_ratio), + int(pw * self.patch_size / self.down_ratio), + ), + mode="bilinear", + align_corners=True, + ) + + # Apply positional embedding after upsampling + if self.pos_embed: + fused = self._apply_pos_embed(fused, W, H) + + # Generate predictions and confidence + if self.is_gsdpt: + # GSDPT: output features, predictions, and confidence + out = self.scratch.output_conv2(fused) + preds, conf = self.activate_head(out, activation=self.activation) + preds = preds.reshape(B, S, *preds.shape[1:]) + conf = conf.reshape(B, S, *conf.shape[1:]) + + # Merge direct image features + img_flat = images.reshape(B * S, -1, H, W) + img_feat = self.input_merger(img_flat) + fused = fused + img_feat + fused = fused.reshape(B, S, *fused.shape[1:]) + return fused, preds, conf + else: + # Standard: output predictions and confidence + out = self.scratch.output_conv2(fused) + preds, conf = self.activate_head(out, activation=self.activation) + preds = preds.reshape(B, S, *preds.shape[1:]) + conf = conf.reshape(B, S, *conf.shape[1:]) + return preds, conf + + def _apply_pos_embed( + self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1 + ) -> torch.Tensor: + """Apply positional embedding to tensor x.""" + patch_w = x.shape[-1] + patch_h = x.shape[-2] + pos_embed = create_uv_grid( + patch_w, patch_h, aspect_ratio=W / H, dtype=x.dtype, device=x.device + ) + pos_embed = position_grid_to_embed(pos_embed, x.shape[1]) + pos_embed = pos_embed * ratio + pos_embed = pos_embed.permute(2, 0, 1)[None].expand( + x.shape[0], -1, -1, -1 + ) + return x + pos_embed + + def scratch_forward(self, features: List[torch.Tensor]) -> torch.Tensor: + """Forward pass through the fusion blocks. + + Args: + features (List[Tensor]): List of feature maps from different layers. + + Returns: + Tensor: Fused feature map. + """ + layer_1, layer_2, layer_3, layer_4 = features + + layer_1_rn = self.scratch.layer1_rn(layer_1) + layer_2_rn = self.scratch.layer2_rn(layer_2) + layer_3_rn = self.scratch.layer3_rn(layer_3) + layer_4_rn = self.scratch.layer4_rn(layer_4) + + out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) + del layer_4_rn, layer_4 + + out = self.scratch.refinenet3( + out, layer_3_rn, size=layer_2_rn.shape[2:] + ) + del layer_3_rn, layer_3 + + out = self.scratch.refinenet2( + out, layer_2_rn, size=layer_1_rn.shape[2:] + ) + del layer_2_rn, layer_2 + + out = self.scratch.refinenet1(out, layer_1_rn) + del layer_1_rn, layer_1 + + out = self.scratch.output_conv1(out) + return out + + def activate_head( + self, out_head: torch.Tensor, activation: str = "inv_log+expp1" + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Process network output to extract attribute (e.g. points, depth, etc.) and confidence + values. + + Args: + out_head: Network output tensor (B, C, H, W) + activation: Activation type for processing (e.g., "inv_log+expp1") + + Returns: + Tuple of (attribute tensor, confidence tensor) + """ + # Parse activation string + act_attr, act_conf = ( + activation.split("+") + if "+" in activation + else (activation, "expp1") + ) + + # (B,C,H,W) -> (B,H,W,C) + feat = out_head.permute(0, 2, 3, 1) + attr, conf = feat[..., :-1], feat[..., -1] + + # Map point activations to lambdas for clarity and conciseness + attr_activations = { + "norm_exp": lambda x: ( + (x / x.norm(dim=-1, keepdim=True).clamp(min=1e-8)) + * torch.expm1(x.norm(dim=-1, keepdim=True)) + ), + "norm": lambda x: x / x.norm(dim=-1, keepdim=True), + "exp": torch.exp, + "relu": F.relu, + "inv_log": self._apply_inverse_log_transform, + "xy_inv_log": lambda x: torch.cat( + [ + x[..., :2] * self._apply_inverse_log_transform(x[..., 2:]), + self._apply_inverse_log_transform(x[..., 2:]), + ], + dim=-1, + ), + "sigmoid": torch.sigmoid, + "linear": lambda x: x, + } + + if act_attr not in attr_activations: + raise ValueError(f"Unknown attribute activation: {act_attr}") + attr_out = attr_activations[act_attr](attr) + + # Confidence activation mapping + conf_activations = { + "expp1": lambda c: 1 + c.exp(), + "expp0": torch.exp, + "sigmoid": torch.sigmoid, + } + if act_conf not in conf_activations: + raise ValueError(f"Unknown confidence activation: {act_conf}") + conf_out = conf_activations[act_conf](conf) + + return attr_out, conf_out + + def _apply_inverse_log_transform( + self, input_tensor: torch.Tensor + ) -> torch.Tensor: + """ + Apply inverse logarithm transform: sign(y) * (exp(|y|) - 1) + + Args: + input_tensor: Input tensor + + Returns: + Transformed tensor + """ + return torch.sign(input_tensor) * (torch.expm1(torch.abs(input_tensor))) + + +################################################################################ +# DPT Modules +################################################################################ + + +def _make_fusion_block( + features: int, size: int = None, has_residual: bool = True, groups: int = 1 +) -> nn.Module: + return FeatureFusionBlock( + features, + nn.ReLU(inplace=True), + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=size, + has_residual=has_residual, + groups=groups, + ) + + +def _make_scratch( + in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False +) -> nn.Module: + scratch = nn.Module() + out_shape1 = out_shape + out_shape2 = out_shape + out_shape3 = out_shape + if len(in_shape) >= 4: + out_shape4 = out_shape + + if expand: + out_shape1 = out_shape + out_shape2 = out_shape * 2 + out_shape3 = out_shape * 4 + if len(in_shape) >= 4: + out_shape4 = out_shape * 8 + + scratch.layer1_rn = nn.Conv2d( + in_shape[0], + out_shape1, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + scratch.layer2_rn = nn.Conv2d( + in_shape[1], + out_shape2, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + scratch.layer3_rn = nn.Conv2d( + in_shape[2], + out_shape3, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + if len(in_shape) >= 4: + scratch.layer4_rn = nn.Conv2d( + in_shape[3], + out_shape4, + kernel_size=3, + stride=1, + padding=1, + bias=False, + groups=groups, + ) + return scratch + + +class ResidualConvUnit(nn.Module): + """Residual convolution module with skip connection.""" + + def __init__(self, features, activation, bn, groups=1): + """Initialize ResidualConvUnit. + + Args: + features (int): Number of input/output feature channels + activation: Activation function to use + bn (bool): Whether to use batch normalization (currently unused) + groups (int): Number of groups for grouped convolution + """ + super().__init__() + + self.bn = bn + self.groups = groups + self.conv1 = nn.Conv2d( + features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups, + ) + self.conv2 = nn.Conv2d( + features, + features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + groups=self.groups, + ) + + self.norm1 = None + self.norm2 = None + + self.activation = activation + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + """Forward pass with residual connection. + + Args: + x (tensor): Input tensor of shape (B, C, H, W) + + Returns: + tensor: Output tensor of shape (B, C, H, W) with residual added + """ + + out = self.activation(x) + out = self.conv1(out) + if self.norm1 is not None: + out = self.norm1(out) + + out = self.activation(out) + out = self.conv2(out) + if self.norm2 is not None: + out = self.norm2(out) + + return self.skip_add.add(out, x) + + +class FeatureFusionBlock(nn.Module): + """Feature fusion block.""" + + def __init__( + self, + features, + activation, + deconv=False, + bn=False, + expand=False, + align_corners=True, + size=None, + has_residual=True, + groups=1, + ): + """Initialize FeatureFusionBlock. + + Args: + features (int): Number of input/output feature channels + activation: Activation function to use + deconv (bool): Whether to use deconvolution + bn (bool): Whether to use batch normalization + expand (bool): Whether to expand features (halve output channels) + align_corners (bool): Align corners for interpolation + size: Target size for upsampling + has_residual (bool): Whether to include residual connection + groups (int): Number of groups for grouped convolution + """ + super(FeatureFusionBlock, self).__init__() + + self.deconv = deconv + self.align_corners = align_corners + self.groups = groups + self.expand = expand + out_features = features + if self.expand == True: + out_features = features // 2 + + self.out_conv = nn.Conv2d( + features, + out_features, + kernel_size=1, + stride=1, + padding=0, + bias=True, + groups=self.groups, + ) + + if has_residual: + self.resConfUnit1 = ResidualConvUnit( + features, activation, bn, groups=self.groups + ) + + self.has_residual = has_residual + self.resConfUnit2 = ResidualConvUnit( + features, activation, bn, groups=self.groups + ) + + self.skip_add = nn.quantized.FloatFunctional() + self.size = size + + def forward(self, *xs, size=None): + """Forward pass through the feature fusion block. + + Args: + *xs: Variable number of input tensors. First tensor is the main input, + second tensor (if present) is used for residual connection. + size: Optional target size for upsampling. If None, uses self.size or scale_factor=2. + + Returns: + torch.Tensor: Fused and upsampled output tensor. + """ + output = xs[0] + + if self.has_residual: + res = self.resConfUnit1(xs[1]) + output = self.skip_add.add(output, res) + + output = self.resConfUnit2(output) + + if (size is None) and (self.size is None): + modifier = {"scale_factor": 2} + elif size is None: + modifier = {"size": self.size} + else: + modifier = {"size": size} + + output = custom_interpolate( + output, + **modifier, + mode="bilinear", + align_corners=self.align_corners, + ) + output = self.out_conv(output) + + return output + + +def custom_interpolate( + x: torch.Tensor, + size: Tuple[int, int] = None, + scale_factor: float = None, + mode: str = "bilinear", + align_corners: bool = True, +) -> torch.Tensor: + """Custom interpolation function to handle large tensors by chunking. + + Avoids INT_MAX overflow issues in nn.functional.interpolate when dealing with + very large input tensors by splitting them into smaller chunks. + + Args: + x: Input tensor to interpolate + size: Target output size (H, W) + scale_factor: Scaling factor if size is not provided + mode: Interpolation mode (default: "bilinear") + align_corners: Whether to align corners in interpolation + + Returns: + Interpolated tensor + """ + if size is None: + size = ( + int(x.shape[-2] * scale_factor), + int(x.shape[-1] * scale_factor), + ) + + INT_MAX = 1610612736 + + input_elements = size[0] * size[1] * x.shape[0] * x.shape[1] + + if input_elements > INT_MAX: + chunks = torch.chunk(x, chunks=(input_elements // INT_MAX) + 1, dim=0) + interpolated_chunks = [ + nn.functional.interpolate( + chunk, size=size, mode=mode, align_corners=align_corners + ) + for chunk in chunks + ] + x = torch.cat(interpolated_chunks, dim=0) + return x.contiguous() + else: + return nn.functional.interpolate( + x, size=size, mode=mode, align_corners=align_corners + ) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/__init__.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2190558e6a4fd88d71ac8c731d2d0314350d5744 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/__init__.py @@ -0,0 +1,5 @@ +from .mlp import Mlp +from .patch_embed import PatchEmbed, PatchEmbed_Mlp +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused +from .block import NestedTensorBlock +from .attention import MemEffAttention diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/attention.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..262a64fac3c99d45d6878b90e05b510f30556e1e --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/attention.py @@ -0,0 +1,98 @@ +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import logging +import os +import warnings + +from torch import Tensor +from torch import nn +import torch.nn.functional as F +import torch + +# from torch.nn.attention import SDPBackend + +XFORMERS_AVAILABLE = False + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + qk_norm: bool = False, + fused_attn: bool = True, # use F.scaled_dot_product_attention or not + rope=None, + ) -> None: + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim**-0.5 + self.fused_attn = fused_attn + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x: Tensor, pos=None) -> Tensor: + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, self.head_dim) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + if self.rope is not None: + q = self.rope(q, pos) + k = self.rope(k, pos) + + # orig_dtype = q.dtype + x = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0 + ) + + # with torch.cuda.amp.autocast(dtype=torch.bfloat16): + # with nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION): + # if x.dtype != orig_dtype: + # x = x.to(orig_dtype) + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class MemEffAttention(Attention): + def forward(self, x: Tensor, attn_bias=None, pos=None) -> Tensor: + assert pos is None + if not XFORMERS_AVAILABLE: + if attn_bias is not None: + raise AssertionError( + "xFormers is required for using nested tensors" + ) + return super().forward(x) + + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + + q, k, v = unbind(qkv, 2) + + x = memory_efficient_attention(q, k, v, attn_bias=attn_bias) + x = x.reshape([B, N, C]) + + x = self.proj(x) + x = self.proj_drop(x) + return x diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/block.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..319bdbad5a092edd7bdb131809231fc20f8a0b14 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/block.py @@ -0,0 +1,301 @@ +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, List, Any, Tuple, Dict + +import torch +from torch import nn, Tensor + +from .attention import Attention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + +XFORMERS_AVAILABLE = False + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2) + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + qk_norm: bool = False, + fused_attn: bool = True, # use F.scaled_dot_product_attention or not + rope=None, + ) -> None: + super().__init__() + + self.norm1 = norm_layer(dim) + + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + qk_norm=qk_norm, + fused_attn=fused_attn, + rope=rope, + ) + + self.ls1 = ( + LayerScale(dim, init_values=init_values) + if init_values + else nn.Identity() + ) + self.drop_path1 = ( + DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + ) + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = ( + LayerScale(dim, init_values=init_values) + if init_values + else nn.Identity() + ) + self.drop_path2 = ( + DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + ) + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor, pos=None) -> Tensor: + def attn_residual_func(x: Tensor, pos=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), pos=pos)) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + # the overhead is compensated only for a drop path rate larger than 0.1 + x = drop_add_residual_stochastic_depth( + x, + pos=pos, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x, pos=pos)) + x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 + else: + x = x + attn_residual_func(x, pos=pos) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, + pos=None, +) -> Tensor: + # 1) extract subset using permutation + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + # 2) apply residual_func to get residual + if pos is not None: + # if necessary, apply rope to the subset + pos = pos[brange] + residual = residual_func(x_subset, pos=pos) + else: + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + # 3) add the residual + x_plus_residual = torch.index_add( + x_flat, + 0, + brange, + residual.to(dtype=x.dtype), + alpha=residual_scale_factor, + ) + return x_plus_residual.view_as(x) + + +def get_branges_scales(x, sample_drop_ratio=0.0): + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + residual_scale_factor = b / sample_subset_size + return brange, residual_scale_factor + + +def add_residual( + x, brange, residual, residual_scale_factor, scaling_vector=None +): + if scaling_vector is None: + x_flat = x.flatten(1) + residual = residual.flatten(1) + x_plus_residual = torch.index_add( + x_flat, + 0, + brange, + residual.to(dtype=x.dtype), + alpha=residual_scale_factor, + ) + else: + x_plus_residual = scaled_index_add( + x, + brange, + residual.to(dtype=x.dtype), + scaling=scaling_vector, + alpha=residual_scale_factor, + ) + return x_plus_residual + + +attn_bias_cache: Dict[Tuple, Any] = {} + + +def get_attn_bias_and_cat(x_list, branges=None): + """This will perform the index select, cat the tensors, and provide the attn_bias from cache.""" + batch_sizes = ( + [b.shape[0] for b in branges] + if branges is not None + else [x.shape[0] for x in x_list] + ) + all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list)) + if all_shapes not in attn_bias_cache.keys(): + seqlens = [] + for b, x in zip(batch_sizes, x_list): + for _ in range(b): + seqlens.append(x.shape[1]) + attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens) + attn_bias._batch_sizes = batch_sizes + attn_bias_cache[all_shapes] = attn_bias + + if branges is not None: + cat_tensors = index_select_cat( + [x.flatten(1) for x in x_list], branges + ).view(1, -1, x_list[0].shape[-1]) + else: + tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list) + cat_tensors = torch.cat(tensors_bs1, dim=1) + + return attn_bias_cache[all_shapes], cat_tensors + + +def drop_add_residual_stochastic_depth_list( + x_list: List[Tensor], + residual_func: Callable[[Tensor, Any], Tensor], + sample_drop_ratio: float = 0.0, + scaling_vector=None, +) -> Tensor: + # 1) generate random set of indices for dropping samples in the batch + branges_scales = [ + get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) + for x in x_list + ] + branges = [s[0] for s in branges_scales] + residual_scale_factors = [s[1] for s in branges_scales] + + # 2) get attention bias and index+concat the tensors + attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges) + + # 3) apply residual_func to get residual, and split the result + residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore + + outputs = [] + for x, brange, residual, residual_scale_factor in zip( + x_list, branges, residual_list, residual_scale_factors + ): + outputs.append( + add_residual( + x, brange, residual, residual_scale_factor, scaling_vector + ).view_as(x) + ) + return outputs + + +class NestedTensorBlock(Block): + def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]: + """x_list contains a list of tensors to nest together and run.""" + assert isinstance(self.attn, MemEffAttention) + + if self.training and self.sample_drop_ratio > 0.0: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.attn(self.norm1(x), attn_bias=attn_bias) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.mlp(self.norm2(x)) + + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=( + self.ls1.gamma if isinstance(self.ls1, LayerScale) else None + ), + ) + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=( + self.ls2.gamma if isinstance(self.ls1, LayerScale) else None + ), + ) + return x_list + else: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias)) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + attn_bias, x = get_attn_bias_and_cat(x_list) + x = x + attn_residual_func(x, attn_bias=attn_bias) + x = x + ffn_residual_func(x) + return attn_bias.split(x) + + def forward(self, x_or_x_list): + if isinstance(x_or_x_list, Tensor): + return super().forward(x_or_x_list) + elif isinstance(x_or_x_list, list): + if not XFORMERS_AVAILABLE: + raise AssertionError( + "xFormers is required for using nested tensors" + ) + return self.forward_nested(x_or_x_list) + else: + raise AssertionError diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/drop_path.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/drop_path.py new file mode 100644 index 0000000000000000000000000000000000000000..7f3dc1996a8f0d1ec9541d781a967e7f74f838a6 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/drop_path.py @@ -0,0 +1,31 @@ +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * ( + x.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/layer_scale.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..85054796555e16d4db3a47605d94939b7664be2d --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/layer_scale.py @@ -0,0 +1,22 @@ +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 + +from typing import Union + +import torch +from torch import Tensor +from torch import nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/mlp.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..96dd129eca9a582a416e3be35ea3174b6268b174 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/mlp.py @@ -0,0 +1,35 @@ +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional + +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/patch_embed.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..21e8a5a61d1d8d6bfa81dccf38d4e15b65dcf0c3 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/patch_embed.py @@ -0,0 +1,196 @@ +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union + +import torch +from torch import Tensor +import torch.nn as nn +import torch.nn.functional as F +from itertools import repeat +import collections.abc + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW + ) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert ( + H % patch_H == 0 + ), f"Input image height {H} is not a multiple of patch height {patch_H}" + assert ( + W % patch_W == 0 + ), f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + +class PatchEmbed_Mlp(PatchEmbed): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten_embedding=True, + ): + super().__init__( + img_size, + patch_size, + in_chans, + embed_dim, + norm_layer, + flatten_embedding, + ) + + self.proj = nn.Sequential( + PixelUnshuffle(patch_size), + Permute((0, 2, 3, 1)), + Mlp(in_chans * patch_size**2, 4 * embed_dim, embed_dim), + Permute((0, 3, 1, 2)), + ) + + +class PixelUnshuffle(nn.Module): + def __init__(self, downscale_factor): + super().__init__() + self.downscale_factor = downscale_factor + + def forward(self, input): + if input.numel() == 0: + # this is not in the original torch implementation + C, H, W = input.shape[-3:] + assert ( + H + and W + and H % self.downscale_factor == W % self.downscale_factor == 0 + ) + return input.view( + *input.shape[:-3], + C * self.downscale_factor**2, + H // self.downscale_factor, + W // self.downscale_factor, + ) + else: + return F.pixel_unshuffle(input, self.downscale_factor) + + +class Permute(torch.nn.Module): + dims: tuple[int, ...] + + def __init__(self, dims: tuple[int, ...]) -> None: + super().__init__() + self.dims = tuple(dims) + + def __repr__(self): + return f"Permute{self.dims}" + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return input.permute(*self.dims) + + +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return x + return tuple(repeat(x, n)) + + return parse + + +to_2tuple = _ntuple(2) + + +class Mlp(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks.""" + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + bias=True, + drop=0.0, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0]) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1]) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.fc2(x) + x = self.drop2(x) + return x diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/rope.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..aa04dc3e7d31d98164f5d35b0517eb06b3027c85 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/rope.py @@ -0,0 +1,210 @@ +# Implementation of 2D Rotary Position Embeddings (RoPE). + +# This module provides a clean implementation of 2D Rotary Position Embeddings, +# which extends the original RoPE concept to handle 2D spatial positions. + +# Inspired by: +# https://github.com/meta-llama/codellama/blob/main/llama/model.py +# https://github.com/naver-ai/rope-vit + + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Dict, Tuple + + +class PositionGetter: + """Generates and caches 2D spatial positions for patches in a grid. + + This class efficiently manages the generation of spatial coordinates for patches + in a 2D grid, caching results to avoid redundant computations. + + Attributes: + position_cache: Dictionary storing precomputed position tensors for different + grid dimensions. + """ + + def __init__(self): + """Initializes the position generator with an empty cache.""" + self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {} + + def __call__( + self, batch_size: int, height: int, width: int, device: torch.device + ) -> torch.Tensor: + """Generates spatial positions for a batch of patches. + + Args: + batch_size: Number of samples in the batch. + height: Height of the grid in patches. + width: Width of the grid in patches. + device: Target device for the position tensor. + + Returns: + Tensor of shape (batch_size, height*width, 2) containing y,x coordinates + for each position in the grid, repeated for each batch item. + """ + if (height, width) not in self.position_cache: + y_coords = torch.arange(height, device=device) + x_coords = torch.arange(width, device=device) + positions = torch.cartesian_prod(y_coords, x_coords) + self.position_cache[height, width] = positions + + cached_positions = self.position_cache[height, width] + return ( + cached_positions.view(1, height * width, 2) + .expand(batch_size, -1, -1) + .clone() + ) + + +class RotaryPositionEmbedding2D(nn.Module): + """2D Rotary Position Embedding implementation. + + This module applies rotary position embeddings to input tokens based on their + 2D spatial positions. It handles the position-dependent rotation of features + separately for vertical and horizontal dimensions. + + Args: + frequency: Base frequency for the position embeddings. Default: 100.0 + scaling_factor: Scaling factor for frequency computation. Default: 1.0 + + Attributes: + base_frequency: Base frequency for computing position embeddings. + scaling_factor: Factor to scale the computed frequencies. + frequency_cache: Cache for storing precomputed frequency components. + """ + + def __init__( + self, + frequency: float = 100.0, + scaling_factor: float = 1.0, + ): + """Initializes the 2D RoPE module.""" + super().__init__() + self.base_frequency = frequency + self.scaling_factor = scaling_factor + self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = ( + {} + ) + + def _compute_frequency_components( + self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Computes frequency components for rotary embeddings. + + Args: + dim: Feature dimension (must be even). + seq_len: Maximum sequence length. + device: Target device for computations. + dtype: Data type for the computed tensors. + + Returns: + Tuple of (cosine, sine) tensors for frequency components. + """ + cache_key = (dim, seq_len, device, dtype) + if cache_key not in self.frequency_cache: + # Compute frequency bands + exponents = torch.arange(0, dim, 2, device=device).float() / dim + inv_freq = 1.0 / (self.base_frequency**exponents) + + # Generate position-dependent frequencies + positions = torch.arange( + seq_len, device=device, dtype=inv_freq.dtype + ) + angles = torch.einsum("i,j->ij", positions, inv_freq) + + # Compute and cache frequency components + angles = angles.to(dtype) + angles = torch.cat((angles, angles), dim=-1) + cos_components = angles.cos().to(dtype) + sin_components = angles.sin().to(dtype) + self.frequency_cache[cache_key] = (cos_components, sin_components) + + return self.frequency_cache[cache_key] + + @staticmethod + def _rotate_features(x: torch.Tensor) -> torch.Tensor: + """Performs feature rotation by splitting and recombining feature dimensions. + + Args: + x: Input tensor to rotate. + + Returns: + Rotated feature tensor. + """ + feature_dim = x.shape[-1] + x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_1d_rope( + self, + tokens: torch.Tensor, + positions: torch.Tensor, + cos_comp: torch.Tensor, + sin_comp: torch.Tensor, + ) -> torch.Tensor: + """Applies 1D rotary position embeddings along one dimension. + + Args: + tokens: Input token features. + positions: Position indices. + cos_comp: Cosine components for rotation. + sin_comp: Sine components for rotation. + + Returns: + Tokens with applied rotary position embeddings. + """ + # Embed positions with frequency components + cos = F.embedding(positions, cos_comp)[:, None, :, :] + sin = F.embedding(positions, sin_comp)[:, None, :, :] + + # Apply rotation + return (tokens * cos) + (self._rotate_features(tokens) * sin) + + def forward( + self, tokens: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + """Applies 2D rotary position embeddings to input tokens. + + Args: + tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim). + The feature dimension (dim) must be divisible by 4. + positions: Position tensor of shape (batch_size, n_tokens, 2) containing + the y and x coordinates for each token. + + Returns: + Tensor of same shape as input with applied 2D rotary position embeddings. + + Raises: + AssertionError: If input dimensions are invalid or positions are malformed. + """ + # Validate inputs + assert tokens.size(-1) % 2 == 0, "Feature dimension must be even" + assert ( + positions.ndim == 3 and positions.shape[-1] == 2 + ), "Positions must have shape (batch_size, n_tokens, 2)" + + # Compute feature dimension for each spatial direction + feature_dim = tokens.size(-1) // 2 + + # Get frequency components + max_position = int(positions.max()) + 1 + cos_comp, sin_comp = self._compute_frequency_components( + feature_dim, max_position, tokens.device, tokens.dtype + ) + + # Split features for vertical and horizontal processing + vertical_features, horizontal_features = tokens.chunk(2, dim=-1) + + # Apply RoPE separately for each dimension + vertical_features = self._apply_1d_rope( + vertical_features, positions[..., 0], cos_comp, sin_comp + ) + horizontal_features = self._apply_1d_rope( + horizontal_features, positions[..., 1], cos_comp, sin_comp + ) + + # Combine processed features + return torch.cat((vertical_features, horizontal_features), dim=-1) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/swiglu_ffn.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/swiglu_ffn.py new file mode 100644 index 0000000000000000000000000000000000000000..65adc6cb3bc5b95502971bc939e195cbd5ae7313 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/swiglu_ffn.py @@ -0,0 +1,67 @@ +import os +from typing import Callable, Optional +import warnings + +from torch import Tensor, nn +import torch.nn.functional as F + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None +# try: +# if XFORMERS_ENABLED: +# from xformers.ops import SwiGLU + +# XFORMERS_AVAILABLE = True +# warnings.warn("xFormers is available (SwiGLU)") +# else: +# warnings.warn("xFormers is disabled (SwiGLU)") +# raise ImportError +# except ImportError: +SwiGLU = SwiGLUFFN +XFORMERS_AVAILABLE = False + +# warnings.warn("xFormers is not available (SwiGLU)") + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/vision_transformer.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..60833f808940283f5cad63c70dc8367b54f895be --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/layers/vision_transformer.py @@ -0,0 +1,455 @@ +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +from functools import partial +import math +import logging +from typing import Sequence, Tuple, Union, Callable + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint +from torch.nn.init import trunc_normal_ +from . import ( + Mlp, + PatchEmbed, + SwiGLUFFNFused, + MemEffAttention, + NestedTensorBlock as Block, +) + +log = logging.getLogger(__name__) + + +def named_apply( + fn: Callable, + module: nn.Module, + name="", + depth_first=True, + include_root=False, +) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, + module=child_module, + name=child_name, + depth_first=depth_first, + include_root=True, + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=None, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + qk_norm=False, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + drop_path_rate (float): stochastic depth rate + drop_path_uniform (bool): apply uniform drop rate across blocks + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings + """ + super().__init__() + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.num_features = self.embed_dim = ( + embed_dim # num_features for consistency with other models + ) + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + self.use_reentrant = False # hardcoded to False + + self.patch_embed = embed_layer( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter( + torch.zeros(1, num_patches + self.num_tokens, embed_dim) + ) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) + if num_register_tokens + else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + + if ffn_layer == "mlp": + log.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + log.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + log.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + qk_norm=qk_norm, + ) + for i in range(depth) + ] + if block_chunks > 0: + self.chunked_blocks = True + chunked_blocks = [] + chunksize = depth // block_chunks + for i in range(0, depth, chunksize): + # this is to keep the block index consistent if we chunk the block list + chunked_blocks.append( + [nn.Identity()] * i + blocks_list[i : i + chunksize] + ) + self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) + else: + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + self.norm = norm_layer(embed_dim) + self.head = nn.Identity() + + self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) + + self.init_weights() + + def init_weights(self): + trunc_normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.cls_token, std=1e-6) + if self.register_tokens is not None: + nn.init.normal_(self.register_tokens, std=1e-6) + named_apply(init_weights_vit_timm, self) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if self.interpolate_offset: + # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sx, sy) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (w0, h0) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + assert (w0, h0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat( + (class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1 + ).to(previous_dtype) + + def prepare_tokens_with_masks(self, x, masks=None): + B, nc, w, h = x.shape + x = self.patch_embed(x) + if masks is not None: + x = torch.where( + masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x + ) + + x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + + return x + + def forward_features_list(self, x_list, masks_list): + x = [ + self.prepare_tokens_with_masks(x, masks) + for x, masks in zip(x_list, masks_list) + ] + + for blk in self.blocks: + if self.training: + x = checkpoint(blk, x, use_reentrant=self.use_reentrant) + else: + x = blk(x) + + all_x = x + output = [] + for x, masks in zip(all_x, masks_list): + x_norm = self.norm(x) + output.append( + { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[ + :, 1 : self.num_register_tokens + 1 + ], + "x_norm_patchtokens": x_norm[ + :, self.num_register_tokens + 1 : + ], + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x, masks=None): + if isinstance(x, list): + return self.forward_features_list(x, masks) + + x = self.prepare_tokens_with_masks(x, masks) + + for blk in self.blocks: + if self.training: + x = checkpoint(blk, x, use_reentrant=self.use_reentrant) + else: + x = blk(x) + + x_norm = self.norm(x) + return { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + + def _get_intermediate_layers_not_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = ( + range(total_block_len - n, total_block_len) + if isinstance(n, int) + else n + ) + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in blocks_to_take: + output.append(x) + assert len(output) == len( + blocks_to_take + ), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def _get_intermediate_layers_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + output, i, total_block_len = [], 0, len(self.blocks[-1]) + # If n is an int, take the n last blocks. If it's a list, take them + blocks_to_take = ( + range(total_block_len - n, total_block_len) + if isinstance(n, int) + else n + ) + for block_chunk in self.blocks: + for blk in block_chunk[i:]: # Passing the nn.Identity() + x = blk(x) + if i in blocks_to_take: + output.append(x) + i += 1 + assert len(output) == len( + blocks_to_take + ), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + norm=True, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + if self.chunked_blocks: + outputs = self._get_intermediate_layers_chunked(x, n) + else: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs] + if reshape: + B, _, w, h = x.shape + outputs = [ + out.reshape(B, w // self.patch_size, h // self.patch_size, -1) + .permute(0, 3, 1, 2) + .contiguous() + for out in outputs + ] + if return_class_token: + return tuple(zip(outputs, class_tokens)) + return tuple(outputs) + + def forward(self, *args, is_training=True, **kwargs): + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + +def vit_small(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): + """Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64.""" + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=40, + num_heads=24, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/rasterization.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/rasterization.py new file mode 100644 index 0000000000000000000000000000000000000000..110033050181054f35200215e94e1cda5f337999 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/rasterization.py @@ -0,0 +1,656 @@ +from typing import Dict, Tuple +import numpy as np +import torch +import torch.nn as nn +from torch import Tensor +from einops import rearrange + +from gsplat.rendering import rasterization +from gsplat.strategy import DefaultStrategy + +from reward_function.HunyuanWorldMirror.src.models.utils.frustum import ( + calculate_unprojected_mask, +) +from reward_function.HunyuanWorldMirror.src.models.utils.geometry import ( + depth_to_world_coords_points, +) +from reward_function.HunyuanWorldMirror.src.models.utils import sh_utils, act_gs + + +class Rasterizer: + def __init__( + self, + rasterization_mode="classic", + packed=True, + abs_grad=True, + with_eval3d=False, + camera_model="pinhole", + sparse_grad=False, + distributed=False, + grad_strategy=DefaultStrategy, + ): + self.rasterization_mode = rasterization_mode + self.packed = packed + self.abs_grad = abs_grad + self.camera_model = camera_model + self.sparse_grad = sparse_grad + self.grad_strategy = grad_strategy + self.distributed = distributed + self.with_eval3d = with_eval3d + + def rasterize_splats( + self, + means, + quats, + scales, + opacities, + colors, + camtoworlds: Tensor, + Ks: Tensor, + width: int, + height: int, + **kwargs, + ) -> Tuple[Tensor, Tensor, Dict]: + render_colors, render_alphas, _ = rasterization( + means=means, + quats=quats, + scales=scales, + opacities=opacities, + colors=colors, + viewmats=torch.linalg.inv(camtoworlds), # [C, 4, 4] + Ks=Ks, # [C, 3, 3] + width=width, + height=height, + packed=self.packed, + absgrad=( + self.abs_grad + if isinstance(self.grad_strategy, DefaultStrategy) + else False + ), + sparse_grad=self.sparse_grad, + rasterize_mode=self.rasterization_mode, + distributed=self.distributed, + camera_model=self.camera_model, + with_eval3d=self.with_eval3d, + render_mode="RGB+ED", + **kwargs, + ) + return render_colors[..., :3], render_colors[..., 3:], render_alphas + + def rasterize_batches( + self, + means, + quats, + scales, + opacities, + colors, + viewmats, + Ks, + width, + height, + **kwargs, + ): + rendered_colors, rendered_depths, rendered_alphas = [], [], [] + batch_size = len(means) + + for i in range(batch_size): + means_i = means[i] # [N, 4] + quats_i = quats[i] # [N, 4] + scales_i = scales[i] # [N, 3] + opacities_i = opacities[i] # [N,] + colors_i = colors[i] # [N, 3] + viewmats_i = viewmats[i] # [V, 4, 4] + Ks_i = Ks[i] # [V, 3, 3] + + render_colors_i, render_depths_i, render_alphas_i = ( + self.rasterize_splats( + means_i, + quats_i, + scales_i, + opacities_i, + colors_i, + viewmats_i, + Ks_i, + width, + height, + **kwargs, + ) + ) + + rendered_colors.append(render_colors_i) # V H W 3 + rendered_depths.append(render_depths_i) # V H W 1 + rendered_alphas.append(render_alphas_i) # V H W 1 + + rendered_colors = torch.stack(rendered_colors, dim=0) # B V H W 3 + rendered_depths = torch.stack(rendered_depths, dim=0) # B V H W 1 + rendered_alphas = torch.stack(rendered_alphas, dim=0) # B V H W 1 + + return rendered_colors, rendered_depths, rendered_alphas + + +class GaussianSplatRenderer(nn.Module): + def __init__( + self, + feature_dim: int = 256, # Output channels of gs_feat_head + sh_degree: int = 0, + enable_prune: bool = True, + voxel_size: float = 0.002, # Default voxel size for prune_gs + enable_conf_filter: bool = False, # Enable confidence filtering + conf_threshold_percent: float = 30.0, # Confidence threshold percentage + max_gaussians: int = 5000000, # Maximum number of Gaussians + ): + super().__init__() + + self.feature_dim = feature_dim + self.sh_degree = sh_degree + self.nums_sh = (sh_degree + 1) ** 2 + self.voxel_size = voxel_size + self.enable_prune = enable_prune + self.enable_conf_filter = enable_conf_filter + self.conf_threshold_percent = conf_threshold_percent + self.max_gaussians = max_gaussians + + # Predict Gaussian parameters from GS features (quaternions/scales/opacities/SH/weights) + splits_and_inits = [ + (4, 1.0, 0.0), # quats + (3, 0.00003, -7.0), # scales + (1, 1.0, -2.0), # opacities + (3 * self.nums_sh, 1.0, 0.0), # residual_sh + (1, 1.0, -2.0), # weights + ] + gaussian_raw_channels = 4 + 3 + 1 + self.nums_sh * 3 + 1 + + self.gs_head = nn.Sequential( + nn.Conv2d( + feature_dim // 2, + feature_dim, + kernel_size=3, + padding=1, + bias=False, + ), + nn.ReLU(True), + nn.Conv2d(feature_dim, gaussian_raw_channels, kernel_size=1), + ) + # Initialize weights and biases of the final layer by segments + final_conv_layer = self.gs_head[-1] + start_channels = 0 + for out_channel, s, b in splits_and_inits: + nn.init.xavier_uniform_( + final_conv_layer.weight[ + start_channels : start_channels + out_channel + ], + s, + ) + nn.init.constant_( + final_conv_layer.bias[ + start_channels : start_channels + out_channel + ], + b, + ) + start_channels += out_channel + + # Rasterizer + self.rasterizer = Rasterizer() + + # ======== Main entry point: Complete GS rendering and fill results back to predictions ======== + def render( + self, + gs_feats: torch.Tensor, # [B, S, 3, H, W] + images: torch.Tensor, # [B, S+V, 3, H, W] + predictions: Dict[ + str, torch.Tensor + ], # From WorldMirror: pose/depth/pts3d etc + views: Dict[str, torch.Tensor], + context_predictions: Dict[str, torch.Tensor], + is_inference: bool = True, + ) -> Dict[str, torch.Tensor]: + """ + Returns predictions with the following fields filled: + - rendered_colors / rendered_depths / (rendered_alphas during training) + - gt_colors / gt_depths / valid_masks + - splats / rendered_extrinsics / rendered_intrinsics + """ + B, _, _, H, W = images.shape + S = context_predictions.get("imgs", images).shape[ + 1 + ] # context view nums + V = images.shape[1] - S # target view nums + + # 1) Predict GS features from tokens, then convert to Gaussian parameters + gs_feats_reshape = rearrange(gs_feats, "b s c h w -> (b s) c h w") + gs_params = self.gs_head(gs_feats_reshape) + gt_colors = images.permute(0, 1, 3, 4, 2) + + # 2) Select rendering cameras + if self.training: + # Using all gt cameras + render_viewmats, render_Ks = self.prepare_cameras(views, S + V) + gt_valid_masks_src = views["valid_mask"][:, :S] # [B, S, H, W] + gt_valid_masks_tgt = views["valid_mask"][:, S:] # [B, V, H, W] + unproject_masks = calculate_unprojected_mask( + views, S + ) # [B, V, H, W] + valid_masks = torch.cat( + [gt_valid_masks_src, (gt_valid_masks_tgt & unproject_masks)], + dim=1, + ) + else: + # Re-predict the camera for novel views and perform translation scale alignment + pred_all_extrinsic, pred_all_intrinsic = self.prepare_cameras( + predictions, S + V + ) + scale_factor = 1.0 + if "camera_poses" in context_predictions: + pred_context_extrinsic, _ = self.prepare_cameras( + context_predictions, S + ) + scale_factor = pred_context_extrinsic[:, :, :3, 3].norm( + dim=-1 + ).mean(dim=1, keepdim=True) / ( + pred_all_extrinsic[:, :S, :3, 3] + .norm(dim=-1) + .mean(dim=1, keepdim=True) + + 1e-6 + ) + + pred_all_extrinsic[..., :3, 3] = pred_all_extrinsic[ + ..., :3, 3 + ] * scale_factor.unsqueeze(-1) + render_viewmats, render_Ks = pred_all_extrinsic, pred_all_intrinsic + valid_masks = views.get( + "valid_mask", + torch.ones(B, S + V, H, W, dtype=bool, device=images.device), + ) + + # 3) Generate splats from gs_params + predictions, and perform voxel merging + if self.training: + splats = self.prepare_splats( + views, + predictions, + images, + gs_params, + S, + position_from="gsdepth+gtcamera", + ) + elif not is_inference: + splats = self.prepare_splats( + views, + predictions, + images, + gs_params, + S, + context_predictions, + position_from="gsdepth+predcamera", + ) + else: + splats = self.prepare_splats( + views, + predictions, + images, + gs_params, + S, + position_from="gsdepth+predcamera", + ) + + # Apply confidence filtering before pruning + if self.enable_conf_filter and "gs_depth_conf" in predictions: + splats = self.apply_confidence_filter( + splats, predictions["gs_depth_conf"] + ) + + if self.enable_prune: + splats = self.prune_gs(splats, voxel_size=self.voxel_size) + + predictions["splats"] = splats + if is_inference: + return predictions + + # 4) Rasterization rendering (training: chunked rendering + novel view valid mask correction; evaluation: view-by-view) + + # Prevent OOM by using chunked rendering + rendered_colors_list, rendered_depths_list, rendered_alphas_list = ( + [], + [], + [], + ) + chunk_size = 4 + for i in range(0, gt_colors.shape[1], chunk_size): + end_idx = min(i + chunk_size, gt_colors.shape[1]) + viewmats_i = render_viewmats[:, i:end_idx] + Ks_i = render_Ks[:, i:end_idx] + + rendered_colors, rendered_depths, rendered_alphas = ( + self.rasterizer.rasterize_batches( + splats["means"], + splats["quats"], + splats["scales"], + splats["opacities"], + splats["sh"] if "sh" in splats else splats["colors"], + viewmats_i.detach(), + Ks_i.detach(), + width=images.shape[-1], + height=images.shape[-2], + sh_degree=( + min(self.sh_degree, 0) if "sh" in splats else None + ), + ) + ) + rendered_colors_list.append(rendered_colors) + rendered_depths_list.append(rendered_depths) + rendered_alphas_list.append(rendered_alphas) + + rendered_colors = torch.cat(rendered_colors_list, dim=1) + rendered_depths = torch.cat(rendered_depths_list, dim=1) + rendered_alphas = torch.cat(rendered_alphas_list, dim=1) + + if self.training and V > 0: + nvs_rendered_mask = rendered_alphas[:, S:, ..., 0].detach() > 0.1 + valid_masks[:, S:] = nvs_rendered_mask & valid_masks[:, S:] + + # 5) return predictions + predictions["rendered_colors"] = rendered_colors + predictions["rendered_depths"] = rendered_depths + predictions["gt_colors"] = gt_colors + predictions["gt_depths"] = views.get("depthmap") + predictions["valid_masks"] = valid_masks.bool() + predictions["rendered_extrinsics"] = render_viewmats + predictions["rendered_intrinsics"] = render_Ks + + return predictions + + def apply_confidence_filter(self, splats, gs_depth_conf): + """Apply confidence filtering to Gaussian splats before pruning. Discard bottom p% + confidence points, keep top (100-p)%. + + Args: + splats: Dictionary containing Gaussian parameters + gs_depth_conf: Confidence tensor [B, S, H, W] + + Returns: + Filtered splats dictionary + """ + if not self.enable_conf_filter or gs_depth_conf is None: + return splats + + device = splats["means"].device + B, N = splats["means"].shape[:2] + + # Flatten confidence: [B, S, H, W] -> [B, N] + conf = gs_depth_conf.flatten(1).to(device) + # Mask invalid/very small values + conf = conf.masked_fill(conf <= 1e-5, float("-inf")) + + # Keep top (100-p)% points, discard bottom p% + if self.conf_threshold_percent > 0: + keep_from_percent = int( + np.ceil(N * (100.0 - self.conf_threshold_percent) / 100.0) + ) + else: + keep_from_percent = N + K = max(1, min(self.max_gaussians, keep_from_percent)) + + # Select top-K indices for each batch (deterministic, no randomness) + topk_idx = torch.topk( + conf, K, dim=1, largest=True, sorted=False + ).indices # [B, K] + + filtered = {} + mask_keys = ["means", "quats", "scales", "opacities", "sh", "weights"] + + for key in splats.keys(): + if key in mask_keys and key in splats: + x = splats[key] + if x.ndim == 2: # [B, N] + filtered[key] = torch.gather(x, 1, topk_idx) + else: + # Expand indices to match tensor dimensions + expand_idx = topk_idx.clone() + for i in range(x.ndim - 2): + expand_idx = expand_idx.unsqueeze(-1) + expand_idx = expand_idx.expand(-1, -1, *x.shape[2:]) + filtered[key] = torch.gather(x, 1, expand_idx) + else: + filtered[key] = splats[key] + + return filtered + + def prune_gs(self, splats, voxel_size=0.002): + """Prune Gaussian splats by merging those in the same voxel. + + Args: + splats: Dictionary containing Gaussian parameters + voxel_size: Size of voxels for spatial grouping + + Returns: + Dictionary with pruned splats + """ + B = splats["means"].shape[0] + merged_splats_list = [] + device = splats["means"].device + + for i in range(B): + # Extract splats for current batch + splats_i = { + k: splats[k][i] + for k in [ + "means", + "quats", + "scales", + "opacities", + "sh", + "weights", + ] + } + + # Compute voxel indices + coords = splats_i["means"] + voxel_indices = (coords / voxel_size).floor().long() + min_indices = voxel_indices.min(dim=0)[0] + voxel_indices = voxel_indices - min_indices + max_dims = voxel_indices.max(dim=0)[0] + 1 + + # Flatten 3D voxel indices to 1D + flat_indices = ( + voxel_indices[:, 0] * max_dims[1] * max_dims[2] + + voxel_indices[:, 1] * max_dims[2] + + voxel_indices[:, 2] + ) + + # Find unique voxels and inverse mapping + unique_voxels, inverse_indices = torch.unique( + flat_indices, return_inverse=True + ) + K = len(unique_voxels) + + # Initialize merged splats + merged = { + "means": torch.zeros((K, 3), device=device), + "quats": torch.zeros((K, 4), device=device), + "scales": torch.zeros((K, 3), device=device), + "opacities": torch.zeros(K, device=device), + "sh": torch.zeros((K, self.nums_sh, 3), device=device), + } + + # Get weights and compute weight sums per voxel + weights = splats_i["weights"] + weight_sums = torch.zeros(K, device=device) + weight_sums.scatter_add_(0, inverse_indices, weights) + weight_sums = torch.clamp(weight_sums, min=1e-8) + + # Merge means (weighted average) + for d in range(3): + merged["means"][:, d].scatter_add_( + 0, inverse_indices, splats_i["means"][:, d] * weights + ) + merged["means"] = merged["means"] / weight_sums.unsqueeze(1) + + # Merge spherical harmonics (weighted average) + for d in range(3): + merged["sh"][:, 0, d].scatter_add_( + 0, inverse_indices, splats_i["sh"][:, 0, d] * weights + ) + merged["sh"] = merged["sh"] / weight_sums.unsqueeze(-1).unsqueeze( + -1 + ) + + # Merge opacities (weighted sum of squares) + merged["opacities"].scatter_add_( + 0, inverse_indices, weights * weights + ) + merged["opacities"] = merged["opacities"] / weight_sums + + # Merge scales (weighted average) + for d in range(3): + merged["scales"][:, d].scatter_add_( + 0, inverse_indices, splats_i["scales"][:, d] * weights + ) + merged["scales"] = merged["scales"] / weight_sums.unsqueeze(1) + + # Merge quaternions (weighted average + normalization) + for d in range(4): + merged["quats"][:, d].scatter_add_( + 0, inverse_indices, splats_i["quats"][:, d] * weights + ) + quat_norms = torch.norm(merged["quats"], dim=1, keepdim=True) + merged["quats"] = merged["quats"] / torch.clamp( + quat_norms, min=1e-8 + ) + + merged_splats_list.append(merged) + + # Reorganize output + output = {} + for key in ["means", "sh", "opacities", "scales", "quats"]: + output[key] = [merged[key] for merged in merged_splats_list] + + return output + + def prepare_splats( + self, + views, + predictions, + images, + gs_params, + context_nums, + context_predictions={}, + position_from="gsdepth+gtcamera", + ): + """Prepare Gaussian splats from model predictions and input data. + + Args: + views: Dictionary containing view data (camera poses, intrinsics, etc.) + predictions: Model predictions including depth, pose_enc, etc. + images: Input images [B, S_all, 3, H, W] + gs_params: Gaussian splatting parameters from model + context_predictions: Optional context predictions for camera poses + position_from: Method to compute 3D positions ("pts3d", "gsdepth+gtcamera", "gsdepth+predcamera) + debug: Whether to use debug mode with ground truth data + + Returns: + splats: Dictionary containing prepared Gaussian splat parameters + """ + B, _, _, H, W = images.shape + S = context_nums + splats = {} + + # Only take parameters from source view branch + gs_params = rearrange(gs_params, "(b s) c h w -> b s h w c", b=B) + splats["gs_feats"] = gs_params.reshape(B, S * H * W, -1) + + # Split Gaussian parameters + quats, scales, opacities, residual_sh, weights = torch.split( + gs_params, [4, 3, 1, self.nums_sh * 3, 1], dim=-1 + ) + + # Apply activation functions to Gaussian parameters + splats["quats"] = act_gs.reg_dense_rotation( + quats.reshape(B, S * H * W, 4) + ) + splats["scales"] = act_gs.reg_dense_scales( + scales.reshape(B, S * H * W, 3) + ).clamp_max(0.3) + splats["opacities"] = act_gs.reg_dense_opacities( + opacities.reshape(B, S * H * W) + ) + residual_sh = act_gs.reg_dense_sh( + residual_sh.reshape(B, S * H * W, self.nums_sh * 3) + ) + + # Handle spherical harmonics (SH) coefficients + new_sh = torch.zeros_like(residual_sh) + new_sh[..., 0, :] = sh_utils.RGB2SH( + images[:, :S].permute(0, 1, 3, 4, 2).reshape(B, S * H * W, 3) + ) + splats["sh"] = new_sh + residual_sh + splats["residual_sh"] = residual_sh + + splats["weights"] = act_gs.reg_dense_weights( + weights.reshape(B, S * H * W) + ) + + # Compute 3D positions based on specified method + if position_from == "pts3d": + pts3d = predictions["pts3d"][:, :S].reshape(B, S * H * W, 3) + splats["means"] = pts3d + + elif position_from == "gsdepth+gtcamera": + depth = predictions["gs_depth"][:, :S].reshape(B * S, H, W) + pose4x4 = views["camera_poses"][:, :S].reshape(B * S, 4, 4) + intrinsic = views["camera_intrs"][:, :S].reshape(B * S, 3, 3) + pts3d, _, _ = depth_to_world_coords_points( + depth, pose4x4, intrinsic + ) + pts3d = pts3d.reshape(B, S * H * W, 3) + splats["means"] = pts3d + + elif position_from == "gsdepth+predcamera": + depth = predictions["gs_depth"][:, :S].reshape(B * S, H, W) + pose4x4 = context_predictions.get( + "camera_poses", predictions["camera_poses"] + )[:, :S].reshape(B * S, 4, 4) + intrinsic = context_predictions.get( + "camera_intrs", predictions["camera_intrs"] + )[:, :S].reshape(B * S, 3, 3) + pts3d, _, _ = depth_to_world_coords_points( + depth, pose4x4.detach(), intrinsic.detach() + ) + pts3d = pts3d.reshape(B, S * H * W, 3) + splats["means"] = pts3d + + else: + raise ValueError(f"Invalid position_from={position_from}") + + return splats + + def prepare_cameras(self, views, nums): + viewmats = views["camera_poses"][:, :nums] + Ks = views["camera_intrs"][:, :nums] + return viewmats, Ks + + +if __name__ == "__main__": + device = "cuda:0" + means = torch.randn((100, 3), device=device) + quats = torch.randn((100, 4), device=device) + scales = torch.rand((100, 3), device=device) * 0.1 + opacities = torch.rand((100,), device=device) + colors = torch.rand((100, 3), device=device) + + viewmats = torch.eye(4, device=device)[None, :, :].repeat(10, 1, 1) + Ks = torch.tensor( + [[300.0, 0.0, 150.0], [0.0, 300.0, 100.0], [0.0, 0.0, 1.0]], + device=device, + )[None, :, :].repeat(10, 1, 1) + width, height = 300, 200 + + rasterizer = Rasterizer() + splats = { + "means": means, + "quats": quats, + "scales": scales, + "opacities": opacities, + "colors": colors, + } + colors, alphas, _ = rasterizer.rasterize_splats( + splats, viewmats, Ks, width, height + ) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/visual_transformer.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/visual_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..850e15d9599241a4e5f341811e5f2dbfcc011c72 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/visual_transformer.py @@ -0,0 +1,571 @@ +import logging +import random +from typing import Tuple, List + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +from reward_function.HunyuanWorldMirror.src.models.layers import ( + PatchEmbed, + PatchEmbed_Mlp, +) +from reward_function.HunyuanWorldMirror.src.models.layers.block import Block +from reward_function.HunyuanWorldMirror.src.models.layers.rope import ( + RotaryPositionEmbedding2D, + PositionGetter, +) +from reward_function.HunyuanWorldMirror.src.models.layers.vision_transformer import ( + vit_small, + vit_base, + vit_large, + vit_giant2, +) + +logger = logging.getLogger(__name__) + +_RESNET_MEAN = [0.485, 0.456, 0.406] +_RESNET_STD = [0.229, 0.224, 0.225] + + +class VisualGeometryTransformer(nn.Module): + """ + The VisualGeometryTransformer applies alternating-attention over input frames, + as described in VGGT: Visual Geometry Grounded Transformer. + + Args: + img_size (int): Image size in pixels. + patch_size (int): Size of each patch for PatchEmbed. + embed_dim (int): Dimension of the token embeddings. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + mlp_ratio (float): Ratio of MLP hidden dim to embedding dim. + num_register_tokens (int): Number of register tokens. + block_fn (nn.Module): The block type used for attention (Block by default). + qkv_bias (bool): Whether to include bias in QKV projections. + proj_bias (bool): Whether to include bias in the output projection. + ffn_bias (bool): Whether to include bias in MLP layers. + patch_embed (str): Type of patch embed. e.g., "conv" or "dinov2_vitl14_reg". + aa_order (list[str]): The order of alternating attention, e.g. ["frame", "global"]. + qk_norm (bool): Whether to apply QK normalization. + rope_freq (int): Base frequency for rotary embedding. -1 to disable. + init_values (float): Init scale for layer scale. + enable_condition (bool): Whether to enable conditioning inputs. + sampling_strategy (str): Sampling strategy for patches. + fixed_patch_embed (bool): Whether to fix patch embedding weights. + condition_strategy (list[str]): Strategy for each conditioning input. + """ + + def __init__( + self, + img_size=518, + patch_size=14, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4.0, + num_register_tokens=4, + block_fn=Block, + qkv_bias=True, + proj_bias=True, + ffn_bias=True, + patch_embed="dinov2_vitl14_reg", + qk_norm=True, + rope_freq=100, + init_values=0.01, + enable_cond=False, + sampling_strategy="uniform", + fixed_patch_embed=False, + condition_strategy=["token", "pow3r", "token"], + intermediate_idxs: List[int] = [4, 11, 17, 23], + ): + super().__init__() + # Store config parameters + self.enable_cond = enable_cond + self.sampling_strategy = sampling_strategy + self.cond_methods = condition_strategy + self.intermediate_idxs = intermediate_idxs + self.depth = depth + self.patch_size = patch_size + + # Initialize patch embedding module + self.patch_embed = self._init_patch_embedding_module( + patch_embed, + img_size, + patch_size, + num_register_tokens, + embed_dim=embed_dim, + is_fixed=fixed_patch_embed, + ) + + # Initialize conditioning embeddings if enabled + if self.enable_cond: + self._init_cond_embeddings( + embed_dim, img_size, patch_size, num_register_tokens + ) + + # Initialize rotary position embedding + self._init_rotary_position_embedding(rope_freq) + + # Initialize transformer blocks + self._init_transformer_blocks( + block_fn, + embed_dim, + num_heads, + mlp_ratio, + qkv_bias, + proj_bias, + ffn_bias, + init_values, + qk_norm, + ) + + # Initialize learnable tokens + self._init_learnable_tokens(embed_dim, num_register_tokens) + + # Calculate patch start index based on conditioning + if self.enable_cond: + self.patch_start_idx = ( + 1 + num_register_tokens + 1 + 1 + ) # camera + register + pose + rays + else: + self.patch_start_idx = 1 + num_register_tokens # camera + register + + # Register normalization constants + for name, value in ( + ("_resnet_mean", _RESNET_MEAN), + ("_resnet_std", _RESNET_STD), + ): + self.register_buffer( + name, + torch.FloatTensor(value).reshape(1, 1, 3, 1, 1), + persistent=False, + ) + + self.use_reentrant = False + + def _init_patch_embedding_module( + self, + patch_embed_type, + img_size, + patch_size, + num_reg_tokens, + interpolate_antialias=True, + interpolate_offset=0.0, + block_chunks=0, + init_values=1.0, + embed_dim=1024, + is_fixed=False, + in_chans=3, + ): + """Create the patch embedding module. + + If 'conv', we use a simple PatchEmbed conv layer. Otherwise, we use a vision transformer. + """ + if "conv" in patch_embed_type: + if "mlp" in patch_embed_type: + patch_embed_module = PatchEmbed_Mlp( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + else: + patch_embed_module = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + else: + vit_models = { + "dinov2_vitl14_reg": vit_large, + "dinov2_vitb14_reg": vit_base, + "dinov2_vits14_reg": vit_small, + "dinov2_vitg2_reg": vit_giant2, + } + + patch_embed_module = vit_models[patch_embed_type]( + img_size=img_size, + patch_size=patch_size, + num_register_tokens=num_reg_tokens, + interpolate_antialias=interpolate_antialias, + interpolate_offset=interpolate_offset, + block_chunks=block_chunks, + init_values=init_values, + ) + + # Disable gradient updates for mask token + if hasattr(patch_embed_module, "mask_token"): + patch_embed_module.mask_token.requires_grad_(False) + + if is_fixed: + for param in patch_embed_module.parameters(): + param.requires_grad_(False) + + return patch_embed_module + + def _init_cond_embeddings( + self, embed_dim, img_size, patch_size, num_reg_tokens + ): + """Initialize conditioning embeddings for camera, depth, and rays.""" + assert self.cond_methods is not None + assert self.cond_methods[0] == "token" + + # Camera pose embedding + if self.cond_methods[0] == "token": + self.pose_embed = nn.Sequential( + nn.Linear(7, embed_dim, bias=True), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=True), + ) + else: + raise NotImplementedError + + # Depth map embedding + if self.cond_methods[1] == "pow3r": + self.depth_embed = self._init_patch_embedding_module( + "conv+mlp", + img_size, + patch_size, + num_reg_tokens, + embed_dim=embed_dim, + in_chans=1, + ) + else: + raise NotImplementedError + + # Ray direction embedding + if self.cond_methods[2] == "token": + self.ray_embed = nn.Sequential( + nn.Linear(4, embed_dim, bias=True), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=True), + ) + else: + raise NotImplementedError + + def _init_rotary_position_embedding(self, rope_freq): + self.rope = ( + RotaryPositionEmbedding2D( + frequency=rope_freq, + ) + if rope_freq > 0 + else None + ) + self.pos_getter = PositionGetter() if self.rope is not None else None + + def _init_transformer_blocks( + self, + block_fn, + embed_dim, + num_heads, + mlp_ratio, + qkv_bias, + proj_bias, + ffn_bias, + init_values, + qk_norm, + ): + self.frame_blocks = nn.ModuleList( + [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + init_values=init_values, + qk_norm=qk_norm, + rope=self.rope, + ) + for _ in range(self.depth) + ] + ) + + self.global_blocks = nn.ModuleList( + [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + init_values=init_values, + qk_norm=qk_norm, + rope=self.rope, + ) + for _ in range(self.depth) + ] + ) + + def _init_learnable_tokens(self, embed_dim, num_reg_tokens): + """Initialize learnable tokens.""" + self.cam_token = nn.Parameter(torch.zeros(1, 2, 1, embed_dim)) + self.reg_token = nn.Parameter( + torch.zeros(1, 2, num_reg_tokens, embed_dim) + ) + nn.init.normal_(self.cam_token, std=1e-6) + nn.init.normal_(self.reg_token, std=1e-6) + + def forward( + self, + images: torch.Tensor, + priors: List | None = None, + cond_flags: List[int] = [0, 0, 0], + ctx_frames: int = None, + ) -> Tuple[List[torch.Tensor], int]: + """ + Args: + images: Input images with shape [B, S, 3, H, W], in range [0, 1] + priors: Optional tuple of (depth, rays, poses) for conditioning + cond_flags: List indicating which conditions to use [pose, depth, rays] + ctx_frames: Number of context frames to use + + Returns: + (list[torch.Tensor], int): List of attention block outputs and patch_start_idx + """ + depth_maps, ray_dirs, poses = ( + priors if priors is not None else (None, None, None) + ) + + # Slice to context frames if specified + if ctx_frames is not None: + for var_name in ["images", "depth_maps", "ray_dirs", "poses"]: + var = locals()[var_name] + if var is not None: + locals()[var_name] = var[:, :ctx_frames].clone() + + # Process image tokens + b, seq_len, ch, h, w = images.shape + if ch != 3: + raise ValueError(f"Expected 3 input channels, got {ch}") + + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + images = (images - self._resnet_mean) / self._resnet_std + images = images.reshape(b * seq_len, ch, h, w) + patch_tokens = self.patch_embed(images) + if isinstance(patch_tokens, dict): + patch_tokens = patch_tokens["x_norm_patchtokens"] + + _, patch_count, embed_dim = patch_tokens.shape + + # Prepare special tokens + cam_tokens = expand_and_flatten_special_tokens( + self.cam_token, b, seq_len + ) + reg_tokens = expand_and_flatten_special_tokens( + self.reg_token, b, seq_len + ) + + # Process all tokens (optional conditioning) + if self.enable_cond: + pose_tokens, depth_tokens, ray_tokens = self._process_conditioning( + depth_maps, + ray_dirs, + poses, + b, + seq_len, + patch_count, + embed_dim, + images, + cond_flags, + ) + # Add condition tokens to patch tokens + patch_tokens = patch_tokens + depth_tokens + all_tokens = torch.cat( + [cam_tokens, reg_tokens, pose_tokens, ray_tokens, patch_tokens], + dim=1, + ) + else: + all_tokens = torch.cat( + [cam_tokens, reg_tokens, patch_tokens], dim=1 + ) + + _, patch_count, embed_dim = all_tokens.shape + + # Position embedding + pos_emb = None + if self.rope is not None: + pos_emb = self.pos_getter( + b * seq_len, + h // self.patch_size, + w // self.patch_size, + device=images.device, + ) + if self.patch_start_idx > 0: + pos_emb = pos_emb + 1 + special_pos = torch.zeros( + b * seq_len, + self.patch_start_idx, + 2, + device=images.device, + dtype=pos_emb.dtype, + ) + pos_emb = torch.cat([special_pos, pos_emb], dim=1) + + # Forward through attention blocks + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + outputs = [] + global_tokens = None + for idx in range(self.depth): + local_tokens = self._process_attention_blocks( + tokens=( + all_tokens if global_tokens is None else global_tokens + ), + b=b, + seq_len=seq_len, + patch_count=patch_count, + embed_dim=embed_dim, + block_idx=idx, + blocks=self.frame_blocks, + block_type="frame", + pos=pos_emb, + ) + global_tokens = self._process_attention_blocks( + tokens=local_tokens, + b=b, + seq_len=seq_len, + patch_count=patch_count, + embed_dim=embed_dim, + block_idx=idx, + blocks=self.global_blocks, + block_type="global", + pos=pos_emb, + ) + + # Combine frame and global intermediates + if idx in self.intermediate_idxs: + combined_out = torch.cat( + [local_tokens, global_tokens], dim=-1 + ) + outputs.append(combined_out) + + return outputs, self.patch_start_idx + + def _process_conditioning( + self, + depth_maps, + ray_dirs, + poses, + b, + seq_len, + patch_count, + embed_dim, + images, + cond_flags, + ): + """Process conditioning inputs.""" + h, w = images.shape[-2:] + + # Process camera pose embedding + use_poses = cond_flags[0] == 1 and poses is not None + if use_poses: + poses = poses.reshape(b * seq_len, -1) + pose_tokens = self.pose_embed(poses).unsqueeze(1) + else: + pose_tokens = torch.zeros( + (b * seq_len, 1, embed_dim), + device=images.device, + dtype=images.dtype, + ) + + # Process depth map embedding + use_depth = cond_flags[1] == 1 and depth_maps is not None + if use_depth: + depth_maps = depth_maps.reshape(b * seq_len, 1, h, w) + depth_tokens = self.depth_embed(depth_maps).reshape( + b * seq_len, patch_count, embed_dim + ) + else: + depth_tokens = torch.zeros( + (b * seq_len, patch_count, embed_dim), + device=images.device, + dtype=images.dtype, + ) + + # Process ray direction embedding + use_rays = cond_flags[2] == 1 and ray_dirs is not None + if use_rays: + ray_dirs = ray_dirs.reshape(b * seq_len, -1) + ray_tokens = self.ray_embed(ray_dirs).unsqueeze(1) + else: + ray_tokens = torch.zeros( + (b * seq_len, 1, embed_dim), + device=images.device, + dtype=images.dtype, + ) + + return pose_tokens, depth_tokens, ray_tokens + + def _process_attention_blocks( + self, + tokens, + b, + seq_len, + patch_count, + embed_dim, + block_idx, + blocks, + block_type, + pos=None, + ): + """Process attention blocks with tokens in shape (B*S, P, C).""" + token_shape = (b, seq_len, patch_count, embed_dim) + if block_type == "frame": # local + target_shape = (b * seq_len, patch_count, embed_dim) + pos_target_shape = ( + (b * seq_len, patch_count, 2) if pos is not None else None + ) + else: # global + target_shape = (b, seq_len * patch_count, embed_dim) + pos_target_shape = ( + (b, seq_len * patch_count, 2) if pos is not None else None + ) + + if tokens.shape != target_shape: + tokens = tokens.reshape(*target_shape) + + if pos is not None and pos.shape != pos_target_shape: + pos = pos.reshape(*pos_target_shape) + + if self.training: + tokens = checkpoint( + blocks[block_idx], + tokens, + pos=pos, + use_reentrant=self.use_reentrant, + ) + else: + tokens = blocks[block_idx](tokens, pos=pos) + + return tokens.reshape(*token_shape) + + +def expand_and_flatten_special_tokens(token_tensor, b, seq_len): + """Processes specialized tokens with shape (1, 2, X, C) for multi-frame processing. Uses first + position for frame 0, second position for remaining frames. + + Args: + token_tensor: Input tensor with shape (1, 2, X, C) + b: Batch size + seq_len: Sequence length + + Returns: + torch.Tensor: Processed tokens with shape (B*S, X, C) + """ + # First frame uses position 0, remaining frames use position 1 + first_frame_tokens = token_tensor[:, 0:1, ...].expand( + b, 1, *token_tensor.shape[2:] + ) + remaining_frame_tokens = token_tensor[:, 1:, ...].expand( + b, seq_len - 1, *token_tensor.shape[2:] + ) + + # Concatenate and flatten + combined_tokens = torch.cat( + [first_frame_tokens, remaining_frame_tokens], dim=1 + ) + return combined_tokens.reshape(b * seq_len, *combined_tokens.shape[2:]) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/worldmirror.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/worldmirror.py new file mode 100644 index 0000000000000000000000000000000000000000..1dc628e32b9f1155fb64fca6cfc8df0aa52dede2 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/models/worldmirror.py @@ -0,0 +1,393 @@ +from typing import Dict, List + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from reward_function.HunyuanWorldMirror.src.models.models.visual_transformer import ( + VisualGeometryTransformer, +) +from reward_function.HunyuanWorldMirror.src.models.heads.camera_head import ( + CameraHead, +) +from reward_function.HunyuanWorldMirror.src.models.heads.dense_head import ( + DPTHead, +) +from reward_function.HunyuanWorldMirror.src.models.models.rasterization import ( + GaussianSplatRenderer, +) +from reward_function.HunyuanWorldMirror.src.models.utils.camera_utils import ( + vector_to_camera_matrices, + extrinsics_to_vector, +) +from reward_function.HunyuanWorldMirror.src.models.utils.priors import ( + normalize_depth, + normalize_poses, +) + +from huggingface_hub import PyTorchModelHubMixin + + +class WorldMirror(nn.Module, PyTorchModelHubMixin): + def __init__( + self, + img_size=518, + patch_size=14, + embed_dim=1024, + gs_dim=256, + enable_cond=True, + enable_cam=True, + enable_pts=True, + enable_depth=True, + enable_norm=True, + enable_gs=True, + patch_embed="dinov2_vitl14_reg", + fixed_patch_embed=False, + sampling_strategy="uniform", + dpt_gradient_checkpoint=False, + condition_strategy=["token", "pow3r", "token"], + ): + + super().__init__() + # Configuration flags + self.img_size = img_size + self.patch_size = patch_size + self.embed_dim = embed_dim + self.gs_dim = gs_dim + self.enable_cam = enable_cam + self.enable_pts = enable_pts + self.enable_depth = enable_depth + self.enable_cond = enable_cond + self.enable_norm = enable_norm + self.enable_gs = enable_gs + self.patch_embed = patch_embed + self.sampling = sampling_strategy + self.dpt_checkpoint = dpt_gradient_checkpoint + self.cond_methods = condition_strategy + self.config = self._store_config() + + # Visual geometry transformer + self.visual_geometry_transformer = VisualGeometryTransformer( + img_size=img_size, + patch_size=patch_size, + embed_dim=embed_dim, + enable_cond=enable_cond, + sampling_strategy=sampling_strategy, + patch_embed=patch_embed, + fixed_patch_embed=fixed_patch_embed, + condition_strategy=condition_strategy, + ) + + # Initialize prediction heads + self._init_heads(embed_dim, patch_size, gs_dim) + + def _store_config(self): + """Save the model configuration.""" + return { + "img_size": self.img_size, + "patch_size": self.patch_size, + "embed_dim": self.embed_dim, + "gs_dim": self.gs_dim, + "enable_cam": self.enable_cam, + "enable_pts": self.enable_pts, + "enable_depth": self.enable_depth, + "enable_norm": self.enable_norm, + "enable_gs": self.enable_gs, + "patch_embed": self.patch_embed, + "sampling_strategy": self.sampling, + "dpt_checkpoint": self.dpt_checkpoint, + "condition_strategy": self.cond_methods, + } + + def _init_heads(self, dim, patch_size, gs_dim): + """Initialize all prediction heads.""" + + # Camera pose prediction head + if self.enable_cam: + self.cam_head = CameraHead(dim_in=2 * dim) + + # 3D point prediction head + if self.enable_pts: + self.pts_head = DPTHead( + dim_in=2 * dim, + output_dim=4, + patch_size=patch_size, + activation="inv_log+expp1", + ) + + # Depth prediction head + if self.enable_depth: + self.depth_head = DPTHead( + dim_in=2 * dim, + output_dim=2, + patch_size=patch_size, + activation="exp+expp1", + ) + + # Surface normal prediction head + if self.enable_norm: + self.norm_head = DPTHead( + dim_in=2 * dim, + output_dim=4, + patch_size=patch_size, + activation="norm+expp1", + ) + + # Gaussian splatting feature head and renderer + if self.enable_gs: + self.gs_head = DPTHead( + dim_in=2 * dim, + output_dim=2, + patch_size=patch_size, + features=gs_dim, + is_gsdpt=True, + activation="exp+expp1", + ) + self.gs_renderer = GaussianSplatRenderer( + sh_degree=0, + enable_prune=True, + voxel_size=0.002, + ) + + def forward( + self, + views: Dict[str, torch.Tensor], + cond_flags: List[int] = [0, 0, 0], + is_inference=True, + ): + """Execute forward pass through the WorldMirror model. + + Args: + views: Input data dictionary + cond_flags: Conditioning flags [depth, rays, camera] + + Returns: + dict: Prediction results dictionary + """ + imgs = views["img"] + + # Enable conditional input during training if enabled, or during inference if any cond_flags are set + use_cond = sum(cond_flags) > 0 + + # Extract priors and process features based on conditional input + if use_cond: + priors = self.extract_priors(views) + token_list, patch_start_idx = self.visual_geometry_transformer( + imgs, priors, cond_flags=cond_flags + ) + else: + token_list, patch_start_idx = self.visual_geometry_transformer(imgs) + + # Execute predictions + with torch.amp.autocast("cuda", dtype=torch.float32): + # Generate all predictions + preds = self._gen_all_preds( + token_list, + imgs, + patch_start_idx, + views, + cond_flags, + is_inference, + ) + + return preds + + def _gen_all_preds( + self, token_list, imgs, patch_start_idx, views, cond_flags, is_inference + ): + """Generate all enabled predictions.""" + preds = {} + + # Camera pose prediction + if self.enable_cam: + cam_seq = self.cam_head(token_list) + cam_params = cam_seq[-1] + preds["camera_params"] = cam_params + c2w_mat, int_mat = self.transform_camera_vector( + cam_params, imgs.shape[-2], imgs.shape[-1] + ) + preds["camera_poses"] = ( + c2w_mat # C2W pose (OpenCV) in world coordinates: [B, S, 4, 4] + ) + preds["camera_intrs"] = ( + int_mat # Camera intrinsic matrix: [B, S, 3, 3] + ) + + # Depth prediction + if self.enable_depth: + depth, depth_conf = self.depth_head( + token_list, + images=imgs, + patch_start_idx=patch_start_idx, + ) + preds["depth"] = depth + preds["depth_conf"] = depth_conf + + # 3D point prediction + if self.enable_pts: + pts, pts_conf = self.pts_head( + token_list, + images=imgs, + patch_start_idx=patch_start_idx, + ) + preds["pts3d"] = pts + preds["pts3d_conf"] = pts_conf + + # Normal prediction + if self.enable_norm: + normals, norm_conf = self.norm_head( + token_list, + images=imgs, + patch_start_idx=patch_start_idx, + ) + preds["normals"] = normals + preds["normals_conf"] = norm_conf + + # 3D Gaussian Splatting + if self.enable_gs: + context_preds = self.prepare_contexts( + views, cond_flags, is_inference + ) + gs_feat, gs_depth, gs_depth_conf = self.gs_head( + context_preds.get("token_list", token_list), + images=context_preds.get("imgs", imgs), + patch_start_idx=patch_start_idx, + ) + + preds["gs_depth"] = gs_depth + preds["gs_depth_conf"] = gs_depth_conf + preds = self.gs_renderer.render( + gs_feats=gs_feat, + images=imgs, + predictions=preds, + views=views, + context_predictions=context_preds, + is_inference=is_inference, + ) + return preds + + def extract_priors(self, views): + """Extract and normalize geometric priors. + + Args: + views: Input view data dictionary. + + Returns: + tuple: (depths, rays, poses) Normalized priors. + """ + h, w = views["img"].shape[-2:] + + # Initialize prior variables + depths = rays = poses = None + + # Extract camera pose + if "camera_poses" in views: + extrinsics = views["camera_poses"][:, :, :3] + extrinsics = normalize_poses(extrinsics) + cam_params = extrinsics_to_vector(extrinsics) + poses = cam_params[:, :, :7] # Shape: [B, S, 7] + + # Extract depth map + if "depthmap" in views: + depth_h, depth_w = views["depthmap"].shape[-2:] + depths = views["depthmap"] + if ( + depth_h != h or depth_w != w + ): # Check if depth dimensions match target resolution + try: + depths = F.interpolate( + depths, + size=(h, w), + mode="bilinear", + align_corners=False, + ) + except: + import pdb + + pdb.set_trace() + depths = normalize_depth(depths) # Shape: [B, S, H, W] + + # Extract ray directions + if "camera_intrs" in views: + intrinsics = views["camera_intrs"][:, :, :3, :3] + fx, fy = intrinsics[:, :, 0, 0] / w, intrinsics[:, :, 1, 1] / h + cx, cy = intrinsics[:, :, 0, 2] / w, intrinsics[:, :, 1, 2] / h + rays = torch.stack([fx, fy, cx, cy], dim=-1) # Shape: [B, S, 4] + + return (depths, rays, poses) + + def transform_camera_vector(self, camera_params, h, w): + ext_mat, int_mat = vector_to_camera_matrices( + camera_params, image_hw=(h, w) + ) + # Create homogeneous transformation matrix + homo_row = torch.tensor([0, 0, 0, 1], device=ext_mat.device).view( + 1, 1, 1, 4 + ) + homo_row = homo_row.repeat(ext_mat.shape[0], ext_mat.shape[1], 1, 1) + w2c_mat = torch.cat([ext_mat, homo_row], dim=2) + c2w_mat = torch.linalg.inv(w2c_mat) + return c2w_mat, int_mat + + def prepare_contexts(self, views, cond_flags, is_inference): + # Generate context views predictions + context_preds = {} + # only for training or evaluation + if is_inference: + return context_preds + + assert self.enable_cam and self.enable_gs + if "is_target" not in views: + context_nums = views["img"].shape[1] + else: + context_nums = (views["is_target"][0] == False).sum().item() + context_imgs = views["img"][:, :context_nums] + + use_cond = sum(cond_flags) > 0 + + # Extract context priors and process features based on context views + with torch.amp.autocast("cuda", enabled=True): + if use_cond: + priors = self.extract_priors(views) + context_priors = ( + prior[:, :context_nums] if prior is not None else None + for prior in priors + ) + context_token_list, _ = self.visual_geometry_transformer( + context_imgs, context_priors, cond_flags=cond_flags + ) + else: + context_token_list, _ = self.visual_geometry_transformer( + context_imgs + ) + + # Execute predictions + with torch.amp.autocast("cuda", enabled=False): + # Context camera pose prediction + context_cam_seq = self.cam_head(context_token_list) + context_cam_params = context_cam_seq[-1] + context_c2w_mat, context_int_mat = self.transform_camera_vector( + context_cam_params, + context_imgs.shape[-2], + context_imgs.shape[-1], + ) + context_preds["camera_poses"] = ( + context_c2w_mat # C2W pose (OpenCV) in world coordinates: [B, S, 4, 4] + ) + context_preds["camera_intrs"] = ( + context_int_mat # Camera intrinsic matrix: [B, S, 3, 3] + ) + context_preds["token_list"] = context_token_list + context_preds["imgs"] = context_imgs + + return context_preds + + +if __name__ == "__main__": + device = "cuda" + model = WorldMirror().to(device).eval() + x = torch.rand(1, 1, 3, 518, 518).to(device) + out = model({"img": x}) + import pdb + + pdb.set_trace() diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/act_gs.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/act_gs.py new file mode 100644 index 0000000000000000000000000000000000000000..7b3c0c82ee18380c1153426897a8e8cd616b4ae2 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/act_gs.py @@ -0,0 +1,27 @@ +import torch +from einops import rearrange + + +def reg_dense_offsets(xyz, shift=6.0): + d = xyz.norm(dim=-1, keepdim=True) + return xyz / d.clamp(min=1e-8) * (torch.exp(d - shift) - torch.exp(-shift)) + + +def reg_dense_scales(scales): + return scales.exp() + + +def reg_dense_rotation(rotations, eps=1e-8): + return rotations / (rotations.norm(dim=-1, keepdim=True) + eps) + + +def reg_dense_sh(sh): + return rearrange(sh, "... (d_sh xyz) -> ... d_sh xyz", xyz=3) + + +def reg_dense_opacities(opacities): + return opacities.sigmoid() + + +def reg_dense_weights(weights): + return weights.sigmoid() diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/camera_utils.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/camera_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f36b2f9606051d61f6c7f1ab777c1d44403fc2b7 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/camera_utils.py @@ -0,0 +1,91 @@ +import torch +from .rotation import quat_to_rotmat, rotmat_to_quat + + +def camera_params_to_vector(ext, intr, image_hw=None): + """Convert camera matrices to a compact vector.""" + # ext: (..., 3, 4): Camera-to-world extrinsic [R|t] + # intr: (..., 3, 3): Intrinsics + # image_hw: (h, w) + R = ext[..., :3, :3] # Rotation part + t = ext[..., :3, 3] # Translation part + q = rotmat_to_quat(R) # Quaternion (wxyz) + h, w = image_hw + fov_v = 2.0 * torch.atan(h * 0.5 / intr[..., 1, 1]) # Vertical FOV + fov_u = 2.0 * torch.atan(w * 0.5 / intr[..., 0, 0]) # Horizontal FOV + vec = torch.stack( + [ + t[..., 0], + t[..., 1], + t[..., 2], + q[..., 0], + q[..., 1], + q[..., 2], + q[..., 3], + fov_v, + fov_u, + ], + dim=-1, + ).float() + return vec + + +def extrinsics_to_vector(ext): + """Convert extrinsics to [t, q] vector.""" + # ext: (..., 3, 4) + R = ext[..., :3, :3] + t = ext[..., :3, 3] + q = rotmat_to_quat(R) + vec = torch.stack( + [ + t[..., 0], + t[..., 1], + t[..., 2], + q[..., 0], + q[..., 1], + q[..., 2], + q[..., 3], + ], + dim=-1, + ).float() + return vec + + +def vector_to_extrinsics(cam_vec): + """Convert [t, q] vector to extrinsic matrix.""" + # cam_vec: (..., 7) + q = cam_vec[..., 3:7] + t = cam_vec[..., :3] + R = quat_to_rotmat(q) + ext = torch.cat([R, t.unsqueeze(-1)], dim=-1) + return ext + + +def vector_to_camera_matrices(cam_vec, image_hw=None, build_intr=True): + """Reconstruct extrinsic and intrinsic matrix from vector.""" + # cam_vec: (..., 9) + intr = None + # Decompose vector + t = cam_vec[..., 0:3] + q = cam_vec[..., 3:7] + fov_v = cam_vec[..., 7] + fov_u = cam_vec[..., 8] + + # Build extrinsic: [R|t] + R = quat_to_rotmat(q) + ext = torch.cat([R, t.unsqueeze(-1)], dim=-1) + + # Build intrinsic if needed + if build_intr: + h, w = image_hw + fy = h * 0.5 / torch.tan(fov_v * 0.5) + fx = w * 0.5 / torch.tan(fov_u * 0.5) + shape = cam_vec.shape[:-1] + (3, 3) + intr = torch.zeros(shape, device=cam_vec.device, dtype=cam_vec.dtype) + intr[..., 0, 0] = fx + intr[..., 1, 1] = fy + intr[..., 0, 2] = w * 0.5 + intr[..., 1, 2] = h * 0.5 + intr[..., 2, 2] = 1.0 + + return ext, intr diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/frustum.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/frustum.py new file mode 100644 index 0000000000000000000000000000000000000000..6de200672f67393e0407cd1541815ef1dc54a490 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/frustum.py @@ -0,0 +1,222 @@ +import einops +import torch + + +# Calculate the loss mask for the target views in the batch +@torch.no_grad() +def calculate_unprojected_mask(views, context_nums): + """Calcuate the loss mask for the target views in the batch.""" + target_depth = views["depthmap"][:, context_nums:] + target_intrinsics = views["camera_intrs"][:, context_nums:] + target_c2w = views["camera_poses"][:, context_nums:] + context_depth = views["depthmap"][:, :context_nums] + context_intrinsics = views["camera_intrs"][:, :context_nums] + context_c2w = views["camera_poses"][:, :context_nums] + + target_intrinsics = target_intrinsics[..., :3, :3] + context_intrinsics = context_intrinsics[..., :3, :3] + + mask = calculate_in_frustum_mask( + target_depth, + target_intrinsics, + target_c2w, + context_depth, + context_intrinsics, + context_c2w, + ) + return mask + + +@torch.no_grad() +def calculate_in_frustum_mask( + depth_1, intrinsics_1, c2w_1, depth_2, intrinsics_2, c2w_2 +): + """A function that takes in the depth, intrinsics and c2w matrices of two sets of views, and + then works out which of the pixels in the first set of views has a direct corresponding pixel in + any of views in the second set. + + Args: + depth_1: (b, v1, h, w) + intrinsics_1: (b, v1, 3, 3) + c2w_1: (b, v1, 4, 4) + depth_2: (b, v2, h, w) + intrinsics_2: (b, v2, 3, 3) + c2w_2: (b, v2, 4, 4) + + Returns: + torch.Tensor: valid mask with shape (b, v1, v2, h, w). + """ + + _, v1, h, w = depth_1.shape + _, v2, _, _ = depth_2.shape + + # Unproject the depth to get the 3D points in world space + points_3d = unproject_depth( + depth_1[..., None], intrinsics_1, c2w_1 + ) # (b, v1, h, w, 3) + + # Project the 3D points into the pixel space of all the second views simultaneously + camera_points = world_space_to_camera_space( + points_3d, c2w_2 + ) # (b, v1, v2, h, w, 3) + points_2d = camera_space_to_pixel_space( + camera_points, intrinsics_2 + ) # (b, v1, v2, h, w, 2) + + # Calculate the depth of each point + rendered_depth = camera_points[..., 2] # (b, v1, v2, h, w) + + # We use three conditions to determine if a point should be masked + + # Condition 1: Check if the points are in the frustum of any of the v2 views + in_frustum_mask = ( + (points_2d[..., 0] > 0) + & (points_2d[..., 0] < w) + & (points_2d[..., 1] > 0) + & (points_2d[..., 1] < h) + ) # (b, v1, v2, h, w) + in_frustum_mask = in_frustum_mask.any(dim=-3) # (b, v1, h, w) + + # Condition 2: Check if the points have non-zero (i.e. valid) depth in the input view + non_zero_depth = depth_1 > 1e-6 + + # Condition 3: Check if the points have matching depth to any of the v2 + # views torch.nn.functional.grid_sample expects the input coordinates to + # be normalized to the range [-1, 1], so we normalize first + points_2d[..., 0] /= w + points_2d[..., 1] /= h + points_2d = points_2d * 2 - 1 + matching_depth = torch.ones_like(rendered_depth, dtype=torch.bool) + for b in range(depth_1.shape[0]): + for i in range(v1): + for j in range(v2): + depth = einops.rearrange(depth_2[b, j], "h w -> 1 1 h w") + coords = einops.rearrange( + points_2d[b, i, j], "h w c -> 1 h w c" + ) + sampled_depths = torch.nn.functional.grid_sample( + depth, coords, align_corners=False + )[0, 0] + matching_depth[b, i, j] = torch.isclose( + rendered_depth[b, i, j], sampled_depths, atol=1e-1 + ) + + matching_depth = matching_depth.any(dim=-3) # (..., v1, h, w) + + mask = in_frustum_mask & non_zero_depth & matching_depth + return mask + + +# --- Projections --- +def homogenize_points(points): + """Append a '1' along the final dimension of the tensor (i.e. convert xyz->xyz1)""" + return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + +def normalize_homogenous_points(points): + """Normalize the point vectors.""" + return points / points[..., -1:] + + +def pixel_space_to_camera_space(pixel_space_points, depth, intrinsics): + """Convert pixel space points to camera space points. + + Args: + pixel_space_points (torch.Tensor): Pixel space points with shape (h, w, 2) + depth (torch.Tensor): Depth map with shape (b, v, h, w, 1) + intrinsics (torch.Tensor): Camera intrinsics with shape (b, v, 3, 3) + + Returns: + torch.Tensor: Camera space points with shape (b, v, h, w, 3). + """ + pixel_space_points = homogenize_points(pixel_space_points) + camera_space_points = torch.einsum( + "b v i j , h w j -> b v h w i", intrinsics.inverse(), pixel_space_points + ) + camera_space_points = camera_space_points * depth + return camera_space_points + + +def camera_space_to_world_space(camera_space_points, c2w): + """Convert camera space points to world space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + camera_space_points = homogenize_points(camera_space_points) + world_space_points = torch.einsum( + "b v i j , b v h w j -> b v h w i", c2w, camera_space_points + ) + return world_space_points[..., :3] + + +def camera_space_to_pixel_space(camera_space_points, intrinsics): + """Convert camera space points to pixel space points. + + Args: + camera_space_points (torch.Tensor): Camera space points with shape (b, v1, v2, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 3, 3) + + Returns: + torch.Tensor: World space points with shape (b, v1, v2, h, w, 2). + """ + camera_space_points = normalize_homogenous_points(camera_space_points) + pixel_space_points = torch.einsum( + "b u i j , b v u h w j -> b v u h w i", intrinsics, camera_space_points + ) + return pixel_space_points[..., :2] + + +def world_space_to_camera_space(world_space_points, c2w): + """Convert world space points to pixel space points. + + Args: + world_space_points (torch.Tensor): World space points with shape (b, v1, h, w, 3) + c2w (torch.Tensor): Camera to world extrinsics matrix with shape (b, v2, 4, 4) + + Returns: + torch.Tensor: Camera space points with shape (b, v1, v2, h, w, 3). + """ + world_space_points = homogenize_points(world_space_points) + camera_space_points = torch.einsum( + "b u i j , b v h w j -> b v u h w i", c2w.inverse(), world_space_points + ) + return camera_space_points[..., :3] + + +def unproject_depth(depth, intrinsics, c2w): + """Turn the depth map into a 3D point cloud in world space. + + Args: + depth: (b, v, h, w, 1) + intrinsics: (b, v, 3, 3) + c2w: (b, v, 4, 4) + + Returns: + torch.Tensor: World space points with shape (b, v, h, w, 3). + """ + + # Compute indices of pixels + h, w = depth.shape[-3], depth.shape[-2] + x_grid, y_grid = torch.meshgrid( + torch.arange(w, device=depth.device, dtype=torch.float32), + torch.arange(h, device=depth.device, dtype=torch.float32), + indexing="xy", + ) # (h, w), (h, w) + + # Compute coordinates of pixels in camera space + pixel_space_points = torch.stack((x_grid, y_grid), dim=-1) # (..., h, w, 2) + camera_points = pixel_space_to_camera_space( + pixel_space_points, depth, intrinsics + ) # (..., h, w, 3) + + # Convert points to world space + world_points = camera_space_to_world_space( + camera_points, c2w + ) # (..., h, w, 3) + + return world_points diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/geometry.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..be1f4f19754254c7ed6f3e587eeb9adf924513f9 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/geometry.py @@ -0,0 +1,145 @@ +import torch +import numpy as np + + +def depth_to_camera_coords(depthmap, camera_intrinsics): + """Convert depth map to 3D camera coordinates. + + Args: + depthmap (BxHxW tensor): Batch of depth maps + camera_intrinsics (Bx3x3 tensor): Camera intrinsics matrix for each camera + + Returns: + X_cam (BxHxWx3 tensor): 3D points in camera coordinates + valid_mask (BxHxW tensor): Mask indicating valid depth pixels + """ + B, H, W = depthmap.shape + device = depthmap.device + dtype = depthmap.dtype + + # Ensure intrinsics are float + camera_intrinsics = camera_intrinsics.float() + + # Extract focal lengths and principal points + fx = camera_intrinsics[:, 0, 0] # (B,) + fy = camera_intrinsics[:, 1, 1] # (B,) + cx = camera_intrinsics[:, 0, 2] # (B,) + cy = camera_intrinsics[:, 1, 2] # (B,) + + # Generate pixel grid + v_grid, u_grid = torch.meshgrid( + torch.arange(H, dtype=dtype, device=device), + torch.arange(W, dtype=dtype, device=device), + indexing="ij", + ) + + # Reshape for broadcasting: (1, H, W) + u_grid = u_grid.unsqueeze(0) + v_grid = v_grid.unsqueeze(0) + + # Compute 3D camera coordinates + # X = (u - cx) * Z / fx + # Y = (v - cy) * Z / fy + # Z = depth + z_cam = depthmap # (B, H, W) + x_cam = (u_grid - cx.view(B, 1, 1)) * z_cam / fx.view(B, 1, 1) + y_cam = (v_grid - cy.view(B, 1, 1)) * z_cam / fy.view(B, 1, 1) + + # Stack to form (B, H, W, 3) + X_cam = torch.stack([x_cam, y_cam, z_cam], dim=-1) + + # Valid depth mask + valid_mask = depthmap > 0.0 + + return X_cam, valid_mask + + +def depth_to_world_coords_points( + depth_map: torch.Tensor, + extrinsic: torch.Tensor, + intrinsic: torch.Tensor, + eps=1e-8, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert a batch of depth maps to world coordinates. + + Args: + depth_map (torch.Tensor): (B, H, W) Depth map + extrinsic (torch.Tensor): (B, 4, 4) Camera extrinsic matrix (camera-to-world transformation) + intrinsic (torch.Tensor): (B, 3, 3) Camera intrinsic matrix + + Returns: + world_coords_points (torch.Tensor): (B, H, W, 3) World coordinates + camera_points (torch.Tensor): (B, H, W, 3) Camera coordinates + point_mask (torch.Tensor): (B, H, W) Valid depth mask + """ + if depth_map is None: + return None, None, None + + # Valid depth mask (B, H, W) + point_mask = depth_map > eps + + # Convert depth map to camera coordinates (B, H, W, 3) + camera_points, _ = depth_to_camera_coords(depth_map, intrinsic) + + # Apply extrinsic matrix (camera -> world) + R_cam_to_world = extrinsic[:, :3, :3] # (B, 3, 3) + t_cam_to_world = extrinsic[:, :3, 3] # (B, 3) + + # Transform (B, H, W, 3) x (B, 3, 3)^T + (B, 3) -> (B, H, W, 3) + world_coords_points = ( + torch.einsum("bhwi,bji->bhwj", camera_points, R_cam_to_world) + + t_cam_to_world[:, None, None, :] + ) + + return world_coords_points, camera_points, point_mask + + +def closed_form_inverse_se3(se3: torch.Tensor) -> torch.Tensor: + """Efficiently invert batched SE(3) matrices of shape (B, 4, 4). + + Args: + se3 (torch.Tensor): (B, 4, 4) Transformation matrices + + Returns: + out (torch.Tensor): (B, 4, 4) Inverse transformation matrices + """ + assert se3.ndim == 3 and se3.shape[1:] == ( + 4, + 4, + ), f"se3 must be (B, 4, 4), got {se3.shape}" + R = se3[:, :3, :3] # (B, 3, 3) + t = se3[:, :3, 3] # (B, 3) + Rt = R.transpose(1, 2) # (B, 3, 3) + t_inv = -torch.bmm(Rt, t.unsqueeze(-1)).squeeze(-1) # (B, 3) + out = se3.new_zeros(se3.shape) + out[:, :3, :3] = Rt + out[:, :3, 3] = t_inv + out[:, 3, 3] = 1.0 + return out + + +def create_pixel_coordinate_grid(num_frames, height, width): + """ + Creates a grid of pixel coordinates and frame indices for all frames. + Returns: + tuple: A tuple containing: + - points_xyf (numpy.ndarray): Array of shape (num_frames, height, width, 3) + with x, y coordinates and frame indices + """ + # Create coordinate grids for a single frame + y_grid, x_grid = np.indices((height, width), dtype=np.float32) + x_grid = x_grid[np.newaxis, :, :] + y_grid = y_grid[np.newaxis, :, :] + + # Broadcast to all frames + x_coords = np.broadcast_to(x_grid, (num_frames, height, width)) + y_coords = np.broadcast_to(y_grid, (num_frames, height, width)) + + # Create frame indices and broadcast + f_idx = np.arange(num_frames, dtype=np.float32)[:, np.newaxis, np.newaxis] + f_coords = np.broadcast_to(f_idx, (num_frames, height, width)) + + # Stack coordinates and frame indices + points_xyf = np.stack((x_coords, y_coords, f_coords), axis=-1) + + return points_xyf diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/grid.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/grid.py new file mode 100644 index 0000000000000000000000000000000000000000..382d0a149b4151c15f4482df68d5079b2e57cadd --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/grid.py @@ -0,0 +1,102 @@ +import torch + + +def position_grid_to_embed( + pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100 +) -> torch.Tensor: + """Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC) + + Args: + pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates + embed_dim: Output channel dimension for embeddings + omega_0: Base frequency for sinusoidal encoding + + Returns: + Tensor of shape (H, W, embed_dim) with positional embeddings + """ + H, W, grid_dim = pos_grid.shape + assert grid_dim == 2 + assert embed_dim % 2 == 0 + + device = pos_grid.device + pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2) + + # Generate frequency bands + omega = torch.arange( + embed_dim // 4, + dtype=torch.float32 if device.type == "mps" else torch.double, + device=device, + ) + omega /= embed_dim / 4.0 + omega = 1.0 / omega_0**omega # (D/4,) + + # Process x and y coordinates separately + pos_x = pos_flat[:, 0].reshape(-1) # (H*W,) + pos_y = pos_flat[:, 1].reshape(-1) # (H*W,) + + # Compute outer products + out_x = torch.einsum("m,d->md", pos_x, omega) # (H*W, D/4) + out_y = torch.einsum("m,d->md", pos_y, omega) # (H*W, D/4) + + # Apply sin and cos + emb_x = torch.cat([torch.sin(out_x), torch.cos(out_x)], dim=1) # (H*W, D/2) + emb_y = torch.cat([torch.sin(out_y), torch.cos(out_y)], dim=1) # (H*W, D/2) + + # Combine x and y embeddings + emb = torch.cat([emb_x, emb_y], dim=-1) # (H*W, D) + + return emb.float().view(H, W, embed_dim) # [H, W, D] + + +# Inspired by https://github.com/microsoft/moge +def create_uv_grid( + width: int, + height: int, + aspect_ratio: float = None, + dtype: torch.dtype = None, + device: torch.device = None, +) -> torch.Tensor: + """Create a normalized UV grid of shape (width, height, 2). + + The grid spans horizontally and vertically according to an aspect ratio, + ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right + corner is at (x_span, y_span), normalized by the diagonal of the plane. + + Args: + width (int): Number of points horizontally. + height (int): Number of points vertically. + aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height. + dtype (torch.dtype, optional): Data type of the resulting tensor. + device (torch.device, optional): Device on which the tensor is created. + + Returns: + torch.Tensor: A (width, height, 2) tensor of UV coordinates. + """ + # Derive aspect ratio if not explicitly provided + if aspect_ratio is None: + aspect_ratio = float(width) / float(height) + + # Compute normalized spans for X and Y + diag_factor = (aspect_ratio**2 + 1.0) ** 0.5 + span_x = aspect_ratio / diag_factor + span_y = 1.0 / diag_factor + + # Establish the linspace boundaries + left_x = -span_x * (width - 1) / width + right_x = span_x * (width - 1) / width + top_y = -span_y * (height - 1) / height + bottom_y = span_y * (height - 1) / height + + # Generate 1D coordinates + x_coords = torch.linspace( + left_x, right_x, steps=width, dtype=dtype, device=device + ) + y_coords = torch.linspace( + top_y, bottom_y, steps=height, dtype=dtype, device=device + ) + + # Create 2D meshgrid (width x height) and stack into UV + uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy") + uv_grid = torch.stack((uu, vv), dim=-1) + + return uv_grid diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/priors.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/priors.py new file mode 100644 index 0000000000000000000000000000000000000000..2085e4787dce84501601b34797e6c3f1ec40fa0a --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/priors.py @@ -0,0 +1,171 @@ +import torch + + +def normalize_poses(extrinsics, padding=0.1, return_stats=False): + """Normalize camera positions to unit cube, processing each batch separately. + + Args: + extrinsics: Camera extrinsic matrices with shape (B, S, 3, 4) + padding: Boundary space within [0,1] range to prevent values near boundaries + return_stats: Whether to return normalization statistics + + Returns: + normalized_extrinsics: Normalized extrinsic matrices + (optional) stats: Dictionary containing scale and translation information + """ + B, S, _, _ = extrinsics.shape + device = extrinsics.device + + # Check input validity and handle NaN/Inf values + for i in range(B): + if torch.isnan(extrinsics[i]).any() or torch.isinf(extrinsics[i]).any(): + print(f"Warning: dataset sample has NaN/Inf in extrinsics") + extrinsics[i] = torch.nan_to_num( + extrinsics[i], nan=0.0, posinf=1e6, neginf=-1e6 + ) + + normalized_extrinsics = extrinsics.clone() + + # Store normalization parameters if needed + if return_stats: + stats = { + "scale_factors": torch.zeros(B, device=device), + "translation_vectors": torch.zeros(B, 3, device=device), + } + + for b in range(B): + # Extract camera positions for this batch + positions = extrinsics[b, :, :3, 3] # (S, 3) + + # Filter valid positions to ignore outliers + valid_mask = torch.isfinite(positions).all(dim=1) # (S,) + + if valid_mask.sum() == 0: + # No valid positions, use default values + print(f"Warning: Batch {b} has no valid camera positions") + normalized_extrinsics[b, :, :3, 3] = 0.5 # Place at center + if return_stats: + stats["scale_factors"][b] = 1.0 + stats["translation_vectors"][b] = 0.0 + continue + + valid_positions = positions[valid_mask] + + # Calculate bounds using percentiles for robustness + if valid_positions.shape[0] > 10: + # Use 5% and 95% percentiles instead of min/max + min_pos = torch.quantile(valid_positions, 0.05, dim=0) + max_pos = torch.quantile(valid_positions, 0.95, dim=0) + else: + # Too few samples, use min/max + min_pos = torch.min(valid_positions, dim=0)[0] + max_pos = torch.max(valid_positions, dim=0)[0] + + # Calculate scale factor considering all dimensions + pos_range = max_pos - min_pos + + # Add small epsilon to prevent dimension collapse + eps = torch.maximum( + torch.tensor(1e-6, device=device), torch.abs(max_pos) * 1e-6 + ) + pos_range = torch.maximum(pos_range, eps) + + # Use maximum range as scale factor for uniform scaling + scale_factor = torch.max(pos_range) + scale_factor = torch.clamp(scale_factor, min=1e-6, max=1e6) + + # Calculate center point for centering + center = (min_pos + max_pos) / 2.0 + + # Normalize: center first, then scale with padding + actual_scale = scale_factor / (1 - 2 * padding) + normalized_positions = (positions - center) / actual_scale + 0.5 + + # Ensure all values are within valid range + normalized_positions = torch.clamp(normalized_positions, 0.0, 1.0) + + # Handle invalid positions by setting them to scene center + invalid_mask = ~torch.isfinite(positions).all(dim=1) + if invalid_mask.any(): + normalized_positions[invalid_mask] = 0.5 + + normalized_extrinsics[b, :, :3, 3] = normalized_positions + + if return_stats: + stats["scale_factors"][b] = actual_scale + stats["translation_vectors"][b] = center + + # Final validation + assert torch.isfinite( + normalized_extrinsics + ).all(), "Output contains non-finite values" + + if return_stats: + return normalized_extrinsics, stats + return normalized_extrinsics + + +def normalize_depth(depth, eps=1e-6, min_percentile=1, max_percentile=99): + """Normalize depth values to [0, 1] range using percentile-based scaling. + + Args: + depth: Input depth tensor with shape (B, S, H, W) + eps: Small epsilon value to prevent division by zero + min_percentile: Lower percentile for robust min calculation (default: 1) + max_percentile: Upper percentile for robust max calculation (default: 99) + + Returns: + normalized_depth: Depth tensor normalized to [0, 1] range with same shape (B, S, H, W) + """ + B, S, H, W = depth.shape + depth = depth.flatten(0, 1) # [B*S, H, W] + + # Handle invalid values + depth = torch.nan_to_num(depth, nan=0.0, posinf=1e6, neginf=0.0) + + normalized_list = [] + for i in range(depth.shape[0]): + depth_img = depth[i] # [H, W] + depth_flat = depth_img.flatten() + + # Filter out zero values if needed + non_zero_mask = depth_flat > 0 + if non_zero_mask.sum() > 0: + values_to_use = depth_flat[non_zero_mask] + else: + values_to_use = depth_flat + + # Only calculate percentiles when there are enough values + if ( + values_to_use.numel() > 100 + ): # Ensure enough samples for percentile calculation + # Calculate min and max percentiles + depth_min = torch.quantile(values_to_use, min_percentile / 100.0) + depth_max = torch.quantile(values_to_use, max_percentile / 100.0) + else: + # If too few samples, use min/max values + depth_min = values_to_use.min() + depth_max = values_to_use.max() + + # Handle case where max equals min + if depth_max == depth_min: + depth_max = depth_min + 1.0 + + # Use relative epsilon + scale = torch.abs(depth_max - depth_min) + eps_val = max(eps, scale.item() * eps) + + # Perform normalization + depth_norm_img = (depth_img - depth_min) / ( + depth_max - depth_min + eps_val + ) + + # Ensure output is within [0,1] range + depth_norm_img = torch.clamp(depth_norm_img, 0.0, 1.0) + + normalized_list.append(depth_norm_img) + + # Recombine all normalized images + depth_norm = torch.stack(normalized_list) + + return depth_norm.reshape(B, S, H, W) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/rotation.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..3e52781347729481b96f93b530f691945dace6ae --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/rotation.py @@ -0,0 +1,140 @@ +# Modified from PyTorch3D, https://github.com/facebookresearch/pytorch3d + +import torch +import numpy as np +import torch.nn.functional as F + + +def quat_to_rotmat(quaternions: torch.Tensor) -> torch.Tensor: + """ + Quaternion Order: XYZW or say ijkr, scalar-last + + Convert rotations given as quaternions to rotation matrices. + Args: + quaternions: quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + i, j, k, r = torch.unbind(quaternions, -1) + # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def rotmat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack( + [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1 + ), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack( + [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1 + ), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack( + [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1 + ), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack( + [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1 + ), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : + ].reshape(batch_dim + (4,)) + + # Convert from rijk to ijkr + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0.""" + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/sh_utils.py b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/sh_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7a5b6c22d357f2d9860e374ce44c777dd27ecbb8 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/models/utils/sh_utils.py @@ -0,0 +1,128 @@ +# Copyright 2021 The PlenOctree Authors. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +C0 = 0.28209479177387814 +C1 = 0.4886025119029199 +C2 = [ + 1.0925484305920792, + -1.0925484305920792, + 0.31539156525252005, + -1.0925484305920792, + 0.5462742152960396, +] +C3 = [ + -0.5900435899266435, + 2.890611442640554, + -0.4570457994644658, + 0.3731763325901154, + -0.4570457994644658, + 1.445305721320277, + -0.5900435899266435, +] +C4 = [ + 2.5033429417967046, + -1.7701307697799304, + 0.9461746957575601, + -0.6690465435572892, + 0.10578554691520431, + -0.6690465435572892, + 0.47308734787878004, + -1.7701307697799304, + 0.6258357354491761, +] + + +def eval_sh(deg, sh, dirs): + """Evaluate spherical harmonics at unit directions using hardcoded SH polynomials. + + Works with torch/np/jnp. + ... Can be 0 or more batch dimensions. + Args: + deg: int SH deg. Currently, 0-3 supported + sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] + dirs: jnp.ndarray unit directions [..., 3] + Returns: + [..., C] + """ + assert deg <= 4 and deg >= 0 + coeff = (deg + 1) ** 2 + assert sh.shape[-1] >= coeff + + result = C0 * sh[..., 0] + if deg > 0: + x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] + result = ( + result + - C1 * y * sh[..., 1] + + C1 * z * sh[..., 2] + - C1 * x * sh[..., 3] + ) + + if deg > 1: + xx, yy, zz = x * x, y * y, z * z + xy, yz, xz = x * y, y * z, x * z + result = ( + result + + C2[0] * xy * sh[..., 4] + + C2[1] * yz * sh[..., 5] + + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + + C2[3] * xz * sh[..., 7] + + C2[4] * (xx - yy) * sh[..., 8] + ) + + if deg > 2: + result = ( + result + + C3[0] * y * (3 * xx - yy) * sh[..., 9] + + C3[1] * xy * z * sh[..., 10] + + C3[2] * y * (4 * zz - xx - yy) * sh[..., 11] + + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + + C3[5] * z * (xx - yy) * sh[..., 14] + + C3[6] * x * (xx - 3 * yy) * sh[..., 15] + ) + + if deg > 3: + result = ( + result + + C4[0] * xy * (xx - yy) * sh[..., 16] + + C4[1] * yz * (3 * xx - yy) * sh[..., 17] + + C4[2] * xy * (7 * zz - 1) * sh[..., 18] + + C4[3] * yz * (7 * zz - 3) * sh[..., 19] + + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + + C4[5] * xz * (7 * zz - 3) * sh[..., 21] + + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + + C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + + C4[8] + * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) + * sh[..., 24] + ) + return result + + +def RGB2SH(rgb): + return (rgb - 0.5) / C0 + + +def SH2RGB(sh): + return sh * C0 + 0.5 diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/__init__.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/build_pycolmap_recon.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/build_pycolmap_recon.py new file mode 100644 index 0000000000000000000000000000000000000000..d2385099bdb2054d57d6a27a59894b812ef10081 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/build_pycolmap_recon.py @@ -0,0 +1,139 @@ +import numpy as np +import pycolmap + + +def _create_camera_params( + frame_idx, cam_matrices, model_type, distortion_coeffs=None +): + """Build camera parameter array for different model types.""" + if model_type == "PINHOLE": + return np.array( + [ + cam_matrices[frame_idx][0, 0], + cam_matrices[frame_idx][1, 1], + cam_matrices[frame_idx][0, 2], + cam_matrices[frame_idx][1, 2], + ] + ) + elif model_type == "SIMPLE_PINHOLE": + focal_avg = ( + cam_matrices[frame_idx][0, 0] + cam_matrices[frame_idx][1, 1] + ) / 2 + return np.array( + [ + focal_avg, + cam_matrices[frame_idx][0, 2], + cam_matrices[frame_idx][1, 2], + ] + ) + elif model_type == "SIMPLE_RADIAL": + raise NotImplementedError("SIMPLE_RADIAL model not supported") + else: + raise ValueError(f"Unsupported camera model: {model_type}") + + +def _setup_camera_object( + frame_idx, cam_matrices, img_dims, model_type, use_shared +): + """Create and configure camera object.""" + if use_shared and frame_idx > 0: + return None + + params = _create_camera_params(frame_idx, cam_matrices, model_type) + return pycolmap.Camera( + model=model_type, + width=img_dims[0], + height=img_dims[1], + params=params, + camera_id=frame_idx + 1, + ) + + +def _process_frame_points(scene_points, point_coords, frame_idx): + """Extract and process 2D points belonging to specific frame.""" + frame_mask = point_coords[:, 2].astype(np.int32) == frame_idx + valid_indices = np.nonzero(frame_mask)[0] + + point2d_list = [] + for idx, batch_idx in enumerate(valid_indices): + point3d_id = batch_idx + 1 + xy_coords = point_coords[batch_idx][:2] + point2d_list.append(pycolmap.Point2D(xy_coords, point3d_id)) + + # Update track information + track = scene_points.points3D[point3d_id].track + track.add_element(frame_idx + 1, idx) + + return point2d_list + + +def build_pycolmap_reconstruction( + points, + pixel_coords, + point_colors, + poses, + intrinsics, + image_size, + shared_camera_model=False, + camera_model="SIMPLE_PINHOLE", +): + """Convert numpy arrays to pycolmap reconstruction format. + + Creates 3D scene structure without track optimization. Suitable for initialization of neural + rendering methods. + """ + num_frames = len(poses) + num_points = len(points) + + scene = pycolmap.Reconstruction() + + # Add 3D points to scene + for pt_idx in range(num_points): + scene.add_point3D( + points[pt_idx], pycolmap.Track(), point_colors[pt_idx] + ) + + current_camera = None + + # Process each frame + for frame_idx in range(num_frames): + # Setup camera if needed + if current_camera is None or not shared_camera_model: + current_camera = _setup_camera_object( + frame_idx, + intrinsics, + image_size, + camera_model, + shared_camera_model, + ) + scene.add_camera(current_camera) + + # Create pose transformation + rotation_matrix = poses[frame_idx][:3, :3] + translation_vec = poses[frame_idx][:3, 3] + world_to_cam = pycolmap.Rigid3d( + pycolmap.Rotation3d(rotation_matrix), translation_vec + ) + + # Create image object + frame_image = pycolmap.Image( + id=frame_idx + 1, + name=f"frame_{frame_idx + 1}", + camera_id=current_camera.camera_id, + cam_from_world=world_to_cam, + ) + + # Process 2D points for this frame + frame_points = _process_frame_points(scene, pixel_coords, frame_idx) + + # Set image points and registration status + try: + frame_image.points2D = pycolmap.ListPoint2D(frame_points) + frame_image.registered = True + except: + print(f"Warning: Frame {frame_idx + 1} has no valid points") + frame_image.registered = False + + scene.add_image(frame_image) + + return scene diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/color_map.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/color_map.py new file mode 100644 index 0000000000000000000000000000000000000000..0a942140b1055a9c443ca33afa365baf44efe8f5 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/color_map.py @@ -0,0 +1,52 @@ +# References: mvsplat + +import torch +from colorspacious import cspace_convert +from einops import rearrange +from jaxtyping import Float +from matplotlib import cm +from torch import Tensor + + +def apply_color_map( + x: torch.Tensor, + color_map: str = "inferno", +) -> torch.Tensor: + cmap = cm.get_cmap(color_map) + + # Convert to NumPy so that Matplotlib color maps can be used. + mapped = cmap(x.detach().clip(min=0, max=1).cpu().numpy())[..., :3] + + # Convert back to the original format. + return torch.tensor(mapped, device=x.device, dtype=torch.float32) + + +def apply_color_map_to_image( + image: torch.Tensor, # Float[Tensor, "*batch height width"], + color_map: str = "inferno", +): # -> Float[Tensor, "*batch 3 height with"] + image = apply_color_map(image, color_map) + return rearrange(image, "... h w c -> ... c h w") + + +def apply_color_map_2d( + x, # Float[Tensor, "*#batch"], + y, # Float[Tensor, "*#batch"], +): # -> Float[Tensor, "*batch 3"] + red = cspace_convert((189, 0, 0), "sRGB255", "CIELab") + blue = cspace_convert((0, 45, 255), "sRGB255", "CIELab") + white = cspace_convert((255, 255, 255), "sRGB255", "CIELab") + x_np = x.detach().clip(min=0, max=1).cpu().numpy()[..., None] + y_np = y.detach().clip(min=0, max=1).cpu().numpy()[..., None] + + # Interpolate between red and blue on the x axis. + interpolated = x_np * red + (1 - x_np) * blue + + # Interpolate between color and white on the y axis. + interpolated = y_np * interpolated + (1 - y_np) * white + + # Convert to RGB. + rgb = cspace_convert(interpolated, "CIELab", "sRGB1") + return torch.tensor(rgb, device=x.device, dtype=torch.float32).clip( + min=0, max=1 + ) diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/cropping.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/cropping.py new file mode 100644 index 0000000000000000000000000000000000000000..74055201a3f8edfa46f578de2b6630bd30fa8e22 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/cropping.py @@ -0,0 +1,530 @@ +"""Utility functions for cropping and resizing data while maintaining proper cameras. + +References: DUSt3R +""" + +import cv2 +import numpy as np +import PIL.Image + +try: + lanczos = PIL.Image.Resampling.LANCZOS + bicubic = PIL.Image.Resampling.BICUBIC +except AttributeError: + lanczos = PIL.Image.LANCZOS + bicubic = PIL.Image.BICUBIC + +from reward_function.HunyuanWorldMirror.src.utils.geometry import ( + colmap_to_opencv_intrinsics, + opencv_to_colmap_intrinsics, +) + + +class ImageList: + """Convenience class to apply the same operation to a whole set of images. + + This class wraps a list of PIL.Image objects and provides methods to perform operations on all + images simultaneously. + """ + + def __init__(self, images): + if not isinstance(images, (tuple, list, set)): + images = [images] + self.images = [] + for image in images: + if not isinstance(image, PIL.Image.Image): + image = PIL.Image.fromarray(image) + self.images.append(image) + + def __len__(self): + """Return the number of images in the list.""" + return len(self.images) + + def to_pil(self): + """Convert ImageList back to PIL Image(s). + + Returns: + PIL.Image.Image or tuple: Single PIL Image if list contains one image, + or tuple of PIL Images if multiple images + """ + return tuple(self.images) if len(self.images) > 1 else self.images[0] + + @property + def size(self): + """Get the size of images in the list. + + Returns: + tuple: (width, height) of the images + + Raises: + AssertionError: If images have different sizes + """ + sizes = [im.size for im in self.images] + assert all( + sizes[0] == s for s in sizes + ), "All images must have the same size" + return sizes[0] + + def resize(self, *args, **kwargs): + """Resize all images with the same parameters. + + Args: + *args, **kwargs: Arguments passed to PIL.Image.resize() + + Returns: + ImageList: New ImageList containing resized images + """ + return ImageList(self._dispatch("resize", *args, **kwargs)) + + def crop(self, *args, **kwargs): + """Crop all images with the same parameters. + + Args: + *args, **kwargs: Arguments passed to PIL.Image.crop() + + Returns: + ImageList: New ImageList containing cropped images + """ + return ImageList(self._dispatch("crop", *args, **kwargs)) + + def _dispatch(self, func, *args, **kwargs): + """Apply a PIL.Image method to all images in the list. + + Args: + func (str): Name of the PIL.Image method to call + *args, **kwargs: Arguments to pass to the method + + Returns: + list: List of results from applying the method to each image + """ + return [getattr(im, func)(*args, **kwargs) for im in self.images] + + +def rescale_image_and_other_optional_info( + image, + output_resolution, + depthmap=None, + camera_intrinsics=None, + force=True, + additional_quantities_to_be_resized_with_nearest=None, +): + """Rescale the image and depthmap to the output resolution. If the image is larger than the + output resolution, it is rescaled with lanczos interpolation. If force is false and the image is + smaller than the output resolution, it is not rescaled. If force is true and the image is + smaller than the output resolution, it is rescaled with bicubic interpolation. Depth and other + quantities are rescaled with nearest interpolation. + + Args: + image (PIL.Image.Image or np.ndarray): The input image to be rescaled. + output_resolution (tuple): The desired output resolution as a tuple (width, height). + depthmap (np.ndarray, optional): The depth map associated with the image. Defaults to None. + camera_intrinsics (np.ndarray, optional): The camera intrinsics matrix. Defaults to None. + force (bool, optional): If True, force rescaling even if the image is smaller than the output resolution. Defaults to True. + additional_quantities_to_be_resized_with_nearest (list of np.ndarray, optional): Additional quantities to be rescaled using nearest interpolation. Defaults to None. + + Returns: + tuple: A tuple containing: + - The rescaled image (PIL.Image.Image) + - The rescaled depthmap (numpy.ndarray or None) + - The updated camera intrinsics (numpy.ndarray or None) + - The list of rescaled additional quantities (list of numpy.ndarray or None) + """ + image = ImageList(image) + input_resolution = np.array(image.size) # (W, H) + output_resolution = np.array(output_resolution) + if depthmap is not None: + assert tuple(depthmap.shape[:2]) == image.size[::-1] + if additional_quantities_to_be_resized_with_nearest is not None: + assert all( + tuple(additional_quantity.shape[:2]) == image.size[::-1] + for additional_quantity in additional_quantities_to_be_resized_with_nearest + ) + + # Define output resolution + assert output_resolution.shape == (2,) + scale_final = max(output_resolution / image.size) + 1e-8 + if ( + scale_final >= 1 and not force + ): # image is already smaller than what is asked + output = ( + image.to_pil(), + depthmap, + camera_intrinsics, + additional_quantities_to_be_resized_with_nearest, + ) + return output + output_resolution = np.floor(input_resolution * scale_final).astype(int) + + # First rescale the image so that it contains the crop + image = image.resize( + tuple(output_resolution), + resample=lanczos if scale_final < 1 else bicubic, + ) + if depthmap is not None: + depthmap = cv2.resize( + depthmap, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + if additional_quantities_to_be_resized_with_nearest is not None: + resized_additional_quantities = [] + for quantity in additional_quantities_to_be_resized_with_nearest: + resized_additional_quantities.append( + cv2.resize( + quantity, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + ) + additional_quantities_to_be_resized_with_nearest = ( + resized_additional_quantities + ) + + # No offset here; simple rescaling + if camera_intrinsics is not None: + camera_intrinsics = camera_matrix_of_crop( + camera_intrinsics, + input_resolution, + output_resolution, + scaling=scale_final, + ) + + # Return + return ( + image.to_pil(), + depthmap, + camera_intrinsics, + additional_quantities_to_be_resized_with_nearest, + ) + + +def camera_matrix_of_crop( + input_camera_matrix, + input_resolution, + output_resolution, + scaling=1, + offset_factor=0.5, + offset=None, +): + """Calculate the camera matrix for a cropped image. + + Args: + input_camera_matrix (numpy.ndarray): Original camera intrinsics matrix + input_resolution (tuple or numpy.ndarray): Original image resolution as (width, height) + output_resolution (tuple or numpy.ndarray): Target image resolution as (width, height) + scaling (float, optional): Scaling factor for the image. Defaults to 1. + offset_factor (float, optional): Factor to determine crop offset. Defaults to 0.5 (centered). + offset (tuple or numpy.ndarray, optional): Explicit offset to use. If None, calculated from offset_factor. + + Returns: + numpy.ndarray: Updated camera matrix for the cropped image + """ + # Margins to offset the origin + margins = np.asarray(input_resolution) * scaling - output_resolution + assert np.all(margins >= 0.0) + if offset is None: + offset = offset_factor * margins + + # Generate new camera parameters + output_camera_matrix_colmap = opencv_to_colmap_intrinsics( + input_camera_matrix + ) + output_camera_matrix_colmap[:2, :] *= scaling + output_camera_matrix_colmap[:2, 2] -= offset + output_camera_matrix = colmap_to_opencv_intrinsics( + output_camera_matrix_colmap + ) + + return output_camera_matrix + + +def crop_image_and_other_optional_info( + image, + crop_bbox, + depthmap=None, + camera_intrinsics=None, + additional_quantities=None, +): + """Return a crop of the input view and associated data. + + Args: + image (PIL.Image.Image or numpy.ndarray): The input image to be cropped + crop_bbox (tuple): Crop bounding box as (left, top, right, bottom) + depthmap (numpy.ndarray, optional): Depth map associated with the image + camera_intrinsics (numpy.ndarray, optional): Camera intrinsics matrix + additional_quantities (list of numpy.ndarray, optional): Additional data arrays to crop + + Returns: + tuple: A tuple containing: + - The cropped image + - The cropped depth map (if provided or None) + - Updated camera intrinsics (if provided or None) + - List of cropped additional quantities (if provided or None) + """ + image = ImageList(image) + left, top, right, bottom = crop_bbox + + image = image.crop((left, top, right, bottom)) + if depthmap is not None: + depthmap = depthmap[top:bottom, left:right] + if additional_quantities is not None: + additional_quantities = [ + quantity[top:bottom, left:right] + for quantity in additional_quantities + ] + + if camera_intrinsics is not None: + camera_intrinsics = camera_intrinsics.copy() + camera_intrinsics[0, 2] -= left + camera_intrinsics[1, 2] -= top + + return (image.to_pil(), depthmap, camera_intrinsics, additional_quantities) + + +def crop_image_depthmap(image, depthmap, camera_intrinsics, crop_bbox): + """Return a crop of the input view.""" + image = ImageList(image) + l, t, r, b = crop_bbox + + image = image.crop((l, t, r, b)) + if depthmap is not None: + depthmap = depthmap[t:b, l:r] + + camera_intrinsics = camera_intrinsics.copy() + camera_intrinsics[0, 2] -= l + camera_intrinsics[1, 2] -= t + + return image.to_pil(), depthmap, camera_intrinsics + + +def crop_image_normalmap( + image, normalmap, camera_intrinsics, valid_mask, crop_bbox +): + """Return a crop of the input view.""" + image = ImageList(image) + l, t, r, b = crop_bbox + + image = image.crop((l, t, r, b)) + normalmap = normalmap[t:b, l:r] + valid_mask = valid_mask[t:b, l:r] + + camera_intrinsics = camera_intrinsics.copy() + camera_intrinsics[0, 2] -= l + camera_intrinsics[1, 2] -= t + + return image.to_pil(), normalmap, camera_intrinsics, valid_mask + + +def rescale_image_depthmap( + image, depthmap, camera_intrinsics, output_resolution, force=True +): + """Jointly rescale a (image, depthmap) so that (out_width, out_height) >= output_res.""" + image = ImageList(image) + input_resolution = np.array(image.size) # (W,H) + output_resolution = np.array(output_resolution) + if depthmap is not None: + # can also use this with masks instead of depthmaps + assert tuple(depthmap.shape[:2]) == image.size[::-1] + + # define output resolution + assert output_resolution.shape == (2,) + scale_final = max(output_resolution / image.size) + 1e-8 + if ( + scale_final >= 1 and not force + ): # image is already smaller than what is asked + return (image.to_pil(), depthmap, camera_intrinsics) + output_resolution = np.floor(input_resolution * scale_final).astype(int) + + # first rescale the image so that it contains the crop + image = image.resize( + output_resolution, resample=lanczos if scale_final < 1 else bicubic + ) + if depthmap is not None: + depthmap = cv2.resize( + depthmap, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + + # no offset here; simple rescaling + camera_intrinsics = camera_matrix_of_crop( + camera_intrinsics, + input_resolution, + output_resolution, + scaling=scale_final, + ) + + return image.to_pil(), depthmap, camera_intrinsics + + +def rescale_image_normalmap( + image, + normalmap, + camera_intrinsics, + normal_mask, + output_resolution, + force=True, +): + """Jointly rescale a (image, normalmap, normal_mask) so that (out_width, out_height) >= + output_res.""" + image = ImageList(image) + input_resolution = np.array(image.size) # (W,H) + output_resolution = np.array(output_resolution) + if normalmap is not None: + assert tuple(normalmap.shape[:2]) == image.size[::-1] + if normal_mask is not None: + assert tuple(normal_mask.shape[:2]) == image.size[::-1] + + # define output resolution + assert output_resolution.shape == (2,) + scale_final = max(output_resolution / image.size) + 1e-8 + if ( + scale_final >= 1 and not force + ): # image is already smaller than what is asked + return (image.to_pil(), normalmap, camera_intrinsics, normal_mask) + output_resolution = np.floor(input_resolution * scale_final).astype(int) + + # first rescale the image so that it contains the crop + image = image.resize( + output_resolution, resample=lanczos if scale_final < 1 else bicubic + ) + if normalmap is not None: + normalmap = cv2.resize( + normalmap, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + if normal_mask is not None: + # Convert bool to uint8 for cv2.resize, then back to bool + normal_mask_uint8 = normal_mask.astype(np.uint8) + normal_mask_uint8 = cv2.resize( + normal_mask_uint8, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + normal_mask = normal_mask_uint8.astype(bool) + + # no offset here; simple rescaling + camera_intrinsics = camera_matrix_of_crop( + camera_intrinsics, + input_resolution, + output_resolution, + scaling=scale_final, + ) + + return image.to_pil(), normalmap, camera_intrinsics, normal_mask + + +def bbox_from_intrinsics_in_out( + input_camera_matrix, output_camera_matrix, output_resolution +): + """Calculate the bounding box for cropping based on input and output camera intrinsics. + + Args: + input_camera_matrix (numpy.ndarray): Original camera intrinsics matrix + output_camera_matrix (numpy.ndarray): Target camera intrinsics matrix + output_resolution (tuple): Target resolution as (width, height) + + Returns: + tuple: Crop bounding box as (left, top, right, bottom) + """ + out_width, out_height = output_resolution + left, top = np.int32( + np.round(input_camera_matrix[:2, 2] - output_camera_matrix[:2, 2]) + ) + crop_bbox = (left, top, left + out_width, top + out_height) + return crop_bbox + + +def crop_resize_if_necessary( + image, + resolution, + depthmap=None, + intrinsics=None, + additional_quantities=None, +): + """First downsample image using LANCZOS and then crop if necessary to achieve target resolution. + + This function performs high-quality downsampling followed by cropping to achieve the + desired output resolution while maintaining proper camera intrinsics. + + Args: + image (PIL.Image.Image or numpy.ndarray): The input image to be processed + resolution (tuple): Target resolution as (width, height) + depthmap (numpy.ndarray, optional): Depth map associated with the image + intrinsics (numpy.ndarray, optional): Camera intrinsics matrix + additional_quantities (list of numpy.ndarray, optional): Additional data arrays to process + + Returns: + tuple: A tuple containing the processed image and any provided additional data + (depthmap, intrinsics, additional_quantities) that have been similarly processed + """ + # Convert image to PIL.Image.Image if necessary + if not isinstance(image, PIL.Image.Image): + image = PIL.Image.fromarray(image) + + # Get width and height of image + original_width, original_height = image.size + + # High-quality Lanczos down-scaling + target_rescale_resolution = np.array(resolution) + image, depthmap, intrinsics, additional_quantities = ( + rescale_image_and_other_optional_info( + image=image, + output_resolution=target_rescale_resolution, + depthmap=depthmap, + camera_intrinsics=intrinsics, + additional_quantities_to_be_resized_with_nearest=additional_quantities, + ) + ) + + # Actual cropping (if necessary) + if intrinsics is not None: + new_intrinsics = camera_matrix_of_crop( + input_camera_matrix=intrinsics, + input_resolution=image.size, + output_resolution=resolution, + offset_factor=0.5, + ) + crop_bbox = bbox_from_intrinsics_in_out( + input_camera_matrix=intrinsics, + output_camera_matrix=new_intrinsics, + output_resolution=resolution, + ) + else: + # Create a centered crop if no intrinsics are available + w, h = image.size + target_w, target_h = resolution + left = (w - target_w) // 2 + top = (h - target_h) // 2 + crop_bbox = (left, top, left + target_w, top + target_h) + + image, depthmap, new_intrinsics, additional_quantities = ( + crop_image_and_other_optional_info( + image=image, + crop_bbox=crop_bbox, + depthmap=depthmap, + camera_intrinsics=intrinsics, + additional_quantities=additional_quantities, + ) + ) + + # Return the output + output = (image,) + if depthmap is not None: + output += (depthmap,) + if new_intrinsics is not None: + output += (new_intrinsics,) + if additional_quantities is not None: + output += (additional_quantities,) + return output diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/geometry.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..ca5df234ed68811eb609b293650c15748b3d6d3a --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/geometry.py @@ -0,0 +1,539 @@ +"""Utilities for geometry operations. + +References: DUSt3R, MoGe +""" + +from numbers import Number +from typing import Tuple, Union + +import numpy as np +from reward_function.HunyuanWorldMirror.src.utils.warnings import no_warnings + + +def colmap_to_opencv_intrinsics(K): + """Modify camera intrinsics to follow a different convention. + + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] -= 0.5 + K[1, 2] -= 0.5 + + return K + + +def opencv_to_colmap_intrinsics(K): + """Modify camera intrinsics to follow a different convention. + + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] += 0.5 + K[1, 2] += 0.5 + + return K + + +def angle_diff_vec3_numpy(v1: np.ndarray, v2: np.ndarray, eps: float = 1e-12): + """Compute angle difference between 3D vectors using NumPy. + + Args: + v1 (np.ndarray): First vector of shape (..., 3) + v2 (np.ndarray): Second vector of shape (..., 3) + eps (float, optional): Small epsilon value for numerical stability. Defaults to 1e-12. + + Returns: + np.ndarray: Angle differences in radians + """ + return np.arctan2( + np.linalg.norm(np.cross(v1, v2, axis=-1), axis=-1) + eps, + (v1 * v2).sum(axis=-1), + ) + + +@no_warnings(category=RuntimeWarning) +def points_to_normals( + point: np.ndarray, mask: np.ndarray = None, edge_threshold: float = None +) -> np.ndarray: + """Calculate normal map from point map. Value range is [-1, 1]. + + Args: + point (np.ndarray): shape (height, width, 3), point map + mask (optional, np.ndarray): shape (height, width), dtype=bool. Mask of valid depth pixels. Defaults to None. + edge_threshold (optional, float): threshold for the angle (in degrees) between the normal and the view direction. Defaults to None. + + Returns: + normal (np.ndarray): shape (height, width, 3), normal map. + """ + height, width = point.shape[-3:-1] + has_mask = mask is not None + + if mask is None: + mask = np.ones_like(point[..., 0], dtype=bool) + mask_pad = np.zeros((height + 2, width + 2), dtype=bool) + mask_pad[1:-1, 1:-1] = mask + mask = mask_pad + + pts = np.zeros((height + 2, width + 2, 3), dtype=point.dtype) + pts[1:-1, 1:-1, :] = point + up = pts[:-2, 1:-1, :] - pts[1:-1, 1:-1, :] + left = pts[1:-1, :-2, :] - pts[1:-1, 1:-1, :] + down = pts[2:, 1:-1, :] - pts[1:-1, 1:-1, :] + right = pts[1:-1, 2:, :] - pts[1:-1, 1:-1, :] + normal = np.stack( + [ + np.cross(up, left, axis=-1), + np.cross(left, down, axis=-1), + np.cross(down, right, axis=-1), + np.cross(right, up, axis=-1), + ] + ) + normal = normal / (np.linalg.norm(normal, axis=-1, keepdims=True) + 1e-12) + + valid = ( + np.stack( + [ + mask[:-2, 1:-1] & mask[1:-1, :-2], + mask[1:-1, :-2] & mask[2:, 1:-1], + mask[2:, 1:-1] & mask[1:-1, 2:], + mask[1:-1, 2:] & mask[:-2, 1:-1], + ] + ) + & mask[None, 1:-1, 1:-1] + ) + if edge_threshold is not None: + view_angle = angle_diff_vec3_numpy(pts[None, 1:-1, 1:-1, :], normal) + view_angle = np.minimum(view_angle, np.pi - view_angle) + valid = valid & (view_angle < np.deg2rad(edge_threshold)) + + normal = (normal * valid[..., None]).sum(axis=0) + normal = normal / (np.linalg.norm(normal, axis=-1, keepdims=True) + 1e-12) + + if has_mask: + normal_mask = valid.any(axis=0) + normal = np.where(normal_mask[..., None], normal, 0) + return normal, normal_mask + else: + return normal + + +def sliding_window_1d( + x: np.ndarray, window_size: int, stride: int, axis: int = -1 +): + """Create a sliding window view of the input array along a specified axis. + + This function creates a memory-efficient view of the input array with sliding windows + of the specified size and stride. The window dimension is appended to the end of the + output array's shape. This is useful for operations like convolution, pooling, or + any analysis that requires examining local neighborhoods in the data. + + Args: + x (np.ndarray): Input array with shape (..., axis_size, ...) + window_size (int): Size of the sliding window + stride (int): Stride of the sliding window (step size between consecutive windows) + axis (int, optional): Axis to perform sliding window over. Defaults to -1 (last axis) + + Returns: + np.ndarray: View of the input array with shape (..., n_windows, ..., window_size), + where n_windows = (axis_size - window_size + 1) // stride + + Raises: + AssertionError: If window_size is larger than the size of the specified axis + + Example: + >>> x = np.array([1, 2, 3, 4, 5, 6]) + >>> sliding_window_1d(x, window_size=3, stride=2) + array([[1, 2, 3], + [3, 4, 5]]) + """ + assert ( + x.shape[axis] >= window_size + ), f"kernel_size ({window_size}) is larger than axis_size ({x.shape[axis]})" + axis = axis % x.ndim + shape = ( + *x.shape[:axis], + (x.shape[axis] - window_size + 1) // stride, + *x.shape[axis + 1 :], + window_size, + ) + strides = ( + *x.strides[:axis], + stride * x.strides[axis], + *x.strides[axis + 1 :], + x.strides[axis], + ) + x_sliding = np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides) + return x_sliding + + +def sliding_window_nd( + x: np.ndarray, + window_size: Tuple[int, ...], + stride: Tuple[int, ...], + axis: Tuple[int, ...], +) -> np.ndarray: + """Create sliding windows along multiple dimensions of the input array. + + This function applies sliding_window_1d sequentially along multiple axes to create + N-dimensional sliding windows. This is useful for operations that need to examine + local neighborhoods in multiple dimensions simultaneously. + + Args: + x (np.ndarray): Input array + window_size (Tuple[int, ...]): Size of the sliding window for each axis + stride (Tuple[int, ...]): Stride of the sliding window for each axis + axis (Tuple[int, ...]): Axes to perform sliding window over + + Returns: + np.ndarray: Array with sliding windows along the specified dimensions. + The window dimensions are appended to the end of the shape. + + Note: + The length of window_size, stride, and axis tuples must be equal. + + Example: + >>> x = np.random.rand(10, 10) + >>> windows = sliding_window_nd(x, window_size=(3, 3), stride=(2, 2), axis=(-2, -1)) + >>> # Creates 3x3 sliding windows with stride 2 in both dimensions + """ + axis = [axis[i] % x.ndim for i in range(len(axis))] + for i in range(len(axis)): + x = sliding_window_1d(x, window_size[i], stride[i], axis[i]) + return x + + +def sliding_window_2d( + x: np.ndarray, + window_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + axis: Tuple[int, int] = (-2, -1), +) -> np.ndarray: + """Create 2D sliding windows over the input array. + + Convenience function for creating 2D sliding windows, commonly used for image + processing operations like convolution, pooling, or patch extraction. + + Args: + x (np.ndarray): Input array + window_size (Union[int, Tuple[int, int]]): Size of the 2D sliding window. + If int, same size is used for both dimensions. + stride (Union[int, Tuple[int, int]]): Stride of the 2D sliding window. + If int, same stride is used for both dimensions. + axis (Tuple[int, int], optional): Two axes to perform sliding window over. + Defaults to (-2, -1) (last two dimensions). + + Returns: + np.ndarray: Array with 2D sliding windows. The window dimensions (height, width) + are appended to the end of the shape. + + Example: + >>> image = np.random.rand(100, 100) + >>> patches = sliding_window_2d(image, window_size=8, stride=4) + >>> # Creates 8x8 patches with stride 4 from the image + """ + if isinstance(window_size, int): + window_size = (window_size, window_size) + if isinstance(stride, int): + stride = (stride, stride) + return sliding_window_nd(x, window_size, stride, axis) + + +def max_pool_1d( + x: np.ndarray, + kernel_size: int, + stride: int, + padding: int = 0, + axis: int = -1, +): + """Perform 1D max pooling on the input array. + + Max pooling reduces the dimensionality of the input by taking the maximum value + within each sliding window. This is commonly used in neural networks and signal + processing for downsampling and feature extraction. + + Args: + x (np.ndarray): Input array + kernel_size (int): Size of the pooling kernel + stride (int): Stride of the pooling operation + padding (int, optional): Amount of padding to add on both sides. Defaults to 0. + axis (int, optional): Axis to perform max pooling over. Defaults to -1. + + Returns: + np.ndarray: Max pooled array with reduced size along the specified axis + + Note: + - For floating point arrays, padding is done with np.nan values + - For integer arrays, padding is done with the minimum value of the dtype + - np.nanmax is used to handle NaN values in the computation + + Example: + >>> x = np.array([1, 3, 2, 4, 5, 1, 2]) + >>> max_pool_1d(x, kernel_size=3, stride=2) + array([3, 5, 2]) + """ + axis = axis % x.ndim + if padding > 0: + fill_value = np.nan if x.dtype.kind == "f" else np.iinfo(x.dtype).min + padding_arr = np.full( + (*x.shape[:axis], padding, *x.shape[axis + 1 :]), + fill_value=fill_value, + dtype=x.dtype, + ) + x = np.concatenate([padding_arr, x, padding_arr], axis=axis) + a_sliding = sliding_window_1d(x, kernel_size, stride, axis) + max_pool = np.nanmax(a_sliding, axis=-1) + return max_pool + + +def max_pool_nd( + x: np.ndarray, + kernel_size: Tuple[int, ...], + stride: Tuple[int, ...], + padding: Tuple[int, ...], + axis: Tuple[int, ...], +) -> np.ndarray: + """Perform N-dimensional max pooling on the input array. + + This function applies max_pool_1d sequentially along multiple axes to perform + multi-dimensional max pooling. This is useful for downsampling multi-dimensional + data while preserving the most important features. + + Args: + x (np.ndarray): Input array + kernel_size (Tuple[int, ...]): Size of the pooling kernel for each axis + stride (Tuple[int, ...]): Stride of the pooling operation for each axis + padding (Tuple[int, ...]): Amount of padding for each axis + axis (Tuple[int, ...]): Axes to perform max pooling over + + Returns: + np.ndarray: Max pooled array with reduced size along the specified axes + + Note: + The length of kernel_size, stride, padding, and axis tuples must be equal. + Max pooling is applied sequentially along each axis in the order specified. + + Example: + >>> x = np.random.rand(10, 10, 10) + >>> pooled = max_pool_nd(x, kernel_size=(2, 2, 2), stride=(2, 2, 2), + ... padding=(0, 0, 0), axis=(-3, -2, -1)) + >>> # Reduces each dimension by half with 2x2x2 max pooling + """ + for i in range(len(axis)): + x = max_pool_1d(x, kernel_size[i], stride[i], padding[i], axis[i]) + return x + + +def max_pool_2d( + x: np.ndarray, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + padding: Union[int, Tuple[int, int]], + axis: Tuple[int, int] = (-2, -1), +): + """Perform 2D max pooling on the input array. + + Convenience function for 2D max pooling, commonly used in computer vision + and image processing for downsampling images while preserving important features. + + Args: + x (np.ndarray): Input array + kernel_size (Union[int, Tuple[int, int]]): Size of the 2D pooling kernel. + If int, same size is used for both dimensions. + stride (Union[int, Tuple[int, int]]): Stride of the 2D pooling operation. + If int, same stride is used for both dimensions. + padding (Union[int, Tuple[int, int]]): Amount of padding for both dimensions. + If int, same padding is used for both dimensions. + axis (Tuple[int, int], optional): Two axes to perform max pooling over. + Defaults to (-2, -1) (last two dimensions). + + Returns: + np.ndarray: 2D max pooled array with reduced size along the specified axes + + Example: + >>> image = np.random.rand(64, 64) + >>> pooled = max_pool_2d(image, kernel_size=2, stride=2, padding=0) + >>> # Reduces image size from 64x64 to 32x32 with 2x2 max pooling + """ + if isinstance(kernel_size, Number): + kernel_size = (kernel_size, kernel_size) + if isinstance(stride, Number): + stride = (stride, stride) + if isinstance(padding, Number): + padding = (padding, padding) + axis = tuple(axis) + return max_pool_nd(x, kernel_size, stride, padding, axis) + + +@no_warnings(category=RuntimeWarning) +def depth_edge( + depth: np.ndarray, + atol: float = None, + rtol: float = None, + kernel_size: int = 3, + mask: np.ndarray = None, +) -> np.ndarray: + """Compute the edge mask from depth map. The edge is defined as the pixels whose neighbors have + large difference in depth. + + Args: + depth (np.ndarray): shape (..., height, width), linear depth map + atol (float): absolute tolerance + rtol (float): relative tolerance + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + if mask is None: + diff = max_pool_2d( + depth, kernel_size, stride=1, padding=kernel_size // 2 + ) + max_pool_2d(-depth, kernel_size, stride=1, padding=kernel_size // 2) + else: + diff = max_pool_2d( + np.where(mask, depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + max_pool_2d( + np.where(mask, -depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + + edge = np.zeros_like(depth, dtype=bool) + if atol is not None: + edge |= diff > atol + + if rtol is not None: + edge |= diff / depth > rtol + return edge + + +def depth_aliasing( + depth: np.ndarray, + atol: float = None, + rtol: float = None, + kernel_size: int = 3, + mask: np.ndarray = None, +) -> np.ndarray: + """ + Compute the map that indicates the aliasing of x depth map. The aliasing is defined as the pixels which neither close to the maximum nor the minimum of its neighbors. + Args: + depth (np.ndarray): shape (..., height, width), linear depth map + atol (float): absolute tolerance + rtol (float): relative tolerance + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + if mask is None: + diff_max = ( + max_pool_2d(depth, kernel_size, stride=1, padding=kernel_size // 2) + - depth + ) + diff_min = ( + max_pool_2d(-depth, kernel_size, stride=1, padding=kernel_size // 2) + + depth + ) + else: + diff_max = ( + max_pool_2d( + np.where(mask, depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + - depth + ) + diff_min = ( + max_pool_2d( + np.where(mask, -depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + + depth + ) + diff = np.minimum(diff_max, diff_min) + + edge = np.zeros_like(depth, dtype=bool) + if atol is not None: + edge |= diff > atol + if rtol is not None: + edge |= diff / depth > rtol + return edge + + +@no_warnings(category=RuntimeWarning) +def normals_edge( + normals: np.ndarray, + tol: float, + kernel_size: int = 3, + mask: np.ndarray = None, +) -> np.ndarray: + """Compute the edge mask from normal map. + + Args: + normal (np.ndarray): shape (..., height, width, 3), normal map + tol (float): tolerance in degrees + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + assert ( + normals.ndim >= 3 and normals.shape[-1] == 3 + ), "normal should be of shape (..., height, width, 3)" + normals = normals / ( + np.linalg.norm(normals, axis=-1, keepdims=True) + 1e-12 + ) + + padding = kernel_size // 2 + normals_window = sliding_window_2d( + np.pad( + normals, + ( + *([(0, 0)] * (normals.ndim - 3)), + (padding, padding), + (padding, padding), + (0, 0), + ), + mode="edge", + ), + window_size=kernel_size, + stride=1, + axis=(-3, -2), + ) + if mask is None: + angle_diff = np.arccos( + (normals[..., None, None] * normals_window).sum(axis=-3) + ).max(axis=(-2, -1)) + else: + mask_window = sliding_window_2d( + np.pad( + mask, + ( + *([(0, 0)] * (mask.ndim - 3)), + (padding, padding), + (padding, padding), + ), + mode="edge", + ), + window_size=kernel_size, + stride=1, + axis=(-3, -2), + ) + angle_diff = np.where( + mask_window, + np.arccos((normals[..., None, None] * normals_window).sum(axis=-3)), + 0, + ).max(axis=(-2, -1)) + + angle_diff = max_pool_2d( + angle_diff, kernel_size, stride=1, padding=kernel_size // 2 + ) + edge = angle_diff > np.deg2rad(tol) + return edge diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/gs_effects.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/gs_effects.py new file mode 100644 index 0000000000000000000000000000000000000000..e9a87fd6cabd44bd263663b073754d55d7e61969 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/gs_effects.py @@ -0,0 +1,327 @@ +from math import atan2, cos, exp, floor, sin, sqrt + +import numpy as np +import torch + + +def fract(x): + """Get fractional part of a number.""" + if isinstance(x, torch.Tensor): + return x - torch.floor(x) + return x - floor(x) + + +class GSEffects: + """Convert GLSL GS render effects to PyTorch - vectorized for batch processing""" + + def __init__(self, start_time=0.0, end_time=10.0): + """Initialize effects with time range. + + Args: + start_time: Animation start time + end_time: Animation end time + """ + self.start_time = start_time + self.end_time = end_time + + @staticmethod + def smoothstep(edge0, edge1, x): + """GLSL smoothstep function (vectorized)""" + if isinstance(x, torch.Tensor): + result = torch.zeros_like(x, dtype=x.dtype) + mask_low = x < edge0 + mask_high = x > edge1 + mask_mid = ~(mask_low | mask_high) + + t = (x[mask_mid] - edge0) / (edge1 - edge0) + result[mask_mid] = t * t * (3.0 - 2.0 * t) + result[mask_low] = 0.0 + result[mask_high] = 1.0 + return result + else: + if x < edge0: + return 0.0 + if x > edge1: + return 1.0 + t = (x - edge0) / (edge1 - edge0) + return t * t * (3.0 - 2.0 * t) + + @staticmethod + def step(edge, x): + """GLSL step function (vectorized)""" + if isinstance(x, torch.Tensor): + return (x >= edge).to(x.dtype) + if isinstance(edge, torch.Tensor): + return (x >= edge).to(edge.dtype) + return 1.0 if x >= edge else 0.0 + + @staticmethod + def mix(x, y, a): + """GLSL mix function (linear interpolation, vectorized)""" + return x * (1.0 - a) + y * a + + @staticmethod + def clamp(x, min_val, max_val): + """Clamp value between min and max (vectorized)""" + if isinstance(x, torch.Tensor): + return torch.clamp(x, min_val, max_val) + return max(min_val, min(max_val, x)) + + @staticmethod + def length_xz(pos): + """Calculate length of XZ components (vectorized)""" + if pos.dim() == 1: + return torch.sqrt(pos[0] ** 2 + pos[2] ** 2) + return torch.sqrt(pos[:, 0] ** 2 + pos[:, 2] ** 2) + + @staticmethod + def length_vec(v): + """Calculate vector length (vectorized)""" + if v.dim() == 1: + return torch.sqrt(torch.sum(v**2)) + return torch.sqrt(torch.sum(v**2, dim=1)) + + @staticmethod + def hash(p): + """Pseudo-random hash function (vectorized)""" + p = fract(p * 0.3183099 + 0.1) + p = p * 17.0 + return torch.stack( + [ + fract(p[:, 0] * p[:, 1] * p[:, 2]), + fract(p[:, 0] + p[:, 1] * p[:, 2]), + fract(p[:, 0] * p[:, 1] + p[:, 2]), + ], + dim=1, + ) + + @staticmethod + def noise(p): + """3D Perlin-style noise function (vectorized)""" + i = torch.floor(p).to(torch.long) + f = fract(p) + f = f * f * (3.0 - 2.0 * f) + + def get_hash_offset(offset): + return GSEffects.hash(i.to(p.dtype) + offset) + + n000 = get_hash_offset( + torch.tensor([0, 0, 0], dtype=p.dtype, device=p.device) + ) + n100 = get_hash_offset( + torch.tensor([1, 0, 0], dtype=p.dtype, device=p.device) + ) + n010 = get_hash_offset( + torch.tensor([0, 1, 0], dtype=p.dtype, device=p.device) + ) + n110 = get_hash_offset( + torch.tensor([1, 1, 0], dtype=p.dtype, device=p.device) + ) + n001 = get_hash_offset( + torch.tensor([0, 0, 1], dtype=p.dtype, device=p.device) + ) + n101 = get_hash_offset( + torch.tensor([1, 0, 1], dtype=p.dtype, device=p.device) + ) + n011 = get_hash_offset( + torch.tensor([0, 1, 1], dtype=p.dtype, device=p.device) + ) + n111 = get_hash_offset( + torch.tensor([1, 1, 1], dtype=p.dtype, device=p.device) + ) + + x0 = GSEffects.mix(n000, n100, f[:, 0:1]) + x1 = GSEffects.mix(n010, n110, f[:, 0:1]) + x2 = GSEffects.mix(n001, n101, f[:, 0:1]) + x3 = GSEffects.mix(n011, n111, f[:, 0:1]) + + y0 = GSEffects.mix(x0, x1, f[:, 1:2]) + y1 = GSEffects.mix(x2, x3, f[:, 1:2]) + + return GSEffects.mix(y0, y1, f[:, 2:3]) + + @staticmethod + def rot_2d(angle): + """2D rotation (vectorized)""" + if isinstance(angle, torch.Tensor): + s = torch.sin(angle) + c = torch.cos(angle) + rot = torch.stack( + [torch.stack([c, -s], dim=-1), torch.stack([s, c], dim=-1)], + dim=-2, + ).squeeze() + else: + s = np.sin(angle) + c = np.cos(angle) + rot = torch.tensor([[c, -s], [s, c]]).cuda().float() + return rot + + def twister(self, pos, scale, t): + h = self.hash(pos)[:, 0:1] + 0.1 + pos_xz_len = self.length_xz(pos) + s = self.smoothstep(0.0, 8.0, t * t * 0.1 - pos_xz_len * 2.0 + 2.0)[ + :, None + ] + mask = torch.linalg.norm(scale, dim=-1, keepdim=True) < 0.05 + pos_y = torch.where( + mask, (-10.0 + pos[:, 1:2]) * (s ** (2 * h)), pos[:, 1:2] + ) + pos_xz = pos[:, [0, 2]] * torch.exp( + -1 * torch.linalg.norm(pos[:, [0, 2]], dim=-1, keepdim=True) + ) + pos_xz = torch.einsum( + "n i, n i j -> n j", + pos_xz, + self.rot_2d(t * 0.2 + pos[:, 1:2] * 20.0 * (1 - s)), + ) + pos_new = torch.cat([pos_xz[:, 0:1], pos_y, pos_xz[:, 1:2]], dim=-1) + return pos_new, s**4 + + def rain(self, pos, scale, t): + h = self.hash(pos) + pos_xz_len = self.length_xz(pos) + s = self.smoothstep(0.0, 5.0, t * t * 0.1 - pos_xz_len * 2.0 + 1.0) ** ( + 0.5 + h[:, 0] + ) + y = pos[:, 1:2] + pos_y = torch.minimum(-10.0 + s[:, None] * 15.0, pos[:, 1:2]) + pos_x = pos[:, 0:1] + pos_y * 0.2 + pos_xz = torch.cat([pos_x, pos[:, 2:3]], dim=-1) + pos_xz = pos_xz * torch.matmul( + self.rot_2d(t * 0.3), torch.ones_like(pos_xz).unsqueeze(-1) + ).squeeze(-1) + pos_new = torch.cat([pos_xz[:, 0:1], pos_y, pos_xz[:, 1:2]], dim=-1) + a = self.smoothstep(-10.0, y.squeeze(), pos_y.squeeze())[:, None] + return pos_new, a + + def apply_effect(self, gsplat, t, effect_type, ignore_scale=False): + """Apply the effect shader logic (vectorized for batch processing) + + Args: + gsplat: Dictionary with: + 'means': (n, 3) tensor + 'scales': (n, 3) tensor + 'colors': (n, 3) tensor + 'quats': (n, 4) tensor + 'opacities': (n,) tensor + t: Current time (normalized based on start_time and end_time) + effect_type: 2=Spread + + Returns: + Modified gsplat dictionary + """ + # Normalize time to animation range + normalized_t = t - self.start_time + device = gsplat["means"].device + dtype = gsplat["means"].dtype + + output = { + "means": gsplat["means"].clone(), + "quats": gsplat["quats"].clone(), + "scales": gsplat["scales"].clone(), + "opacities": gsplat["opacities"].clone(), + "colors": gsplat["colors"].clone(), + } + + s = self.smoothstep(0.0, 10.0, normalized_t - 3.2) * 10.0 + scales = output["scales"] + local_pos = output["means"].clone() + l = self.length_xz(local_pos) + smoothstep_val = None + + if effect_type == 2: # Spread Effect + border = torch.abs(s - l - 0.5) + decay = 1.0 - 0.2 * torch.exp(-20.0 * border) + # decay = 1.0 - 0.7 * torch.exp(-10.0 * border) + local_pos = local_pos * decay[:, None] + + smoothstep_val = self.smoothstep(s - 0.5, s, l + 0.5) + # final_scales = self.mix(scales, 0.002, smoothstep_val[:, None]) + if not ignore_scale: + final_scales = self.mix(scales, 1e-9, smoothstep_val[:, None]) + else: + final_scales = scales + + noise_input = torch.stack( + [ + local_pos[:, 0] * 2.0 + normalized_t * 0.5, + local_pos[:, 1] * 2.0 + normalized_t * 0.5, + local_pos[:, 2] * 2.0 + normalized_t * 0.5, + ], + dim=1, + ) + noise_val = self.noise(noise_input) + + output["means"] = ( + local_pos + 0.0 * noise_val * smoothstep_val[:, None] + ) + output["scales"] = final_scales + + at = torch.atan2(local_pos[:, 0], local_pos[:, 2]) / 3.1416 + output["colors"] *= self.step(at, normalized_t - 3.1416)[:, None] + output["colors"] += ( + torch.exp(-20.0 * border) + + torch.exp(-50.0 * torch.abs(normalized_t - at - 3.1416)) * 0.5 + )[:, None] + output["opacities"] *= self.step(at, normalized_t - 3.1416) + output["opacities"] += ( + torch.exp(-20.0 * border) + + torch.exp(-50.0 * torch.abs(normalized_t - at - 3.1416)) * 0.5 + ) + + # ===== New feature: Randomly mask points based on smoothstep_val ===== + # Higher smoothstep_val means higher probability of masking + mask_prob = ( + smoothstep_val.squeeze() + if smoothstep_val.dim() > 1 + else smoothstep_val + ) + if not hasattr(self, "random_vals"): + self.random_vals = torch.rand( + mask_prob.shape, device=device, dtype=dtype + ) + mask = ( + self.random_vals < mask_prob * 0.8 + ) # True indicates the point is masked + + # Apply mask to various attributes + if not ignore_scale: + output["means"][mask] *= 0 # Or can be set to other values + output["scales"][ + mask + ] *= 0 # Set scales to 0 to make points invisible + output["opacities"][ + mask + ] *= 0 # Set opacity to 0 to make points transparent + + return output, smoothstep_val + + +# Usage example +if __name__ == "__main__": + # Create effects processor with time range from 0 to 10 seconds + effects = GSEffects(start_time=0.0, end_time=10.0) + + # Sample gsplat data (batch) + n_points = 100 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + sample_gsplat = { + "means": torch.randn(n_points, 3, dtype=torch.float32, device=device), + "quats": torch.randn(n_points, 4, dtype=torch.float32, device=device), + "scales": torch.rand(n_points, 3, dtype=torch.float32, device=device), + "opacities": torch.rand(n_points, dtype=torch.float32, device=device), + "colors": torch.rand(n_points, 3, dtype=torch.float32, device=device), + } + + # Apply Magic effect at different time points + for t in [0.0, 2.5, 5.0, 7.5, 10.0]: + result = effects.apply_effect(sample_gsplat, t, effect_type=2) + print(f"\nTime: {t}s") + print(f"Center shape: {result['means'].shape}") + print(f"Center[0]: {result['means'][0]}") + print(f"Scales shape: {result['scales'].shape}") + print(f"Scales[0]: {result['scales'][0]}") + print(f"RGB shape: {result['colors'].shape}") + print(f"RGB[0]: {result['colors'][0]}") + print(f"Opacity shape: {result['opacities'].shape}") + print(f"Opacity[0]: {result['opacities'][0]}") diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/inference_utils.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/inference_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ca90b8a8c71d7805ea821299f1971f40d9c4c8fd --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/inference_utils.py @@ -0,0 +1,319 @@ +import torch +from PIL import Image +from torchvision import transforms + +import glob +import os +from reward_function.HunyuanWorldMirror.src.utils.video_utils import ( + video_to_image_frames, +) + +IMAGE_EXTS = ["*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.webp"] +VIDEO_EXTS = [".mp4", ".avi", ".mov", ".webm", ".gif"] + + +def load_and_preprocess_images( + image_file_paths, preprocessing_mode="crop", output_size=518 +): + """Transform raw image files into model-ready tensor batches with standardized dimensions. + + This utility function handles the complete pipeline from file paths to batched tensors, + ensuring compatibility with neural network requirements while preserving image quality. + + Args: + image_file_paths (list): Collection of file system paths pointing to image files + preprocessing_mode (str, optional): Image transformation strategy: + - "crop" (default): Resize width to 518px, center-crop height if oversized + - "pad": Scale largest dimension to 518px, pad smaller dimension to square + output_size (int, optional): Target dimension for model input (default: 518) + + Returns: + torch.Tensor: Processed image batch with shape (1, N, 3, H, W) ready for model inference + + Raises: + ValueError: When input validation fails (empty list or invalid mode) + + Implementation Details: + - Automatic alpha channel handling: RGBA images composited onto white backgrounds + - Dimension normalization: All outputs divisible by 14 for patch-based processing + - Batch consistency: Different-sized images padded to uniform dimensions + - Memory optimization: Efficient tensor operations with minimal data copying + - Quality preservation: Bicubic resampling maintains visual fidelity + """ + # Input validation and parameter setup + if len(image_file_paths) == 0: + raise ValueError("At least 1 image is required") + + if preprocessing_mode not in ["crop", "pad"]: + raise ValueError("preprocessing_mode must be either 'crop' or 'pad'") + + processed_image_list = [] + image_dimension_set = set() + tensor_converter = transforms.ToTensor() + model_target_size = output_size + + # Individual image processing pipeline + for image_file_path in image_file_paths: + # File system to memory conversion + loaded_image = Image.open(image_file_path) + + # Transparency handling for RGBA images + if loaded_image.mode == "RGBA": + # Generate white canvas matching image dimensions + white_background = Image.new( + "RGBA", loaded_image.size, (255, 255, 255, 255) + ) + # Blend transparent pixels with white background + loaded_image = Image.alpha_composite(white_background, loaded_image) + + # Format standardization to RGB + loaded_image = loaded_image.convert("RGB") + + original_width, original_height = loaded_image.size + + # Dimension calculation based on preprocessing strategy + if preprocessing_mode == "pad": + # Proportional scaling to fit largest dimension within target + if original_width >= original_height: + scaled_width = model_target_size + scaled_height = ( + round( + original_height * (scaled_width / original_width) / 14 + ) + * 14 + ) # Patch compatibility + else: + scaled_height = model_target_size + scaled_width = ( + round( + original_width * (scaled_height / original_height) / 14 + ) + * 14 + ) # Patch compatibility + else: # preprocessing_mode == "crop" + # Width normalization with proportional height adjustment + scaled_width = model_target_size + scaled_height = ( + round(original_height * (scaled_width / original_width) / 14) + * 14 + ) + + # High-quality image resizing + loaded_image = loaded_image.resize( + (scaled_width, scaled_height), Image.Resampling.BICUBIC + ) + image_tensor = tensor_converter( + loaded_image + ) # Normalize to [0, 1] range + + # Height trimming for crop mode (center-based) + if preprocessing_mode == "crop" and scaled_height > model_target_size: + crop_start_y = (scaled_height - model_target_size) // 2 + image_tensor = image_tensor[ + :, crop_start_y : crop_start_y + model_target_size, : + ] + + # Square padding for pad mode (centered) + if preprocessing_mode == "pad": + height_padding_needed = model_target_size - image_tensor.shape[1] + width_padding_needed = model_target_size - image_tensor.shape[2] + + if height_padding_needed > 0 or width_padding_needed > 0: + padding_top = height_padding_needed // 2 + padding_bottom = height_padding_needed - padding_top + padding_left = width_padding_needed // 2 + padding_right = width_padding_needed - padding_left + + # White padding application (value=1.0 for normalized images) + image_tensor = torch.nn.functional.pad( + image_tensor, + (padding_left, padding_right, padding_top, padding_bottom), + mode="constant", + value=1.0, + ) + + image_dimension_set.add((image_tensor.shape[1], image_tensor.shape[2])) + processed_image_list.append(image_tensor) + + # Cross-image dimension harmonization + if len(image_dimension_set) > 1: + print( + f"Warning: Found images with different shapes: {image_dimension_set}" + ) + # Calculate maximum dimensions across the batch + maximum_height = max(dimension[0] for dimension in image_dimension_set) + maximum_width = max(dimension[1] for dimension in image_dimension_set) + + # Uniform padding to achieve batch consistency + uniformly_sized_images = [] + for image_tensor in processed_image_list: + height_padding_needed = maximum_height - image_tensor.shape[1] + width_padding_needed = maximum_width - image_tensor.shape[2] + + if height_padding_needed > 0 or width_padding_needed > 0: + padding_top = height_padding_needed // 2 + padding_bottom = height_padding_needed - padding_top + padding_left = width_padding_needed // 2 + padding_right = width_padding_needed - padding_left + + image_tensor = torch.nn.functional.pad( + image_tensor, + (padding_left, padding_right, padding_top, padding_bottom), + mode="constant", + value=1.0, + ) + uniformly_sized_images.append(image_tensor) + processed_image_list = uniformly_sized_images + + # Batch tensor construction + batched_images = torch.stack( + processed_image_list + ) # Concatenate along batch dimension + + # Single image batch dimension handling + if len(image_file_paths) == 1: + # Ensure proper 4D tensor structure (batch, channels, height, width) + if batched_images.dim() == 3: + batched_images = batched_images.unsqueeze(0) + + return batched_images.unsqueeze(0) + + +def _handle_alpha_channel(img_data): + """Process RGBA images by blending with white background.""" + if img_data.mode == "RGBA": + white_bg = Image.new("RGBA", img_data.size, (255, 255, 255, 255)) + img_data = Image.alpha_composite(white_bg, img_data) + return img_data.convert("RGB") + + +def _calculate_resize_dims( + orig_w, orig_h, max_dim, resize_strategy, patch_size=14 +): + """Calculate new dimensions based on resize strategy.""" + if resize_strategy == "pad": + if orig_w >= orig_h: + new_w = max_dim + new_h = round(orig_h * (new_w / orig_w) / patch_size) * patch_size + else: + new_h = max_dim + new_w = round(orig_w * (new_h / orig_h) / patch_size) * patch_size + else: # crop strategy + new_w = max_dim + new_h = round(orig_h * (new_w / orig_w) / patch_size) * patch_size + return new_w, new_h + + +def _apply_padding(tensor_img, target_dim): + """Apply padding to make tensor square.""" + h_pad = target_dim - tensor_img.shape[1] + w_pad = target_dim - tensor_img.shape[2] + + if h_pad > 0 or w_pad > 0: + pad_top, pad_bottom = h_pad // 2, h_pad - h_pad // 2 + pad_left, pad_right = w_pad // 2, w_pad - w_pad // 2 + return torch.nn.functional.pad( + tensor_img, + (pad_left, pad_right, pad_top, pad_bottom), + mode="constant", + value=1.0, + ) + return tensor_img + + +def prepare_images_to_tensor( + file_paths, resize_strategy="crop", target_size=518 +): + """Process image files into uniform tensor batch for model input. + + Args: + file_paths (list): Paths to image files + resize_strategy (str): "crop" or "pad" processing mode + target_size (int): Target size for processing + + Returns: + torch.Tensor: Processed image batch (1, N, 3, H, W) + """ + if not file_paths: + raise ValueError("At least 1 image is required") + + if resize_strategy not in ["crop", "pad"]: + raise ValueError("Strategy must be 'crop' or 'pad'") + + tensor_list = [] + dimension_set = set() + converter = transforms.ToTensor() + + # Process each image file + for file_path in file_paths: + img_data = Image.open(file_path) + img_data = _handle_alpha_channel(img_data) + + orig_w, orig_h = img_data.size + new_w, new_h = _calculate_resize_dims( + orig_w, orig_h, target_size, resize_strategy + ) + + # Resize and convert to tensor + img_data = img_data.resize((new_w, new_h), Image.Resampling.BICUBIC) + tensor_img = converter(img_data) + + # Apply center crop for crop strategy + if resize_strategy == "crop" and new_h > target_size: + crop_start = (new_h - target_size) // 2 + tensor_img = tensor_img[:, crop_start : crop_start + target_size, :] + + # Apply padding for pad strategy + if resize_strategy == "pad": + tensor_img = _apply_padding(tensor_img, target_size) + + dimension_set.add((tensor_img.shape[1], tensor_img.shape[2])) + tensor_list.append(tensor_img) + + # Handle mixed dimensions + if len(dimension_set) > 1: + print(f"Warning: Mixed image dimensions found: {dimension_set}") + max_h = max(dims[0] for dims in dimension_set) + max_w = max(dims[1] for dims in dimension_set) + + tensor_list = [ + ( + _apply_padding(img, max(max_h, max_w)) + if img.shape[1] != max_h or img.shape[2] != max_w + else img + ) + for img in tensor_list + ] + + batch_tensor = torch.stack(tensor_list) + + # Ensure proper batch dimensions + if batch_tensor.dim() == 3: + batch_tensor = batch_tensor.unsqueeze(0) + + return batch_tensor.unsqueeze(0) + + +def extract_load_and_preprocess_images( + image_folder_or_video_path, fps=1, target_size=518, mode="crop" +): + # Support multiple image formats + if ( + image_folder_or_video_path.is_file() + and image_folder_or_video_path.suffix.lower() in VIDEO_EXTS + ): + frame_paths = video_to_image_frames( + str(image_folder_or_video_path), fps=fps + ) + img_paths = sorted(frame_paths) + else: + img_paths = [] + for ext in IMAGE_EXTS: + img_paths.extend( + glob.glob(os.path.join(str(image_folder_or_video_path), ext)) + ) + img_paths = sorted(img_paths) + images = prepare_images_to_tensor( + img_paths, resize_strategy=mode, target_size=target_size + ) + return images diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/render_utils.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/render_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f564677401f3749d4ab5044fe884a2e09b1b74d5 --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/render_utils.py @@ -0,0 +1,480 @@ +from pathlib import Path + +import numpy as np +import torch +import moviepy.editor as mpy + +from reward_function.HunyuanWorldMirror.src.models.models.rasterization import ( + GaussianSplatRenderer, +) +from reward_function.HunyuanWorldMirror.src.models.utils.sh_utils import ( + RGB2SH, + SH2RGB, +) +from reward_function.HunyuanWorldMirror.src.utils.gs_effects import GSEffects +from reward_function.HunyuanWorldMirror.src.utils.color_map import ( + apply_color_map_to_image, +) +from tqdm import tqdm + + +def rotation_matrix_to_quaternion(R): + """Convert rotation matrix to quaternion.""" + trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + + q = torch.zeros(R.shape[:-2] + (4,), device=R.device, dtype=R.dtype) + + # Case where trace > 0 + mask1 = trace > 0 + s = torch.sqrt(trace[mask1] + 1.0) * 2 # s=4*qw + q[mask1, 0] = 0.25 * s # qw + q[mask1, 1] = (R[mask1, 2, 1] - R[mask1, 1, 2]) / s # qx + q[mask1, 2] = (R[mask1, 0, 2] - R[mask1, 2, 0]) / s # qy + q[mask1, 3] = (R[mask1, 1, 0] - R[mask1, 0, 1]) / s # qz + + # Case where R[0,0] > R[1,1] and R[0,0] > R[2,2] + mask2 = ( + (~mask1) & (R[..., 0, 0] > R[..., 1, 1]) & (R[..., 0, 0] > R[..., 2, 2]) + ) + s = ( + torch.sqrt(1.0 + R[mask2, 0, 0] - R[mask2, 1, 1] - R[mask2, 2, 2]) * 2 + ) # s=4*qx + q[mask2, 0] = (R[mask2, 2, 1] - R[mask2, 1, 2]) / s # qw + q[mask2, 1] = 0.25 * s # qx + q[mask2, 2] = (R[mask2, 0, 1] + R[mask2, 1, 0]) / s # qy + q[mask2, 3] = (R[mask2, 0, 2] + R[mask2, 2, 0]) / s # qz + + # Case where R[1,1] > R[2,2] + mask3 = (~mask1) & (~mask2) & (R[..., 1, 1] > R[..., 2, 2]) + s = ( + torch.sqrt(1.0 + R[mask3, 1, 1] - R[mask3, 0, 0] - R[mask3, 2, 2]) * 2 + ) # s=4*qy + q[mask3, 0] = (R[mask3, 0, 2] - R[mask3, 2, 0]) / s # qw + q[mask3, 1] = (R[mask3, 0, 1] + R[mask3, 1, 0]) / s # qx + q[mask3, 2] = 0.25 * s # qy + q[mask3, 3] = (R[mask3, 1, 2] + R[mask3, 2, 1]) / s # qz + + # Remaining case + mask4 = (~mask1) & (~mask2) & (~mask3) + s = ( + torch.sqrt(1.0 + R[mask4, 2, 2] - R[mask4, 0, 0] - R[mask4, 1, 1]) * 2 + ) # s=4*qz + q[mask4, 0] = (R[mask4, 1, 0] - R[mask4, 0, 1]) / s # qw + q[mask4, 1] = (R[mask4, 0, 2] + R[mask4, 2, 0]) / s # qx + q[mask4, 2] = (R[mask4, 1, 2] + R[mask4, 2, 1]) / s # qy + q[mask4, 3] = 0.25 * s # qz + + return q + + +def quaternion_to_rotation_matrix(q): + """Convert quaternion to rotation matrix.""" + w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3] + + # Normalize quaternion + norm = torch.sqrt(w * w + x * x + y * y + z * z) + w, x, y, z = w / norm, x / norm, y / norm, z / norm + + R = torch.zeros(q.shape[:-1] + (3, 3), device=q.device, dtype=q.dtype) + + R[..., 0, 0] = 1 - 2 * (y * y + z * z) + R[..., 0, 1] = 2 * (x * y - w * z) + R[..., 0, 2] = 2 * (x * z + w * y) + R[..., 1, 0] = 2 * (x * y + w * z) + R[..., 1, 1] = 1 - 2 * (x * x + z * z) + R[..., 1, 2] = 2 * (y * z - w * x) + R[..., 2, 0] = 2 * (x * z - w * y) + R[..., 2, 1] = 2 * (y * z + w * x) + R[..., 2, 2] = 1 - 2 * (x * x + y * y) + + return R + + +def slerp_quaternions(q1, q2, t): + """Spherical linear interpolation between quaternions.""" + # Compute dot product + dot = (q1 * q2).sum(dim=-1, keepdim=True) + + # If dot product is negative, slerp won't take the shorter path. + # Note that q and -q represent the same rotation, so we can flip one. + mask = dot < 0 + q2 = torch.where(mask, -q2, q2) + dot = torch.where(mask, -dot, dot) + + # If the inputs are too close for comfort, linearly interpolate + # and normalize the result. + DOT_THRESHOLD = 0.9995 + mask_linear = dot > DOT_THRESHOLD + + result = torch.zeros_like(q1) + + # Linear interpolation for close quaternions + if mask_linear.any(): + result_linear = q1 + t * (q2 - q1) + norm = torch.norm(result_linear, dim=-1, keepdim=True) + result_linear = result_linear / norm + result = torch.where(mask_linear, result_linear, result) + + # Spherical interpolation for distant quaternions + mask_slerp = ~mask_linear + if mask_slerp.any(): + theta_0 = torch.acos(torch.abs(dot)) + sin_theta_0 = torch.sin(theta_0) + + theta = theta_0 * t + sin_theta = torch.sin(theta) + + s0 = torch.cos(theta) - dot * sin_theta / sin_theta_0 + s1 = sin_theta / sin_theta_0 + + result_slerp = (s0 * q1) + (s1 * q2) + result = torch.where(mask_slerp, result_slerp, result) + + return result + + +def render_interpolated_video( + gs_renderer: GaussianSplatRenderer, + splats: dict, + camtoworlds: torch.Tensor, + intrinsics: torch.Tensor, + hw: tuple[int, int], + out_path: Path, + interp_per_pair: int = 20, + loop_reverse: bool = True, + effects: GSEffects = None, + effect_type: int = 2, + save_mode: str = "split", +) -> None: + # camtoworlds: [B, S, 4, 4], intrinsics: [B, S, 3, 3] + b, s, _, _ = camtoworlds.shape + h, w = hw + + # Build interpolated trajectory + def build_interpolated_traj(index, nums): + exts, ints = [], [] + tmp_camtoworlds = camtoworlds[:, index] + tmp_intrinsics = intrinsics[:, index] + for i in range(len(index) - 1): + exts.append(tmp_camtoworlds[:, i : i + 1]) + ints.append(tmp_intrinsics[:, i : i + 1]) + # Extract rotation and translation + R0, t0 = tmp_camtoworlds[:, i, :3, :3], tmp_camtoworlds[:, i, :3, 3] + R1, t1 = ( + tmp_camtoworlds[:, i + 1, :3, :3], + tmp_camtoworlds[:, i + 1, :3, 3], + ) + + # Convert rotations to quaternions + q0 = rotation_matrix_to_quaternion(R0) + q1 = rotation_matrix_to_quaternion(R1) + + # Interpolate using smooth quaternion slerp + for j in range(1, nums + 1): + alpha = j / (nums + 1) + + # Linear interpolation for translation + t_interp = (1 - alpha) * t0 + alpha * t1 + + # Spherical interpolation for rotation + q_interp = slerp_quaternions(q0, q1, alpha) + R_interp = quaternion_to_rotation_matrix(q_interp) + + # Create interpolated extrinsic matrix + ext = torch.eye( + 4, device=R_interp.device, dtype=R_interp.dtype + )[None].repeat(b, 1, 1) + ext[:, :3, :3] = R_interp + ext[:, :3, 3] = t_interp + + # Linear interpolation for intrinsics + K0 = tmp_intrinsics[:, i] + K1 = tmp_intrinsics[:, i + 1] + K = (1 - alpha) * K0 + alpha * K1 + + exts.append(ext[:, None]) + ints.append(K[:, None]) + + exts = torch.cat(exts, dim=1)[:1] + ints = torch.cat(ints, dim=1)[:1] + return exts, ints + + # Build wobble trajectory + def build_wobble_traj(nums, delta): + assert s == 1 + t = torch.linspace( + 0, 1, nums, dtype=torch.float32, device=camtoworlds.device + ) + t = (torch.cos(torch.pi * (t + 1)) + 1) / 2 + tf = torch.eye(4, dtype=torch.float32, device=camtoworlds.device) + radius = delta * 0.15 + tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone() + radius = radius[..., None] + radius = radius * t + tf[..., 0, 3] = torch.sin(2 * torch.pi * t) * radius + tf[..., 1, 3] = -torch.cos(2 * torch.pi * t) * radius + exts = camtoworlds @ tf + ints = intrinsics.repeat(1, exts.shape[1], 1, 1) + return exts, ints + + if s > 1: + all_ext, all_int = build_interpolated_traj( + [i for i in range(s)], interp_per_pair + ) + else: + all_ext, all_int = build_wobble_traj( + interp_per_pair * 12, + splats["means"][0].median(dim=0).values.norm(dim=-1)[None], + ) + + rendered_rgbs, rendered_depths = [], [] + chunk = 40 if effects is None else 1 + t = 0 + t_skip = 0 + if effects is not None: + try: + pruned_splats = gs_renderer.prune_gs(splats, gs_renderer.voxel_size) + except: + pruned_splats = splats + # indices = [x for x in range(0, all_ext.shape[1], 2)][:4] + # add_ext, add_int = build_interpolated_traj(indices, 150) + # add_ext = torch.flip(add_ext, dims=[1]) + # add_int = torch.flip(add_int, dims=[1]) + add_ext = all_ext[:, :1, :, :].repeat(1, 320, 1, 1) + add_int = all_int[:, :1, :, :].repeat(1, 320, 1, 1) + shift = pruned_splats["means"][0].median(dim=0).values + scale_factor = ( + (pruned_splats["means"][0] - shift) + .abs() + .quantile(0.95, dim=0) + .max() + ) + all_ext[0, :, :3, -1] = (all_ext[0, :, :3, -1] - shift) / scale_factor + add_ext[0, :, :3, -1] = (add_ext[0, :, :3, -1] - shift) / scale_factor + flag = None + try: + raw_splats = gs_renderer.rasterizer.runner.splats + except: + pass + for st in range(0, add_ext.shape[1]): + ed = min(st + 1, add_ext.shape[1]) + assert gs_renderer.sh_degree == 0 + if flag is not None and (flag < 0.99).any(): + break + sample_gsplat = { + "means": (pruned_splats["means"][0] - shift) / scale_factor, + "quats": pruned_splats["quats"][0], + "scales": pruned_splats["scales"][0] / scale_factor, + "opacities": pruned_splats["opacities"][0], + "colors": SH2RGB(pruned_splats["sh"][0].reshape(-1, 3)), + } + effects_splats, flag = effects.apply_effect( + sample_gsplat, t, effect_type=effect_type + ) + t += 0.04 + effects_splats["sh"] = RGB2SH(effects_splats["colors"]).reshape( + -1, 1, 3 + ) + try: + gs_renderer.rasterizer.runner.splats + effects_splats["sh0"] = effects_splats["sh"][:, :1, :] + effects_splats["shN"] = effects_splats["sh"][:, 1:, :] + effects_splats["scales"] = effects_splats["scales"].log() + effects_splats["opacities"] = torch.logit( + torch.clamp(effects_splats["opacities"], 1e-6, 1 - 1e-6) + ) + gs_renderer.rasterizer.runner.splats = effects_splats + colors, depths, _ = gs_renderer.rasterizer.rasterize_batches( + None, + None, + None, + None, + None, + add_ext[:, st:ed].to(torch.float32), + add_int[:, st:ed].to(torch.float32), + width=w, + height=h, + sh_degree=gs_renderer.sh_degree, + ) + except: + colors, depths, _ = gs_renderer.rasterizer.rasterize_batches( + effects_splats["means"][None], + effects_splats["quats"][None], + effects_splats["scales"][None], + effects_splats["opacities"][None], + effects_splats["sh"][None], + add_ext[:, st:ed].to(torch.float32), + add_int[:, st:ed].to(torch.float32), + width=w, + height=h, + sh_degree=( + gs_renderer.sh_degree if "sh" in pruned_splats else None + ), + ) + + if st > add_ext.shape[1] * 0.14: + t_skip = t if t_skip == 0 else t_skip + # break + rendered_rgbs.append(colors) + rendered_depths.append(depths) + # if (flag == 0).all(): + # break + t_st = t + t_ed = 0 + loop_dir = 1 + ignore_scale = False + for st in tqdm(range(0, all_ext.shape[1], chunk)): + ed = min(st + chunk, all_ext.shape[1]) + if effects is not None: + try: + sample_gsplat = { + "means": (pruned_splats["means"][0] - shift) / scale_factor, + "quats": pruned_splats["quats"][0], + "scales": pruned_splats["scales"][0] / scale_factor, + "opacities": pruned_splats["opacities"][0], + "colors": SH2RGB(pruned_splats["sh"][0].reshape(-1, 3)), + } + except: + sample_gsplat = { + "means": (pruned_splats["means"][0] - shift) / scale_factor, + "quats": pruned_splats["quats"][0], + "scales": pruned_splats["scales"][0] / scale_factor, + "opacities": pruned_splats["opacities"][0], + "colors": SH2RGB(pruned_splats["sh"][0].reshape(-1, 3)), + } + effects_splats, flag = effects.apply_effect( + sample_gsplat, + t, + effect_type=effect_type, + ignore_scale=ignore_scale, + ) + if loop_dir < 0: + t -= 0.04 + else: + t += 0.04 + if flag.mean() < 0.01 and t_ed == 0: + t_ed = t + effects_splats["sh"] = RGB2SH(effects_splats["colors"]).reshape( + -1, 1, 3 + ) + effects_splats["sh0"] = effects_splats["sh"][:, :1, :] + effects_splats["shN"] = effects_splats["sh"][:, 1:, :] + try: + gs_renderer.rasterizer.runner.splats + effects_splats["sh0"] = effects_splats["sh"][:, :1, :] + effects_splats["shN"] = effects_splats["sh"][:, 1:, :] + effects_splats["scales"] = effects_splats["scales"].log() + effects_splats["opacities"] = torch.logit( + torch.clamp(effects_splats["opacities"], 1e-6, 1 - 1e-6) + ) + gs_renderer.rasterizer.runner.splats = effects_splats + colors, depths, _ = gs_renderer.rasterizer.rasterize_batches( + None, + None, + None, + None, + None, + all_ext[:, st:ed].to(torch.float32), + all_int[:, st:ed].to(torch.float32), + width=w, + height=h, + sh_degree=gs_renderer.sh_degree, + ) + except: + colors, depths, _ = gs_renderer.rasterizer.rasterize_batches( + effects_splats["means"][None], + effects_splats["quats"][None], + effects_splats["scales"][None], + effects_splats["opacities"][None], + effects_splats["sh"][None], + all_ext[:, st:ed].to(torch.float32), + all_int[:, st:ed].to(torch.float32), + width=w, + height=h, + sh_degree=( + gs_renderer.sh_degree if "sh" in pruned_splats else None + ), + ) + + if ( + t + > (all_ext.shape[1]) * 0.04 + + t_st + - (t_ed - t_st) * 2 + - 15 * 0.04 + or t < t_st + ): + # ignore_scale = True + loop_dir *= -1 + t = t_ed if loop_dir == -1 else t + else: + colors, depths, _ = gs_renderer.rasterizer.rasterize_batches( + splats["means"][:1], + splats["quats"][:1], + splats["scales"][:1], + splats["opacities"][:1], + splats["sh"][:1] if "sh" in splats else splats["colors"][:1], + all_ext[:, st:ed].to(torch.float32), + all_int[:, st:ed].to(torch.float32), + width=w, + height=h, + sh_degree=gs_renderer.sh_degree if "sh" in splats else None, + ) + rendered_rgbs.append(colors) + rendered_depths.append(depths) + + rgbs = torch.cat(rendered_rgbs, dim=1)[0] # [N, H, W, 3] + depths = torch.cat(rendered_depths, dim=1)[0, ..., 0] # [N, H, W] + + def depth_vis(d: torch.Tensor) -> torch.Tensor: + valid = d > 0 + if valid.any(): + near = d[valid].float().quantile(0.01).log() + else: + near = torch.tensor(0.0, device=d.device) + far = d.flatten().float().quantile(0.99).log() + x = d.float().clamp(min=1e-9).log() + x = 1.0 - (x - near) / (far - near + 1e-9) + return apply_color_map_to_image(x, "turbo") + + frames = [] + rgb_frames = [] + depth_frames = [] + + for rgb, dep in zip(rgbs, depths): + rgb_img = rgb.permute(2, 0, 1) # [3, H, W] + depth_img = depth_vis(dep) # [3, H, W] + + if save_mode == "both": + combined = torch.cat([rgb_img, depth_img], dim=1) # [3, 2*H, W] + frames.append(combined) + elif save_mode == "split": + rgb_frames.append(rgb_img) + depth_frames.append(depth_img) + else: + raise ValueError("save_mode must be 'both' or 'split'") + + def _make_video(frames, path): + video = torch.stack(frames).clamp(0, 1) # [N, 3, H, W] + video = video.permute(0, 2, 3, 1) # [N, H, W, 3] for moviepy + video = (video * 255).to(torch.uint8).cpu().numpy() + if loop_reverse and video.shape[0] > 1: + video = np.concatenate([video, video[::-1][1:-1]], axis=0) + clip = mpy.ImageSequenceClip(list(video), fps=30) + clip.write_videofile(str(path), logger=None) + + # Save videos + if save_mode == "both": + _make_video(frames, f"{out_path}.mp4") + elif save_mode == "split": + _make_video(rgb_frames, f"{out_path}_rgb.mp4") + _make_video(depth_frames, f"{out_path}_depth.mp4") + + print(f"Video saved to {out_path} (mode: {save_mode})") + + if effects is not None: + try: + gs_renderer.rasterizer.runner.splats = raw_splats + except: + pass + torch.cuda.empty_cache() diff --git a/worldcompass/reward_function/HunyuanWorldMirror/src/utils/save_utils.py b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/save_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2741ef3020de69bcc1e6b848d44acf0d291e347d --- /dev/null +++ b/worldcompass/reward_function/HunyuanWorldMirror/src/utils/save_utils.py @@ -0,0 +1,210 @@ +"""Utilities for saving images, depths, normals, point clouds, and Gaussian splat data. + +tencent +""" + +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from plyfile import PlyData, PlyElement +from io import BytesIO +import json +import os + + +def save_camera_params(extrinsics, intrinsics, target_dir): + """Save camera parameters (extrinsics and intrinsics) in JSON format. + + Args: + extrinsics: numpy array, shape [N, 4, 4] - extrinsic matrices for N cameras + intrinsics: numpy array, shape [N, 3, 3] - intrinsic matrices for N cameras + target_dir: str - directory to save the parameters + + Returns: + str: path to the saved file + """ + camera_data = { + "num_cameras": int(extrinsics.shape[0]), + "extrinsics": [], + "intrinsics": [], + } + + # Convert each camera's parameters to list format + for i in range(extrinsics.shape[0]): + camera_data["extrinsics"].append( + { + "camera_id": i, + "matrix": extrinsics[i].tolist(), # [4, 4] -> list + } + ) + camera_data["intrinsics"].append( + { + "camera_id": i, + "matrix": intrinsics[i].tolist(), # [3, 3] -> list + } + ) + + # Save as JSON file + camera_params_path = os.path.join(target_dir, "camera_params.json") + with open(camera_params_path, "w") as f: + json.dump(camera_data, f, indent=2) + + return camera_params_path + + +def save_image_png(path: Path, image_tensor: torch.Tensor) -> None: + # image_tensor: [H, W, 3] + img = (image_tensor.detach().cpu() * 255.0).to(torch.uint8).numpy() + Image.fromarray(img).save(str(path)) + + +def save_depth_png(path: Path, depth_tensor: torch.Tensor) -> None: + # depth_tensor: [H, W] + d = depth_tensor.detach() + d = d - d.min() + d = d / (d.max() + 1e-9) + img = (d.clamp(0, 1) * 255.0).to(torch.uint8).cpu().numpy() + Image.fromarray(img, mode="L").save(str(path)) + + +def save_depth_npy(path: Path, depth_tensor: torch.Tensor) -> None: + # depth_tensor: [H, W] + # Save actual depth values in numpy format + d = depth_tensor.detach().cpu().numpy() + np.save(str(path), d) + + +def save_normal_png(path: Path, normal_hwc: torch.Tensor) -> None: + # normal_hwc: [H, W, 3], in [-1, 1] + n = (normal_hwc.detach().cpu() + 1.0) * 0.5 + img = (n.clamp(0, 1) * 255.0).to(torch.uint8).numpy() + Image.fromarray(img).save(str(path)) + + +def save_scene_ply( + path: Path, + points_xyz: torch.Tensor, + point_colors: torch.Tensor, + valid_mask: torch.Tensor = None, +) -> None: + """Save point cloud to PLY format.""" + pts = points_xyz.detach().cpu().to(torch.float32).numpy().reshape(-1, 3) + colors = point_colors.detach().cpu().to(torch.uint8).numpy().reshape(-1, 3) + + # Filter out invalid points (NaN, Inf) + if valid_mask is None: + valid_mask = np.isfinite(pts).all(axis=1) + else: + valid_mask = valid_mask.detach().cpu().numpy().reshape(-1) + pts = pts[valid_mask] + colors = colors[valid_mask] + + # Handle empty point cloud + if len(pts) == 0: + pts = np.array([[0, 0, 0]], dtype=np.float32) + colors = np.array([[255, 255, 255]], dtype=np.uint8) + + # Create PLY data + vertex_dtype = [ + ("x", "f4"), + ("y", "f4"), + ("z", "f4"), + ("red", "u1"), + ("green", "u1"), + ("blue", "u1"), + ] + vertex_elements = np.empty(len(pts), dtype=vertex_dtype) + vertex_elements["x"] = pts[:, 0] + vertex_elements["y"] = pts[:, 1] + vertex_elements["z"] = pts[:, 2] + vertex_elements["red"] = colors[:, 0] + vertex_elements["green"] = colors[:, 1] + vertex_elements["blue"] = colors[:, 2] + + # Write PLY file + PlyData([PlyElement.describe(vertex_elements, "vertex")]).write(str(path)) + + +def save_points_ply( + path: Path, pts_np: np.ndarray, cols_np: np.ndarray +) -> None: + """Save point cloud to PLY format from numpy arrays.""" + vertex_dtype = [ + ("x", "f4"), + ("y", "f4"), + ("z", "f4"), + ("red", "u1"), + ("green", "u1"), + ("blue", "u1"), + ] + vertex_elements = np.empty(len(pts_np), dtype=vertex_dtype) + vertex_elements["x"] = pts_np[:, 0] + vertex_elements["y"] = pts_np[:, 1] + vertex_elements["z"] = pts_np[:, 2] + vertex_elements["red"] = cols_np[:, 0] + vertex_elements["green"] = cols_np[:, 1] + vertex_elements["blue"] = cols_np[:, 2] + + # Write PLY file + PlyData([PlyElement.describe(vertex_elements, "vertex")]).write(str(path)) + + +def save_gs_ply( + path: Path, + means: torch.Tensor, + scales: torch.Tensor, + rotations: torch.Tensor, + rgbs: torch.Tensor, + opacities: torch.Tensor, +) -> None: + """Export Gaussian splat data to PLY format. + + Args: + path: Output PLY file path + means: Gaussian centers [N, 3] + scales: Gaussian scales [N, 3] + rotations: Gaussian rotations as quaternions [N, 4] + rgbs: RGB colors [N, 3] + opacities: Opacity values [N] + """ + # Filter out points with scales greater than the 95th percentile + scale_threshold = torch.quantile(scales.max(dim=-1)[0], 0.95, dim=0) + filter_mask = scales.max(dim=-1)[0] <= scale_threshold + + # Apply the filter to all tensors + means = means[filter_mask].reshape(-1, 3) + scales = scales[filter_mask].reshape(-1, 3) + rotations = rotations[filter_mask].reshape(-1, 4) + rgbs = rgbs[filter_mask].reshape(-1, 3) + opacities = opacities[filter_mask].reshape(-1) + + # Construct attribute names + attributes = ["x", "y", "z", "nx", "ny", "nz"] + for i in range(3): + attributes.append(f"f_dc_{i}") + attributes.append("opacity") + for i in range(3): + attributes.append(f"scale_{i}") + for i in range(4): + attributes.append(f"rot_{i}") + + # Prepare PLY data structure + dtype_full = [(attribute, "f4") for attribute in attributes] + elements = np.empty(means.shape[0], dtype=dtype_full) + + # Concatenate all attributes + attributes_data = ( + means.float().detach().cpu().numpy(), + torch.zeros_like(means).float().detach().cpu().numpy(), + rgbs.detach().cpu().contiguous().numpy(), + opacities[..., None].detach().cpu().numpy(), + scales.log().detach().cpu().numpy(), + rotations.detach().cpu().numpy(), + ) + attributes_data = np.concatenate(attributes_data, axis=1) + elements[:] = list(map(tuple, attributes_data)) + + # Write to PLY file + PlyData([PlyElement.describe(elements, "vertex")]).write(str(path))