From 23316e89e0e257063e5cde7fecb182dcd4c14bf9 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 13 Jun 2026 04:31:55 -0400 Subject: [PATCH 1/3] fix(calib): use sc_matmul in MLP chunked branch _run_mlp_linear_level's chunked path called sc_matmul_enable_triton_mlp, a stale name that no longer exists (NameError at runtime), with obsolete max/min positional args. Switch to sc_matmul(granularity='per_row', ...) to match the attention chunked branch and the non-chunked MLP branch. Unblocks MP threshold calibration in this repo. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/calibrate_mp_thresholds.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/scripts/calibrate_mp_thresholds.py b/scripts/calibrate_mp_thresholds.py index e43a7e8..93321d2 100644 --- a/scripts/calibrate_mp_thresholds.py +++ b/scripts/calibrate_mp_thresholds.py @@ -521,13 +521,10 @@ def _run_mlp_linear_level( x_chunk = x_flat[:, start:end].contiguous() w_chunk = weight[:, start:end].contiguous() config = module._get_sc_config(end - start, sc_prec) - chunk_result = sc_matmul_enable_triton_mlp( + chunk_result = sc_matmul( x_chunk, w_chunk, - x_chunk.max().item(), - x_chunk.min().item(), - w_chunk.max().item(), - w_chunk.min().item(), + granularity="per_row", mode=module.sc_mode, sc_prec=sc_prec, config=config, From 598750bd168b3e7a6792a7f24b4a946dc3a6d112 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 13 Jun 2026 04:31:55 -0400 Subject: [PATCH 2/3] feat(eval): KID-on-pool3 + 9-GPU FP/W8A8 runners + MP calib worker - eval/kid_openai.{py,sh}: KID (poly-kernel unbiased MMD^2, 100x1000 subsets) on the SAME Inception pool3 features as the ADM evaluator FID, so KID is directly comparable to reported FID/sFID/IS; re-derives FID as a check. - eval_openai.sh: ADM FID/sFID/IS/Precision/Recall wrapper (pack PNGs->npz, run evaluator in tfeval env vs the ImageNet-256 reference). - run_quant_9gpu.sh + _quant_node_worker.sh: generalized no-SC baseline runner, 3 nodes x 3 GPU. wbits=16 -> FP16 baseline; wbits=8 -> W8A8 (no SC). - run_fp_9gpu.sh + _fp_node_worker.sh: FP16 baseline 9-GPU runner. - _mp_calib_worker.sh: build fresh MP threshold tables (avg 64/96/128). Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/_fp_node_worker.sh | 46 +++++++++++++ scripts/_mp_calib_worker.sh | 40 +++++++++++ scripts/_quant_node_worker.sh | 45 +++++++++++++ scripts/eval/kid_openai.py | 121 ++++++++++++++++++++++++++++++++++ scripts/eval/kid_openai.sh | 35 ++++++++++ scripts/eval_openai.sh | 54 +++++++++++++++ scripts/run_fp_9gpu.sh | 71 ++++++++++++++++++++ scripts/run_quant_9gpu.sh | 69 +++++++++++++++++++ 8 files changed, 481 insertions(+) create mode 100644 scripts/_fp_node_worker.sh create mode 100644 scripts/_mp_calib_worker.sh create mode 100644 scripts/_quant_node_worker.sh create mode 100644 scripts/eval/kid_openai.py create mode 100644 scripts/eval/kid_openai.sh create mode 100644 scripts/eval_openai.sh create mode 100644 scripts/run_fp_9gpu.sh create mode 100644 scripts/run_quant_9gpu.sh diff --git a/scripts/_fp_node_worker.sh b/scripts/_fp_node_worker.sh new file mode 100644 index 0000000..3b6ee4f --- /dev/null +++ b/scripts/_fp_node_worker.sh @@ -0,0 +1,46 @@ +#!/bin/bash +# Runs on ONE node (inside an salloc step) for the FP baseline. Launches 3 FP +# workers pinned to local GPUs 0/1/2 on index shards BASE_SHARD+{0,1,2}, all +# writing to one shared samples dir. FP = --wbits 16 --abits 16, no SC flags +# (SC layerwise fractions default to 0 -> pure FP16 DiT-XL/2). +set -uo pipefail +IDX_DIR="$1"; BASE_SHARD="$2"; SAMPLES="$3"; LOGDIR="$4"; CKPT="$5" +CFG="$6"; STEPS="$7"; BATCH="$8"; SEED="$9"; NUM_FID="${10}"; NUM_CLASSES="${11}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +conda activate qdit +export PYTHONUNBUFFERED=1 +cd "$REPO" + +pids=() +for g in 0 1 2; do + shard=$((BASE_SHARD + g)) + IDXF="$IDX_DIR/gpu_${shard}.txt" + if [[ ! -s "$IDXF" ]]; then + echo "[$(hostname) gpu$g] shard $shard empty -> skip" + continue + fi + mkdir -p "$LOGDIR/gpu_${shard}" + echo "[$(hostname) gpu$g] shard $shard: $(wc -l < "$IDXF") indices -> log gpu_${shard}.log" + CUDA_VISIBLE_DEVICES=$g python -u scripts/quant_sc_main.py \ + --wbits 16 --abits 16 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps "$STEPS" --cfg-scale "$CFG" \ + --batch-size "$BATCH" \ + --generate-fid-samples --balanced_classes \ + --num-classes "$NUM_CLASSES" \ + --balanced_total_samples "$NUM_FID" \ + --num-fid-samples "$NUM_FID" \ + --target_indices_path "$IDXF" \ + --samples_dir_override "$SAMPLES" \ + --seed "$SEED" \ + --results-dir "$LOGDIR/gpu_${shard}" \ + --ckpt "$CKPT" \ + > "$LOGDIR/gpu_${shard}.log" 2>&1 & + pids+=($!) +done + +rc=0 +for p in "${pids[@]}"; do wait "$p" || rc=1; done +echo "[$(hostname)] node_worker done (base=$BASE_SHARD) rc=$rc" +exit $rc diff --git a/scripts/_mp_calib_worker.sh b/scripts/_mp_calib_worker.sh new file mode 100644 index 0000000..d86495f --- /dev/null +++ b/scripts/_mp_calib_worker.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Build fresh MP calibration threshold tables for avg_stoc_len = 64/96/128, +# in the CURRENT repo (not reusing scmp_diffusion_prev's April tables). +# Runs 3 calibrations pinned to local GPUs 0/1/2. bitrev, cosine/FP-teacher, +# mp_levels 256..16, budget_ratio = target/256. Mirrors calib_sweep_targets_bitrev.sh. +# Args: +set -uo pipefail +CKPT="$1"; OUTDIR="$2" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev +cd "$REPO" +mkdir -p "$OUTDIR" + +declare -a T=( "64:0.25" "96:0.375" "128:0.5" ) +pids=(); g=0 +for e in "${T[@]}"; do + TGT=${e%:*}; BR=${e#*:} + echo "[gpu$g] calibrate avg=$TGT budget_ratio=$BR" + CUDA_VISIBLE_DEVICES=$g python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 6 \ + --timestep_buckets 4 --layer_buckets 4 \ + --teacher_cfg_scale 0.0 \ + --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${TGT}_l256_ref192.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${TGT}_summary.csv" \ + > "$OUTDIR/calib_avg${TGT}.log" 2>&1 & + pids+=($!); g=$((g+1)) +done +rc=0 +for p in "${pids[@]}"; do wait "$p" || rc=1; done +echo "calib worker done rc=$rc" +exit $rc diff --git a/scripts/_quant_node_worker.sh b/scripts/_quant_node_worker.sh new file mode 100644 index 0000000..c7cb1ce --- /dev/null +++ b/scripts/_quant_node_worker.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# Generalized per-node worker: 3 GPUs on index shards BASE_SHARD+{0,1,2}, all +# writing to one shared samples dir. NO stochastic computing (SC layerwise +# fractions default to 0). WBITS/ABITS control quantization: +# 16/16 -> pure FP16 ; 8/8 -> W8A8 fake-quant (exact matmuls, no SC). +set -uo pipefail +IDX_DIR="$1"; BASE_SHARD="$2"; SAMPLES="$3"; LOGDIR="$4"; CKPT="$5" +CFG="$6"; STEPS="$7"; BATCH="$8"; SEED="$9"; NUM_FID="${10}"; NUM_CLASSES="${11}" +WBITS="${12}"; ABITS="${13}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +conda activate qdit +export PYTHONUNBUFFERED=1 +cd "$REPO" + +pids=() +for g in 0 1 2; do + shard=$((BASE_SHARD + g)) + IDXF="$IDX_DIR/gpu_${shard}.txt" + if [[ ! -s "$IDXF" ]]; then + echo "[$(hostname) gpu$g] shard $shard empty -> skip"; continue + fi + mkdir -p "$LOGDIR/gpu_${shard}" + echo "[$(hostname) gpu$g] shard $shard: $(wc -l < "$IDXF") indices (w${WBITS}a${ABITS}, no SC)" + CUDA_VISIBLE_DEVICES=$g python -u scripts/quant_sc_main.py \ + --wbits "$WBITS" --abits "$ABITS" --w_sym --a_sym \ + --image-size 256 --num-sampling-steps "$STEPS" --cfg-scale "$CFG" \ + --batch-size "$BATCH" \ + --generate-fid-samples --balanced_classes \ + --num-classes "$NUM_CLASSES" \ + --balanced_total_samples "$NUM_FID" \ + --num-fid-samples "$NUM_FID" \ + --target_indices_path "$IDXF" \ + --samples_dir_override "$SAMPLES" \ + --seed "$SEED" \ + --results-dir "$LOGDIR/gpu_${shard}" \ + --ckpt "$CKPT" \ + > "$LOGDIR/gpu_${shard}.log" 2>&1 & + pids+=($!) +done +rc=0 +for p in "${pids[@]}"; do wait "$p" || rc=1; done +echo "[$(hostname)] node_worker done (base=$BASE_SHARD, w${WBITS}a${ABITS}) rc=$rc" +exit $rc diff --git a/scripts/eval/kid_openai.py b/scripts/eval/kid_openai.py new file mode 100644 index 0000000..86306f0 --- /dev/null +++ b/scripts/eval/kid_openai.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +""" +KID (Kernel Inception Distance) on the SAME Inception-V3 pool3 features that the +OpenAI/ADM evaluator.py uses for FID -- so KID is directly comparable to the +FID / sFID / IS / Precision / Recall that evaluator.py already reports. + +Metric: polynomial-kernel (degree 3, coef0 1, gamma 1/d) UNBIASED MMD^2, +averaged over `n_subsets` random subsets of size `subset_size` +(torch-fidelity / clean-fid default: 100 x 1000). Reported as mean +/- std. + +As a consistency check we also recompute FID from the same pool3 features; it +should match evaluator.py's FID to within float noise, proving the KID is on +the same feature space as the reported FID. + +Usage: + python kid_openai.py [ ...] +Env override: + EVALUATOR=/path/to/evaluator.py (default: scmp_diffusion_prev ADM evaluator) +""" +import os +import sys + +import numpy as np +from scipy import linalg + +EVALUATOR_PY = os.environ.get( + "EVALUATOR", + "/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/" + "scmp_diffusion_prev/Q-DiT/models/evaluations/evaluator.py", +) +sys.path.insert(0, os.path.dirname(EVALUATOR_PY)) + +import tensorflow.compat.v1 as tf # noqa: E402 +from evaluator import Evaluator # noqa: E402 + +SUBSET_SIZE = 1000 +N_SUBSETS = 100 +SEED = 0 + + +def polynomial_kernel(x, y): + d = x.shape[1] + return (x.astype(np.float64) @ y.astype(np.float64).T / d + 1.0) ** 3 + + +def mmd2_unbiased(k_xx, k_xy, k_yy): + m = k_xx.shape[0] + n = k_yy.shape[0] + sum_xx = k_xx.sum() - np.trace(k_xx) + sum_yy = k_yy.sum() - np.trace(k_yy) + sum_xy = k_xy.sum() + return sum_xx / (m * (m - 1)) + sum_yy / (n * (n - 1)) - 2.0 * sum_xy / (m * n) + + +def compute_kid(ref, smp, subset_size=SUBSET_SIZE, n_subsets=N_SUBSETS, seed=SEED): + rng = np.random.RandomState(seed) + m = min(subset_size, ref.shape[0], smp.shape[0]) + vals = np.empty(n_subsets, dtype=np.float64) + for i in range(n_subsets): + x = ref[rng.choice(ref.shape[0], m, replace=False)] + y = smp[rng.choice(smp.shape[0], m, replace=False)] + vals[i] = mmd2_unbiased( + polynomial_kernel(x, x), + polynomial_kernel(x, y), + polynomial_kernel(y, y), + ) + return float(vals.mean()), float(vals.std()), m, n_subsets + + +def fid_from_pool3(ref, smp): + mu1, mu2 = ref.mean(0), smp.mean(0) + s1 = np.cov(ref, rowvar=False) + s2 = np.cov(smp, rowvar=False) + diff = mu1 - mu2 + covmean, _ = linalg.sqrtm(s1.dot(s2), disp=False) + if np.iscomplexobj(covmean): + covmean = covmean.real + return float(diff.dot(diff) + np.trace(s1) + np.trace(s2) - 2 * np.trace(covmean)) + + +def main(): + if len(sys.argv) < 4: + sys.exit("usage: kid_openai.py [...]") + ref_npz, out_txt, sample_npzs = sys.argv[1], sys.argv[2], sys.argv[3:] + + config = tf.ConfigProto(allow_soft_placement=True) + config.gpu_options.allow_growth = True + ev = Evaluator(tf.Session(config=config)) + print("warming up TensorFlow...", flush=True) + ev.warmup() + + print(f"computing reference pool3 activations: {ref_npz}", flush=True) + ref_pool3 = ev.read_activations(ref_npz)[0] + print(f" ref pool3 {ref_pool3.shape}", flush=True) + + lines = [ + f"# KID = poly-kernel(deg3, gamma=1/d, coef0=1) unbiased MMD^2, " + f"{N_SUBSETS}x{SUBSET_SIZE} subsets, same pool3 features as evaluator.py FID", + f"# ref={ref_npz} n_ref={ref_pool3.shape[0]}", + ] + for npz in sample_npzs: + print(f"computing sample pool3 activations: {npz}", flush=True) + smp_pool3 = ev.read_activations(npz)[0] + kid_mean, kid_std, m, ns = compute_kid(ref_pool3, smp_pool3) + fid_chk = fid_from_pool3(ref_pool3, smp_pool3) + line = ( + f"{os.path.basename(npz)}\t" + f"KID {kid_mean:.6e} +/- {kid_std:.2e}\t" + f"KIDx1e3 {kid_mean * 1e3:.4f}\t" + f"[{ns}x{m}]\tFID_check {fid_chk:.4f}\tn {smp_pool3.shape[0]}" + ) + print(line, flush=True) + lines.append(line) + + with open(out_txt, "w") as f: + f.write("\n".join(lines) + "\n") + print(f"[ok] wrote {out_txt}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/kid_openai.sh b/scripts/eval/kid_openai.sh new file mode 100644 index 0000000..2386f05 --- /dev/null +++ b/scripts/eval/kid_openai.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# ============================================================ +# Compute KID on the SAME pool3 features as the OpenAI/ADM evaluator FID, +# reusing the scmp_diffusion_prev evaluator + ImageNet-256 reference, in the +# `tfeval` conda env (TF 2.15 + bundled CUDA libs). Companion to eval_openai.sh. +# +# Usage: +# bash scripts/eval/kid_openai.sh [ ...] +# Overrides: PREV, REF_NPZ, EVALUATOR +# ============================================================ +set -euo pipefail + +OUT_TXT="${1:?usage: kid_openai.sh [...]}" +shift +[[ $# -ge 1 ]] || { echo "ERROR: need >=1 sample npz" >&2; exit 1; } + +PREV="${PREV:-/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev}" +REF_NPZ="${REF_NPZ:-${PREV}/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz}" +export EVALUATOR="${EVALUATOR:-${PREV}/Q-DiT/models/evaluations/evaluator.py}" +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +for f in "$EVALUATOR" "$REF_NPZ" "$HERE/kid_openai.py"; do + [[ -e "$f" ]] || { echo "ERROR: missing $f" >&2; exit 1; } +done + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +conda activate tfeval +NV=$(python -c "import os, nvidia; print(os.path.dirname(nvidia.__file__))") +export LD_LIBRARY_PATH=${NV}/cudnn/lib:${NV}/cuda_runtime/lib:${NV}/cuda_cupti/lib:${NV}/cuda_nvrtc/lib:${NV}/cublas/lib:${NV}/cufft/lib:${NV}/curand/lib:${NV}/cusolver/lib:${NV}/cusparse/lib:${NV}/nvjitlink/lib:${LD_LIBRARY_PATH:-} +export TF_CPP_MIN_LOG_LEVEL=2 + +echo "=== KID: ref=${REF_NPZ} ===" +python -u "$HERE/kid_openai.py" "$REF_NPZ" "$OUT_TXT" "$@" +echo "=== done: ${OUT_TXT} ===" +cat "$OUT_TXT" diff --git a/scripts/eval_openai.sh b/scripts/eval_openai.sh new file mode 100644 index 0000000..a5b15ad --- /dev/null +++ b/scripts/eval_openai.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# ============================================================ +# Robust ADM/OpenAI evaluation of a sample directory, reusing the +# scmp_diffusion_prev evaluator + ImageNet-256 reference. +# +# Packs a PNG dir -> samples.npz, then runs the OpenAI evaluator +# (Inception Score / FID / sFID / Precision / Recall) in the `tfeval` +# conda env against VIRTUAL_imagenet256_labeled.npz. +# +# Usage: +# bash scripts/eval_openai.sh +# dir of NNNNNN.png (256x256) +# writes .npz and .openai_eval.txt +# +# Override: +# PREV=/gpfs/.../scmp_diffusion_prev (evaluator + reference live here) +# ============================================================ +set -euo pipefail + +SAMPLES="${1:?usage: eval_openai.sh }" +OUT_PREFIX="${2:?usage: eval_openai.sh }" + +PREV="${PREV:-/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev}" +EVALUATOR="${EVALUATOR:-${PREV}/Q-DiT/models/evaluations/evaluator.py}" +REF_NPZ="${REF_NPZ:-${PREV}/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz}" +PACKER="${PACKER:-${PREV}/imagenet256_ref/parallel_npz.py}" + +for f in "$EVALUATOR" "$REF_NPZ" "$PACKER"; do + [[ -e "$f" ]] || { echo "ERROR: missing $f" >&2; exit 1; } +done +[[ -d "$SAMPLES" ]] || { echo "ERROR: samples dir not found: $SAMPLES" >&2; exit 1; } + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh + +NPZ="${OUT_PREFIX}.npz" +EVAL_TXT="${OUT_PREFIX}.openai_eval.txt" + +# 1) pack PNGs -> npz (qdit env has numpy/PIL/tqdm) +echo "=== packing $(basename "$SAMPLES") -> ${NPZ} ===" +conda activate qdit +python -u "$PACKER" "$SAMPLES" "$NPZ" +conda deactivate + +# 2) OpenAI ADM evaluator in tfeval (TF 2.15 + CUDA libs on LD_LIBRARY_PATH) +echo "=== evaluating ${NPZ} vs reference ===" +conda activate tfeval +NV=$(python -c "import os, nvidia; print(os.path.dirname(nvidia.__file__))") +export LD_LIBRARY_PATH=${NV}/cudnn/lib:${NV}/cuda_runtime/lib:${NV}/cuda_cupti/lib:${NV}/cuda_nvrtc/lib:${NV}/cublas/lib:${NV}/cufft/lib:${NV}/curand/lib:${NV}/cusolver/lib:${NV}/cusparse/lib:${NV}/nvjitlink/lib:${LD_LIBRARY_PATH:-} +export TF_CPP_MIN_LOG_LEVEL=2 +python -u "$EVALUATOR" "$REF_NPZ" "$NPZ" 2>&1 | tee "$EVAL_TXT" +conda deactivate + +echo "=== done: ${EVAL_TXT} ===" +grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "$EVAL_TXT" || true diff --git a/scripts/run_fp_9gpu.sh b/scripts/run_fp_9gpu.sh new file mode 100644 index 0000000..101be26 --- /dev/null +++ b/scripts/run_fp_9gpu.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# ============================================================ +# FP (floating-point, no-SC, no-quant) DiT-XL/2 baseline across all 9 held +# GPUs (3 nodes x 3 GPUs). 9-way class-balanced index shard -> generate -> +# eval (FID/sFID/IS/Prec/Rec via OpenAI ADM evaluator) + KID. +# +# Distributional baseline: same seed-0 RNG + same balanced 1000-class layout as +# the SC sweeps, but a fresh 9-way split (so NOT per-index noise-paired with the +# 3-way SC runs; FID/KID are distributional so this is fine). +# +# Override: CFG STEPS BATCH SEED NUM_FID NUM_CLASSES OUT_BASE JOBS +# ============================================================ +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT="${CKPT:-$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt}" + +CFG="${CFG:-1.5}"; STEPS="${STEPS:-50}"; BATCH="${BATCH:-64}" +SEED="${SEED:-0}"; NUM_FID="${NUM_FID:-10000}"; NUM_CLASSES="${NUM_CLASSES:-1000}" +OUT_BASE="${OUT_BASE:-$SCRATCH/scmp_diffusion_fid_fp_cfg${CFG/./}}" +read -ra JOBS <<< "${JOBS:-51660970 51660978 51660979}" # gl1802 gl1804 gl1805 + +SAMPLES="$OUT_BASE/fp/samples"; IDX_DIR="$OUT_BASE/fp/_indices"; LOGDIR="$OUT_BASE/fp/_logs" +mkdir -p "$SAMPLES" "$IDX_DIR" "$LOGDIR" +ALIGN=$((NUM_FID / NUM_CLASSES)) + +echo "============================================================" +echo "FP baseline CFG=$CFG STEPS=$STEPS BATCH=$BATCH SEED=$SEED NUM_FID=$NUM_FID" +echo " OUT_BASE=$OUT_BASE" +echo " JOBS=${JOBS[*]} (9 GPUs) CKPT=$CKPT" +echo "============================================================" +[[ -f "$CKPT" ]] || { echo "ERROR: missing CKPT $CKPT" >&2; exit 1; } + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +echo "=== plan 9 shards (align=$ALIGN) ===" +python -u "$REPO/scripts/_plan_missing_indices.py" "$SAMPLES" "$NUM_FID" 9 "$IDX_DIR" "$ALIGN" + +echo "=== dispatch 3 nodes x 3 GPUs at $(date +%H:%M:%S) ===" +spids=() +for n in 0 1 2; do + base=$((3 * n)) + srun --jobid="${JOBS[$n]}" --overlap --gres=gpu:3 \ + bash "$REPO/scripts/_fp_node_worker.sh" \ + "$IDX_DIR" "$base" "$SAMPLES" "$LOGDIR" "$CKPT" \ + "$CFG" "$STEPS" "$BATCH" "$SEED" "$NUM_FID" "$NUM_CLASSES" \ + > "$LOGDIR/node_${n}.log" 2>&1 & + spids+=($!) + echo " node $n (job ${JOBS[$n]}) shards $base..$((base+2)) dispatched" +done +grc=0 +for p in "${spids[@]}"; do wait "$p" || grc=1; done +FINAL=$(find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l) +echo "=== generation done rc=$grc : $FINAL/$NUM_FID samples at $(date +%H:%M:%S) ===" +if [[ "$FINAL" -lt "$NUM_FID" ]]; then + echo "[partial] re-run this script to resume (idempotent on existing PNGs)"; exit 1 +fi + +echo "=== eval: FID/sFID/IS/Precision/Recall ===" +srun --jobid="${JOBS[0]}" --overlap --gres=gpu:1 \ + bash "$REPO/scripts/eval_openai.sh" "$SAMPLES" "$OUT_BASE/fp" \ + > "$LOGDIR/eval_openai.log" 2>&1 +echo "=== eval: KID (same pool3 features) ===" +srun --jobid="${JOBS[0]}" --overlap --gres=gpu:1 \ + bash "$REPO/scripts/eval/kid_openai.sh" "$OUT_BASE/fp.kid.txt" "$OUT_BASE/fp.npz" \ + > "$LOGDIR/eval_kid.log" 2>&1 + +echo "============================================================" +echo "FP BASELINE RESULTS (CFG=$CFG)" +grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "$OUT_BASE/fp.openai_eval.txt" 2>/dev/null +cat "$OUT_BASE/fp.kid.txt" 2>/dev/null +echo "DONE $(date)" diff --git a/scripts/run_quant_9gpu.sh b/scripts/run_quant_9gpu.sh new file mode 100644 index 0000000..949babb --- /dev/null +++ b/scripts/run_quant_9gpu.sh @@ -0,0 +1,69 @@ +#!/bin/bash +# ============================================================ +# Generalized 9-GPU (3 nodes x 3 GPU) no-SC baseline runner. +# WBITS=16 ABITS=16 -> FP16 baseline +# WBITS=8 ABITS=8 -> W8A8 quantized, NO stochastic computing +# Same seed-0 RNG + balanced 1000-class layout as the SC sweeps (distributional +# baseline; fresh 9-way split, not per-index noise-paired). Generate -> eval +# (FID/sFID/IS/Prec/Rec) + KID. +# +# Override: WBITS ABITS CFG STEPS BATCH SEED NUM_FID NUM_CLASSES TAG OUT_BASE JOBS +# ============================================================ +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT="${CKPT:-$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt}" + +WBITS="${WBITS:-8}"; ABITS="${ABITS:-8}" +CFG="${CFG:-1.5}"; STEPS="${STEPS:-50}"; BATCH="${BATCH:-64}" +SEED="${SEED:-0}"; NUM_FID="${NUM_FID:-10000}"; NUM_CLASSES="${NUM_CLASSES:-1000}" +TAG="${TAG:-w${WBITS}a${ABITS}nosc}" +OUT_BASE="${OUT_BASE:-$SCRATCH/scmp_diffusion_fid_${TAG}_cfg${CFG/./}}" +read -ra JOBS <<< "${JOBS:-51660970 51660978 51660979}" + +SAMPLES="$OUT_BASE/samples"; IDX_DIR="$OUT_BASE/_indices"; LOGDIR="$OUT_BASE/_logs" +mkdir -p "$SAMPLES" "$IDX_DIR" "$LOGDIR" +ALIGN=$((NUM_FID / NUM_CLASSES)) + +echo "============================================================" +echo "$TAG baseline w${WBITS}a${ABITS} NO-SC CFG=$CFG STEPS=$STEPS BATCH=$BATCH SEED=$SEED NUM_FID=$NUM_FID" +echo " OUT_BASE=$OUT_BASE JOBS=${JOBS[*]} (9 GPUs)" +echo "============================================================" +[[ -f "$CKPT" ]] || { echo "ERROR: missing CKPT $CKPT" >&2; exit 1; } + +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +echo "=== plan 9 shards (align=$ALIGN) ===" +python -u "$REPO/scripts/_plan_missing_indices.py" "$SAMPLES" "$NUM_FID" 9 "$IDX_DIR" "$ALIGN" + +echo "=== dispatch 3 nodes x 3 GPUs at $(date +%H:%M:%S) ===" +spids=() +for n in 0 1 2; do + base=$((3 * n)) + srun --jobid="${JOBS[$n]}" --overlap --gres=gpu:3 \ + bash "$REPO/scripts/_quant_node_worker.sh" \ + "$IDX_DIR" "$base" "$SAMPLES" "$LOGDIR" "$CKPT" \ + "$CFG" "$STEPS" "$BATCH" "$SEED" "$NUM_FID" "$NUM_CLASSES" "$WBITS" "$ABITS" \ + > "$LOGDIR/node_${n}.log" 2>&1 & + spids+=($!) + echo " node $n (job ${JOBS[$n]}) shards $base..$((base+2)) dispatched" +done +grc=0 +for p in "${spids[@]}"; do wait "$p" || grc=1; done +FINAL=$(find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l) +echo "=== generation done rc=$grc : $FINAL/$NUM_FID at $(date +%H:%M:%S) ===" +if [[ "$FINAL" -lt "$NUM_FID" ]]; then echo "[partial] re-run to resume"; exit 1; fi + +echo "=== eval: FID/sFID/IS/Precision/Recall ===" +srun --jobid="${JOBS[0]}" --overlap --gres=gpu:1 \ + bash "$REPO/scripts/eval_openai.sh" "$SAMPLES" "$OUT_BASE/${TAG}" \ + > "$LOGDIR/eval_openai.log" 2>&1 +echo "=== eval: KID ===" +srun --jobid="${JOBS[0]}" --overlap --gres=gpu:1 \ + bash "$REPO/scripts/eval/kid_openai.sh" "$OUT_BASE/${TAG}.kid.txt" "$OUT_BASE/${TAG}.npz" \ + > "$LOGDIR/eval_kid.log" 2>&1 + +echo "============================================================" +echo "$TAG RESULTS (w${WBITS}a${ABITS} no-SC, CFG=$CFG)" +grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "$OUT_BASE/${TAG}.openai_eval.txt" 2>/dev/null +cat "$OUT_BASE/${TAG}.kid.txt" 2>/dev/null +echo "DONE $(date)" From b5a4ac4778619d6a520f503897adcfc624890d8a Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 13 Jun 2026 04:36:27 -0400 Subject: [PATCH 3/3] feat(eval): add uniform SC FID sweep runners - run_ngpu_uniform_nohalve_fanout.sh: N-GPU uniform-precision SC FID sweep with work-steal index planning (the cfg=1.5 uniform sl=64/96/128 results came from this). NUM_GPUS env-overridable. - run_3gpu_uniform_nohalve_bitrev.sh / run_3gpu_uniform_halve_bitrev.sh: 3-GPU uniform sweep variants (halve OFF / ON, bitrev scramble). SC configs live under results/ (gitignored, regenerable); FID numbers are in PR #3. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/run_3gpu_uniform_halve_bitrev.sh | 186 +++++++++++++++++++++ scripts/run_3gpu_uniform_nohalve_bitrev.sh | 138 +++++++++++++++ scripts/run_ngpu_uniform_nohalve_fanout.sh | 186 +++++++++++++++++++++ 3 files changed, 510 insertions(+) create mode 100644 scripts/run_3gpu_uniform_halve_bitrev.sh create mode 100755 scripts/run_3gpu_uniform_nohalve_bitrev.sh create mode 100755 scripts/run_ngpu_uniform_nohalve_fanout.sh diff --git a/scripts/run_3gpu_uniform_halve_bitrev.sh b/scripts/run_3gpu_uniform_halve_bitrev.sh new file mode 100644 index 0000000..f3652da --- /dev/null +++ b/scripts/run_3gpu_uniform_halve_bitrev.sh @@ -0,0 +1,186 @@ +#!/bin/bash +# ============================================================ +# 3-GPU uniform-precision FID sample sweep. +# +# sc_prec=8, halve ON (sweepable: explicit stoc_len, rng grid -> 128), +# bitrev scramble, qk & av granularity = per_row, all operators + all +# timesteps SC. Sweeps stoc_len in {256, 192, 128}; NUM_FID samples per +# config, fanned out across 3 GPUs with work-stealing (resumable). +# +# Mirrors scmp_diffusion_prev/Q-DiT/scripts/run_2gpu_calib_fid_sweep.sh. +# +# Usage (from a 3-GPU allocation): +# conda activate qdit +# cd .../scmp_diffusion +# bash scripts/run_3gpu_uniform_halve_bitrev.sh +# +# Override knobs (env vars): +# NUM_FID=10000 samples per config (must be divisible by NUM_CLASSES) +# STOC_LENS=256,192,128 +# BATCH=32 per-GPU batch +# NUM_STEPS=50 sampling steps +# CFG_SCALE=4 classifier-free guidance scale +# OWEN_MODE=bitrev counter | bitrev | random +# CKPT=/scratch/.../pretrained_models/DiT-XL-2-256x256.pt +# OUT_BASE=/scratch/.../scmp_diffusion_fid_halve_bitrev +# CONFIG_DIR=$OUT_BASE/configs (uniform${SL}_all.json live here) +# IMAGENET_REF=/path/to/val if set + pytorch_fid present -> compute FID +# ============================================================ + +set -euo pipefail + +NUM_GPUS="${NUM_GPUS:-3}" +NUM_FID="${NUM_FID:-10000}" +STOC_LENS="${STOC_LENS:-256,192,128}" +BATCH="${BATCH:-32}" +NUM_STEPS="${NUM_STEPS:-50}" +CFG_SCALE="${CFG_SCALE:-4}" +OWEN_MODE="${OWEN_MODE:-bitrev}" +NUM_CLASSES="${NUM_CLASSES:-1000}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi}" +CKPT="${CKPT:-${SCRATCH_ROOT}/pretrained_models/DiT-XL-2-256x256.pt}" +OUT_BASE="${OUT_BASE:-${SCRATCH_ROOT}/scmp_diffusion_fid_halve_bitrev}" +CONFIG_DIR="${CONFIG_DIR:-${OUT_BASE}/configs}" + +mkdir -p "${OUT_BASE}" +SWEEP_LOG="${OUT_BASE}/sweep.log" + +if (( NUM_FID % NUM_CLASSES != 0 )); then + echo "ERROR: NUM_FID (${NUM_FID}) must be divisible by NUM_CLASSES (${NUM_CLASSES})" >&2 + exit 1 +fi +if [[ ! -f "${CKPT}" ]]; then + echo "ERROR: checkpoint not found: ${CKPT}" >&2 + exit 1 +fi + +export SC_OWEN_MODE="${OWEN_MODE}" +export PYTHONUNBUFFERED=1 + +{ +echo "============================================================" +echo "3-GPU uniform halve+bitrev sweep started $(date)" +echo " NUM_GPUS=${NUM_GPUS} NUM_FID=${NUM_FID} STOC_LENS=${STOC_LENS}" +echo " BATCH=${BATCH} NUM_STEPS=${NUM_STEPS} CFG_SCALE=${CFG_SCALE}" +echo " OWEN_MODE=${OWEN_MODE} sc_prec=8 halve=ON qk/av=per_row" +echo " CKPT=${CKPT}" +echo " OUT_BASE=${OUT_BASE}" +echo "============================================================" +} | tee -a "${SWEEP_LOG}" + +cd "${REPO_ROOT}" + +run_config() { + local SL="$1" + local TAG="uniform${SL}" + local CFG_DIR="${OUT_BASE}/${TAG}" + local SAMPLES="${CFG_DIR}/samples" + local IDX_DIR="${CFG_DIR}/_indices" + local SC_JSON="${CONFIG_DIR}/sc_cfg_uniform${SL}_all.json" + mkdir -p "${SAMPLES}" "${IDX_DIR}" + + if [[ ! -f "${SC_JSON}" ]]; then + echo "[skip] ${TAG}: config not found ${SC_JSON}" | tee -a "${SWEEP_LOG}" + return + fi + + local DONE + DONE=$(find "${SAMPLES}" -maxdepth 1 -type f -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' 2>/dev/null | wc -l) + if [[ ${DONE} -ge ${NUM_FID} ]]; then + echo "[skip] ${TAG} already complete (${DONE}/${NUM_FID})" | tee -a "${SWEEP_LOG}" + return + fi + + echo "[plan] ${TAG} ${DONE}/${NUM_FID} done; planning across ${NUM_GPUS} GPUs at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + python -u "${REPO_ROOT}/scripts/_plan_missing_indices.py" \ + "${SAMPLES}" "${NUM_FID}" "${NUM_GPUS}" "${IDX_DIR}" | tee -a "${SWEEP_LOG}" + + local PIDS=() + for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do + local INDICES_FILE="${IDX_DIR}/gpu_${GPU_ID}.txt" + if [[ ! -s "${INDICES_FILE}" ]]; then + echo " [GPU ${GPU_ID}] no work assigned" | tee -a "${SWEEP_LOG}" + continue + fi + local GPU_LOG_DIR="${CFG_DIR}/_logs/gpu_${GPU_ID}" + mkdir -p "${GPU_LOG_DIR}" + CUDA_VISIBLE_DEVICES=${GPU_ID} \ + SC_OWEN_MODE="${OWEN_MODE}" \ + python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 \ + --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_halve \ + --sc_qk_granularity per_row \ + --sc_config "${SC_JSON}" \ + --image-size 256 --num-sampling-steps "${NUM_STEPS}" --cfg-scale "${CFG_SCALE}" \ + --batch-size "${BATCH}" \ + --generate-fid-samples \ + --balanced_classes \ + --num-classes "${NUM_CLASSES}" \ + --balanced_total_samples "${NUM_FID}" \ + --num-fid-samples "${NUM_FID}" \ + --target_indices_path "${INDICES_FILE}" \ + --samples_dir_override "${SAMPLES}" \ + --seed ${GPU_ID} \ + --results-dir "${GPU_LOG_DIR}" \ + --ckpt "${CKPT}" \ + > "${GPU_LOG_DIR}/run.log" 2>&1 & + PIDS+=($!) + sleep 5 + done + + local FAILED=0 + for i in "${!PIDS[@]}"; do + if wait "${PIDS[$i]}"; then + echo " [worker $i] done at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + else + echo " [worker $i] FAILED (rc=$?); see ${CFG_DIR}/_logs/" | tee -a "${SWEEP_LOG}" + FAILED=$((FAILED + 1)) + fi + done + + local FINAL + FINAL=$(find "${SAMPLES}" -maxdepth 1 -type f -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' 2>/dev/null | wc -l) + if [[ ${FINAL} -ge ${NUM_FID} ]]; then + echo "[ok] ${TAG} ${FINAL}/${NUM_FID} samples in ${SAMPLES} at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + else + echo "[partial] ${TAG} ${FINAL}/${NUM_FID}; rerun this script to resume" | tee -a "${SWEEP_LOG}" + fi +} + +IFS=',' read -ra SLS <<< "${STOC_LENS}" +for SL in "${SLS[@]}"; do + run_config "${SL}" +done + +echo "Sample generation finished $(date)" | tee -a "${SWEEP_LOG}" + +# --- robust ADM/OpenAI evaluation (reuses scmp_diffusion_prev evaluator + +# ImageNet-256 reference): Inception Score / FID / sFID / Precision / Recall. +# Set RUN_EVAL=0 to skip (e.g. generate now, evaluate later). +if [[ "${RUN_EVAL:-1}" == "1" ]]; then + for SL in "${SLS[@]}"; do + D="${OUT_BASE}/uniform${SL}/samples" + N=$(find "${D}" -maxdepth 1 -type f -name '*.png' 2>/dev/null | wc -l) + if [[ "${N}" -lt "${NUM_FID}" ]]; then + echo "[eval-skip] uniform${SL}: only ${N}/${NUM_FID} samples" | tee -a "${SWEEP_LOG}" + continue + fi + echo "[eval] uniform${SL} ($(date +%H:%M:%S)) ..." | tee -a "${SWEEP_LOG}" + bash "${REPO_ROOT}/scripts/eval_openai.sh" "${D}" "${OUT_BASE}/uniform${SL}" \ + 2>&1 | tee -a "${SWEEP_LOG}" + done + echo "=== metrics summary ===" | tee -a "${SWEEP_LOG}" + for SL in "${SLS[@]}"; do + T="${OUT_BASE}/uniform${SL}.openai_eval.txt" + [[ -f "${T}" ]] || continue + echo "uniform${SL}:" | tee -a "${SWEEP_LOG}" + grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "${T}" | sed 's/^/ /' | tee -a "${SWEEP_LOG}" + done +else + echo "Eval skipped (RUN_EVAL=0). Run scripts/eval_openai.sh per config later." | tee -a "${SWEEP_LOG}" +fi +echo "DONE $(date)" | tee -a "${SWEEP_LOG}" diff --git a/scripts/run_3gpu_uniform_nohalve_bitrev.sh b/scripts/run_3gpu_uniform_nohalve_bitrev.sh new file mode 100755 index 0000000..8b35aa2 --- /dev/null +++ b/scripts/run_3gpu_uniform_nohalve_bitrev.sh @@ -0,0 +1,138 @@ +#!/bin/bash +# ============================================================ +# 3-GPU parallel stoc_len sweep, halve=OFF, bitrev scramble. +# +# GPU 0 -> uniform128, GPU 1 -> uniform96, GPU 2 -> uniform64. +# Each GPU generates the full balanced NUM_FID set (10/class x 1000) +# for its stoc_len, concurrently. sc_prec=8, halve OFF, bitrev +# scramble, qk & av = per_row, all operators + all timesteps SC. +# Resumable (idempotent on existing PNGs). +# +# Usage (from a >=3-GPU allocation): +# conda activate qdit +# cd .../scmp_diffusion +# bash scripts/run_3gpu_uniform_nohalve_bitrev.sh +# +# Override: NUM_FID, STOC_LENS (must have <= #GPUs entries), BATCH, +# NUM_STEPS, CFG_SCALE, OWEN_MODE, NUM_CLASSES, CKPT, OUT_BASE, CONFIG_DIR, +# RUN_EVAL (1 to eval after generation). +# ============================================================ +set -euo pipefail + +NUM_FID="${NUM_FID:-10000}" +STOC_LENS="${STOC_LENS:-128,96,64}" +BATCH="${BATCH:-64}" +NUM_STEPS="${NUM_STEPS:-50}" +CFG_SCALE="${CFG_SCALE:-4}" +OWEN_MODE="${OWEN_MODE:-bitrev}" +NUM_CLASSES="${NUM_CLASSES:-1000}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi}" +CKPT="${CKPT:-${SCRATCH_ROOT}/pretrained_models/DiT-XL-2-256x256.pt}" +OUT_BASE="${OUT_BASE:-${SCRATCH_ROOT}/scmp_diffusion_fid_halve_bitrev}" +CONFIG_DIR="${CONFIG_DIR:-${OUT_BASE}/configs}" + +mkdir -p "${OUT_BASE}" +SWEEP_LOG="${OUT_BASE}/parallel_sweep.log" + +(( NUM_FID % NUM_CLASSES == 0 )) || { echo "ERROR: NUM_FID must be divisible by NUM_CLASSES" >&2; exit 1; } +[[ -f "${CKPT}" ]] || { echo "ERROR: checkpoint not found: ${CKPT}" >&2; exit 1; } + +IFS=',' read -ra SLS <<< "${STOC_LENS}" + +export SC_OWEN_MODE="${OWEN_MODE}" +export PYTHONUNBUFFERED=1 + +{ +echo "============================================================" +echo "3-GPU parallel stoc_len sweep started $(date)" +echo " one GPU per stoc_len: ${STOC_LENS}" +echo " NUM_FID=${NUM_FID} BATCH=${BATCH} NUM_STEPS=${NUM_STEPS} CFG_SCALE=${CFG_SCALE}" +echo " OWEN_MODE=${OWEN_MODE} sc_prec=8 halve=OFF qk/av=per_row" +echo " CKPT=${CKPT} OUT_BASE=${OUT_BASE}" +echo "============================================================" +} | tee -a "${SWEEP_LOG}" + +cd "${REPO_ROOT}" + +PIDS=() +for GPU_ID in "${!SLS[@]}"; do + SL="${SLS[$GPU_ID]}" + TAG="uniform${SL}" + CFG_DIR="${OUT_BASE}/${TAG}" + SAMPLES="${CFG_DIR}/samples" + IDX_DIR="${CFG_DIR}/_indices" + SC_JSON="${CONFIG_DIR}/sc_cfg_uniform${SL}_all.json" + GPU_LOG_DIR="${CFG_DIR}/_logs/gpu_${GPU_ID}" + mkdir -p "${SAMPLES}" "${IDX_DIR}" "${GPU_LOG_DIR}" + + [[ -f "${SC_JSON}" ]] || { echo "[skip] ${TAG}: config missing ${SC_JSON}" | tee -a "${SWEEP_LOG}"; continue; } + + # All NUM_FID indices assigned to this single GPU (resumable: the runner + # skips indices whose PNG already exists). + INDICES_FILE="${IDX_DIR}/all.txt" + seq 0 $((NUM_FID - 1)) > "${INDICES_FILE}" + + echo "[launch] GPU ${GPU_ID} -> ${TAG} (${NUM_FID} balanced, all on this GPU) at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + CUDA_VISIBLE_DEVICES=${GPU_ID} \ + SC_OWEN_MODE="${OWEN_MODE}" \ + python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 \ + --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec \ + --sc_qk_granularity per_row \ + --sc_config "${SC_JSON}" \ + --image-size 256 --num-sampling-steps "${NUM_STEPS}" --cfg-scale "${CFG_SCALE}" \ + --batch-size "${BATCH}" \ + --generate-fid-samples \ + --balanced_classes \ + --num-classes "${NUM_CLASSES}" \ + --balanced_total_samples "${NUM_FID}" \ + --num-fid-samples "${NUM_FID}" \ + --target_indices_path "${INDICES_FILE}" \ + --samples_dir_override "${SAMPLES}" \ + --seed ${GPU_ID} \ + --results-dir "${GPU_LOG_DIR}" \ + --ckpt "${CKPT}" \ + > "${GPU_LOG_DIR}/run.log" 2>&1 & + PIDS+=($!) + sleep 5 +done + +echo "[wait] ${#PIDS[@]} parallel workers: ${PIDS[*]}" | tee -a "${SWEEP_LOG}" +FAILED=0 +for i in "${!PIDS[@]}"; do + if wait "${PIDS[$i]}"; then + echo " [worker $i] done at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + else + echo " [worker $i] FAILED (rc=$?)" | tee -a "${SWEEP_LOG}" + FAILED=$((FAILED + 1)) + fi +done + +for SL in "${SLS[@]}"; do + N=$(find "${OUT_BASE}/uniform${SL}/samples" -maxdepth 1 -type f -name '*.png' 2>/dev/null | wc -l) + echo "[count] uniform${SL}: ${N}/${NUM_FID}" | tee -a "${SWEEP_LOG}" +done +echo "Generation finished $(date) (failed workers: ${FAILED})" | tee -a "${SWEEP_LOG}" + +# --- robust ADM/OpenAI eval (Inception Score / FID / sFID / Precision / Recall) --- +if [[ "${RUN_EVAL:-1}" == "1" && ${FAILED} -eq 0 ]]; then + for SL in "${SLS[@]}"; do + D="${OUT_BASE}/uniform${SL}/samples" + N=$(find "${D}" -maxdepth 1 -type f -name '*.png' 2>/dev/null | wc -l) + [[ "${N}" -ge "${NUM_FID}" ]] || { echo "[eval-skip] uniform${SL} ${N}/${NUM_FID}" | tee -a "${SWEEP_LOG}"; continue; } + echo "[eval] uniform${SL} ..." | tee -a "${SWEEP_LOG}" + bash "${REPO_ROOT}/scripts/eval_openai.sh" "${D}" "${OUT_BASE}/uniform${SL}" 2>&1 | tee -a "${SWEEP_LOG}" + done + echo "=== metrics summary ===" | tee -a "${SWEEP_LOG}" + for SL in "${SLS[@]}"; do + T="${OUT_BASE}/uniform${SL}.openai_eval.txt" + [[ -f "${T}" ]] || continue + echo "uniform${SL}:" | tee -a "${SWEEP_LOG}" + grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "${T}" | sed 's/^/ /' | tee -a "${SWEEP_LOG}" + done +fi +echo "DONE $(date)" | tee -a "${SWEEP_LOG}" diff --git a/scripts/run_ngpu_uniform_nohalve_fanout.sh b/scripts/run_ngpu_uniform_nohalve_fanout.sh new file mode 100755 index 0000000..17ade16 --- /dev/null +++ b/scripts/run_ngpu_uniform_nohalve_fanout.sh @@ -0,0 +1,186 @@ +#!/bin/bash +# ============================================================ +# 3-GPU uniform-precision FID sample sweep. +# +# sc_prec=8, halve ON (sweepable: explicit stoc_len, rng grid -> 128), +# bitrev scramble, qk & av granularity = per_row, all operators + all +# timesteps SC. Sweeps stoc_len in {256, 192, 128}; NUM_FID samples per +# config, fanned out across 3 GPUs with work-stealing (resumable). +# +# Mirrors scmp_diffusion_prev/Q-DiT/scripts/run_2gpu_calib_fid_sweep.sh. +# +# Usage (from a 3-GPU allocation): +# conda activate qdit +# cd .../scmp_diffusion +# bash scripts/run_3gpu_uniform_halve_bitrev.sh +# +# Override knobs (env vars): +# NUM_FID=10000 samples per config (must be divisible by NUM_CLASSES) +# STOC_LENS=256,192,128 +# BATCH=32 per-GPU batch +# NUM_STEPS=50 sampling steps +# CFG_SCALE=4 classifier-free guidance scale +# OWEN_MODE=bitrev counter | bitrev | random +# CKPT=/scratch/.../pretrained_models/DiT-XL-2-256x256.pt +# OUT_BASE=/scratch/.../scmp_diffusion_fid_halve_bitrev +# CONFIG_DIR=$OUT_BASE/configs (uniform${SL}_all.json live here) +# IMAGENET_REF=/path/to/val if set + pytorch_fid present -> compute FID +# ============================================================ + +set -euo pipefail + +NUM_GPUS="${NUM_GPUS:-3}" +NUM_FID="${NUM_FID:-10000}" +STOC_LENS="${STOC_LENS:-256,192,128}" +BATCH="${BATCH:-32}" +NUM_STEPS="${NUM_STEPS:-50}" +CFG_SCALE="${CFG_SCALE:-4}" +OWEN_MODE="${OWEN_MODE:-bitrev}" +NUM_CLASSES="${NUM_CLASSES:-1000}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi}" +CKPT="${CKPT:-${SCRATCH_ROOT}/pretrained_models/DiT-XL-2-256x256.pt}" +OUT_BASE="${OUT_BASE:-${SCRATCH_ROOT}/scmp_diffusion_fid_halve_bitrev}" +CONFIG_DIR="${CONFIG_DIR:-${OUT_BASE}/configs}" + +mkdir -p "${OUT_BASE}" +SWEEP_LOG="${OUT_BASE}/sweep.log" + +if (( NUM_FID % NUM_CLASSES != 0 )); then + echo "ERROR: NUM_FID (${NUM_FID}) must be divisible by NUM_CLASSES (${NUM_CLASSES})" >&2 + exit 1 +fi +if [[ ! -f "${CKPT}" ]]; then + echo "ERROR: checkpoint not found: ${CKPT}" >&2 + exit 1 +fi + +export SC_OWEN_MODE="${OWEN_MODE}" +export PYTHONUNBUFFERED=1 + +{ +echo "============================================================" +echo "N-GPU uniform NOHALVE bitrev sweep started $(date)" +echo " NUM_GPUS=${NUM_GPUS} NUM_FID=${NUM_FID} STOC_LENS=${STOC_LENS}" +echo " BATCH=${BATCH} NUM_STEPS=${NUM_STEPS} CFG_SCALE=${CFG_SCALE}" +echo " OWEN_MODE=${OWEN_MODE} sc_prec=8 halve=OFF qk/av=per_row" +echo " CKPT=${CKPT}" +echo " OUT_BASE=${OUT_BASE}" +echo "============================================================" +} | tee -a "${SWEEP_LOG}" + +cd "${REPO_ROOT}" + +run_config() { + local SL="$1" + local TAG="uniform${SL}" + local CFG_DIR="${OUT_BASE}/${TAG}" + local SAMPLES="${CFG_DIR}/samples" + local IDX_DIR="${CFG_DIR}/_indices" + local SC_JSON="${CONFIG_DIR}/sc_cfg_uniform${SL}_all.json" + mkdir -p "${SAMPLES}" "${IDX_DIR}" + + if [[ ! -f "${SC_JSON}" ]]; then + echo "[skip] ${TAG}: config not found ${SC_JSON}" | tee -a "${SWEEP_LOG}" + return + fi + + local DONE + DONE=$(find "${SAMPLES}" -maxdepth 1 -type f -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' 2>/dev/null | wc -l) + if [[ ${DONE} -ge ${NUM_FID} ]]; then + echo "[skip] ${TAG} already complete (${DONE}/${NUM_FID})" | tee -a "${SWEEP_LOG}" + return + fi + + echo "[plan] ${TAG} ${DONE}/${NUM_FID} done; planning across ${NUM_GPUS} GPUs at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + python -u "${REPO_ROOT}/scripts/_plan_missing_indices.py" \ + "${SAMPLES}" "${NUM_FID}" "${NUM_GPUS}" "${IDX_DIR}" | tee -a "${SWEEP_LOG}" + + local PIDS=() + for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do + local INDICES_FILE="${IDX_DIR}/gpu_${GPU_ID}.txt" + if [[ ! -s "${INDICES_FILE}" ]]; then + echo " [GPU ${GPU_ID}] no work assigned" | tee -a "${SWEEP_LOG}" + continue + fi + local GPU_LOG_DIR="${CFG_DIR}/_logs/gpu_${GPU_ID}" + mkdir -p "${GPU_LOG_DIR}" + CUDA_VISIBLE_DEVICES=${GPU_ID} \ + SC_OWEN_MODE="${OWEN_MODE}" \ + python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 \ + --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec \ + --sc_qk_granularity per_row \ + --sc_config "${SC_JSON}" \ + --image-size 256 --num-sampling-steps "${NUM_STEPS}" --cfg-scale "${CFG_SCALE}" \ + --batch-size "${BATCH}" \ + --generate-fid-samples \ + --balanced_classes \ + --num-classes "${NUM_CLASSES}" \ + --balanced_total_samples "${NUM_FID}" \ + --num-fid-samples "${NUM_FID}" \ + --target_indices_path "${INDICES_FILE}" \ + --samples_dir_override "${SAMPLES}" \ + --seed "${SEED:-0}" \ + --results-dir "${GPU_LOG_DIR}" \ + --ckpt "${CKPT}" \ + > "${GPU_LOG_DIR}/run.log" 2>&1 & + PIDS+=($!) + sleep 5 + done + + local FAILED=0 + for i in "${!PIDS[@]}"; do + if wait "${PIDS[$i]}"; then + echo " [worker $i] done at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + else + echo " [worker $i] FAILED (rc=$?); see ${CFG_DIR}/_logs/" | tee -a "${SWEEP_LOG}" + FAILED=$((FAILED + 1)) + fi + done + + local FINAL + FINAL=$(find "${SAMPLES}" -maxdepth 1 -type f -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' 2>/dev/null | wc -l) + if [[ ${FINAL} -ge ${NUM_FID} ]]; then + echo "[ok] ${TAG} ${FINAL}/${NUM_FID} samples in ${SAMPLES} at $(date +%H:%M:%S)" | tee -a "${SWEEP_LOG}" + else + echo "[partial] ${TAG} ${FINAL}/${NUM_FID}; rerun this script to resume" | tee -a "${SWEEP_LOG}" + fi +} + +IFS=',' read -ra SLS <<< "${STOC_LENS}" +for SL in "${SLS[@]}"; do + run_config "${SL}" +done + +echo "Sample generation finished $(date)" | tee -a "${SWEEP_LOG}" + +# --- robust ADM/OpenAI evaluation (reuses scmp_diffusion_prev evaluator + +# ImageNet-256 reference): Inception Score / FID / sFID / Precision / Recall. +# Set RUN_EVAL=0 to skip (e.g. generate now, evaluate later). +if [[ "${RUN_EVAL:-1}" == "1" ]]; then + for SL in "${SLS[@]}"; do + D="${OUT_BASE}/uniform${SL}/samples" + N=$(find "${D}" -maxdepth 1 -type f -name '*.png' 2>/dev/null | wc -l) + if [[ "${N}" -lt "${NUM_FID}" ]]; then + echo "[eval-skip] uniform${SL}: only ${N}/${NUM_FID} samples" | tee -a "${SWEEP_LOG}" + continue + fi + echo "[eval] uniform${SL} ($(date +%H:%M:%S)) ..." | tee -a "${SWEEP_LOG}" + bash "${REPO_ROOT}/scripts/eval_openai.sh" "${D}" "${OUT_BASE}/uniform${SL}" \ + 2>&1 | tee -a "${SWEEP_LOG}" + done + echo "=== metrics summary ===" | tee -a "${SWEEP_LOG}" + for SL in "${SLS[@]}"; do + T="${OUT_BASE}/uniform${SL}.openai_eval.txt" + [[ -f "${T}" ]] || continue + echo "uniform${SL}:" | tee -a "${SWEEP_LOG}" + grep -E "^(Inception Score|FID|sFID|Precision|Recall):" "${T}" | sed 's/^/ /' | tee -a "${SWEEP_LOG}" + done +else + echo "Eval skipped (RUN_EVAL=0). Run scripts/eval_openai.sh per config later." | tee -a "${SWEEP_LOG}" +fi +echo "DONE $(date)" | tee -a "${SWEEP_LOG}"