File size: 2,544 Bytes
01badea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | #!/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
|