SC FID eval (KID) + 9-GPU FP/W8A8 runners + MP calibration fix

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

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: <CKPT> <OUTDIR>
+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

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

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,

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 <ref_npz> <out_txt> <sample_npz> [<sample_npz> ...]
+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_npz> [...]")
+    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()

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 <out_txt> <sample_npz> [<sample_npz> ...]
+# Overrides: PREV, REF_NPZ, EVALUATOR
+# ============================================================
+set -euo pipefail
+
+OUT_TXT="${1:?usage: kid_openai.sh <out_txt> <sample_npz> [...]}"
+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"

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 <samples_dir> <out_prefix>
+#     <samples_dir>  dir of NNNNNN.png (256x256)
+#     <out_prefix>   writes <out_prefix>.npz and <out_prefix>.openai_eval.txt
+#
+# Override:
+#   PREV=/gpfs/.../scmp_diffusion_prev   (evaluator + reference live here)
+# ============================================================
+set -euo pipefail
+
+SAMPLES="${1:?usage: eval_openai.sh <samples_dir> <out_prefix>}"
+OUT_PREFIX="${2:?usage: eval_openai.sh <samples_dir> <out_prefix>}"
+
+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