#!/usr/bin/env bash
# NVFP4 quantization of a DiffusionGemma checkpoint with NVIDIA Model Optimizer,
# experts only, as nvidia/diffusiongemma-26B-A4B-it-NVFP4 was made. Runs on a
# rented GPU machine with at least 130 GB of GPU memory (the batch-size probe
# peaked at 127 GB on an H200) and 150 GB of free disk.
#   quantize-diffusiongemma-nvfp4.sh <hf repo> <40-character revision> [calibration samples]
# Writes /root/out and prints its file list and hashes. Run it under nohup.
set -euo pipefail

SOURCE="$1"
REVISION="$2"
SAMPLES="${3:-128}"
# The DiffusionGemma recipe is on main only; release 0.47.0 does not have it.
MODELOPT_COMMIT=d0142c9dcad8f456caa482407960dc9cb5f0e710

cd /root
echo "start $(date -u +%FT%TZ) source=$SOURCE@$REVISION samples=$SAMPLES"
# Triton compiles a helper against Python.h on first use.
apt-get update -q >/dev/null
apt-get install -y -q python3-venv python3-dev git >/dev/null
[[ -d q ]] || python3 -m venv q
. q/bin/activate
pip install -q -U pip
if [[ ! -d modelopt ]]; then
  git init -q modelopt
  git -C modelopt fetch -q --depth 1 https://github.com/NVIDIA/Model-Optimizer "$MODELOPT_COMMIT"
  git -C modelopt checkout -q FETCH_HEAD
  # transformers 5.14 rejects any use_cache argument in the DiffusionGemma decoder,
  # and Model Optimizer sets config.use_cache for its calibration forwards.
  python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/modelopt/torch/utils/dataset_utils.py")
s = p.read_text()
anchor = '    """Yield the top-level config and Step3.7-style nested text config."""\n'
assert s.count(anchor) == 1
p.write_text(s.replace(anchor, anchor + '    if "diffusiongemma" in model.__class__.__name__.lower():\n        return\n'))
PY
fi
pip install -q -e "modelopt[hf]"
pip install -q "huggingface_hub[hf_xet]"
python -c 'import torch, transformers; print("torch", torch.__version__, "cuda", torch.cuda.is_available(), "transformers", transformers.__version__)'
git -C modelopt log --oneline -1

hf download "$SOURCE" --revision "$REVISION" --local-dir /root/src >/dev/null
du -sh /root/src
echo "downloaded $(date -u +%FT%TZ)"

# cnn_dailymail only: NVIDIA's second calibration set is gated on the Hub.
python modelopt/examples/hf_ptq/hf_ptq.py \
  --pyt_ckpt_path /root/src \
  --recipe model_type/diffusion_gemma/ptq/nvfp4_experts_only \
  --dataset cnn_dailymail \
  --calib_size "$SAMPLES" \
  --skip_generate \
  --export_path /root/out
echo "done $(date -u +%FT%TZ)"
du -sh /root/out
(cd /root/out && sha256sum * > /root/out.sha256 && cat /root/out.sha256)
echo FINISHED