unsloth/scripts/build_prequant_checkpoint.py at main · unslothai/unsloth · GitHub

History

268 lines (236 loc) · 11.3 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

# SPDX-License-Identifier: AGPL-3.0-only

"""Build a pre-quantized transformer checkpoint for the Unsloth diffusion fast path.

Quantise a model's dense bf16 DiT transformer ONCE and save the quantized state dict, so

the backend can load the already-quantized weights at runtime (meta-init +

load_state_dict(assign=True)) instead of materialising the dense bf16 on the GPU. That

drops the transformer GPU load peak ~2x and the download ~2x for fp8 (measured on Z-Image:

12.9 -> 6.3 GB peak, 12 -> 6.28 GB on disk), with bit-identical output -- it is the exact

same torchao config + min_features filter the runtime path uses, applied ahead of time.

Run on one CUDA (Blackwell / Ada / Hopper) GPU. fp8 works on torch 2.9+; the FP4/MX schemes

need the newer kernels (see scripts/nvfp4_t211_probe.py).

python scripts/build_prequant_checkpoint.py \

--base Tongyi-MAI/Z-Image-Turbo --family z-image --scheme fp8 \

--out outputs/quant_research/prequant_fp8/transformer_fp8.pt [--upload-repo ORG/REPO]

"""

from __future__ import annotations

import argparse

import sys

import time

from pathlib import Path

from typing import Any, Optional, Sequence

BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"

def convrot_refusal(

group: int, rotatable: Sequence[str], not_divisible: Sequence[str]

) -> Optional[str]:

"""Why a ConvRot build must not be quantised and saved, or None when it is fine.

An empty rotatable set means the group divides no quantized input axis (a group larger than

every Linear, say). The build would still stamp the v2 tag and an empty fqn list, which

``rotation_metadata_error`` refuses at load time, so the only thing it produces is a

multi-gigabyte artifact nothing can ever open. Refuse before the hours, not after."""

if rotatable:

return None

return (

f"ConvRot group {group} divides the in_features of none of the {len(not_divisible)} "

"quantized linears, so the checkpoint would record an empty rotation and be refused at "

"load time. Pick a smaller power-of-4 group, or drop --convrot-groupsize."

)

def upload_destination(

fam: Any,

scheme: str,

rotated: bool,

override: Optional[str] = None,

) -> str:

"""The repo-root filename this build should publish under.

The loader asks for the family's declared ``prequant_filenames`` name first and the derived

``<Model>-<SCHEME>.pt`` second, so a ROTATED artifact published under the legacy

``transformer_<scheme>.pt`` is either never resolved at all, or resolved as the fallback by a

build too old to honour the rotation, which then refuses the v2 tag and drops to the dense

download. A rotated build therefore goes to the declared name or nowhere. Plain builds keep

the legacy name they have always used."""

if override:

return override

from core.inference.diffusion_prequant import prequant_filename

if not rotated:

return prequant_filename(scheme)

from core.inference.diffusion_families import family_prequant_filename

preferred = family_prequant_filename(fam, scheme)

if not preferred:

raise ValueError(

f"family {getattr(fam, 'name', fam)!r} declares no prequant_filenames entry for "

f"{scheme!r}, so a rotated checkpoint has no name the loader would ask for. Add the "

"entry to the family table, or pass --upload-filename."

)

return preferred

def main(argv = None) -> int:

p = argparse.ArgumentParser()

p.add_argument(

"--base", required = True, help = "diffusers base repo (carries the transformer subfolder)"

)

p.add_argument("--family", required = True, help = "diffusion family name/alias (e.g. z-image)")

p.add_argument("--scheme", required = True, help = "quant scheme: int8 | fp8 | nvfp4 | mxfp8")

p.add_argument("--out", required = True, help = "output .pt path for the checkpoint")

p.add_argument("--min-features", type = int, default = 512)

p.add_argument("--dtype", default = "bfloat16", choices = ["bfloat16"])

p.add_argument("--hf-token", default = None)

p.add_argument(

"--convrot-groupsize",

type = int,

default = 0,

help = "bake a ConvRot block-Hadamard activation rotation at this group size (a power of "

"4; 0 = off). Every quantized Linear whose in_features the group divides has its "

"weight rotated before quantize_ so the quantizer sees a flatter distribution; the "

"exact fqn list is recorded in the checkpoint and the loader rotates the "

"activations of that list and nothing else. Writes the v2 format tag.",

)

p.add_argument(

"--upload-repo", default = None, help = "optional HF repo id to upload the checkpoint to"

)

p.add_argument("--upload-revision", default = None)

p.add_argument(

"--upload-filename",

default = None,

help = "repo-root filename to publish under; defaults to the family's declared "

"prequant_filenames entry for a rotated build and the legacy transformer_<scheme>.pt "

"otherwise",

)

args = p.parse_args(argv)

sys.path.insert(0, str(BACKEND))

import torch

import torchao

import diffusers

from core.inference.diffusion_families import detect_family

from core.inference.diffusion_prequant import prequant_format_for

# Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant).

from core.inference.diffusion_transformer_quant import (

FP8_GRANULARITY,

TQ_FP8,

TQ_SCHEMES,

_REQUIRE_BF16_SCHEMES,

_make_quant_config,

_resolve_fast_accum,

exclude_tokens_for_scheme,

make_filter_fn,

)

from torchao.quantization import quantize_

scheme = args.scheme.strip().lower()

if scheme not in TQ_SCHEMES:

print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush = True)

return 2

fam = detect_family(args.base, override = args.family)

if fam is None:

print(f"error: unknown family '{args.family}'", flush = True)

return 2

transformer_cls = getattr(diffusers, fam.transformer_class)

# Resolved BEFORE the load, so a rotated build with nowhere resolvable to publish fails in a

# second rather than after the quantise and the multi-gigabyte save.

upload_dest = None

if args.upload_repo:

try:

upload_dest = upload_destination(

fam,

scheme,

rotated = bool(args.convrot_groupsize),

override = args.upload_filename,

)

except ValueError as exc:

print(f"error: {exc}", flush = True)

return 2

print(f"== build prequant ({fam.name}/{scheme}, min_feat={args.min_features}) ==", flush = True)

print(f" loading dense transformer from {args.base} (subfolder=transformer) ...", flush = True)

t0 = time.time()

transformer = transformer_cls.from_pretrained(

args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token

).to("cuda")

print(f" quantising in place ({scheme}) ...", flush = True)

# Mirror the runtime exclusions: int8 skips the M=1 modulation projections (torch._int_mm needs M>16) plus per-family ones; family=None bakes linears the runtime rejects.

exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name)

# fp8 / mxfp8 need bf16 weights, so skip non-bf16 Linears; nvfp4 handles fp32. Mirrors the runtime gate.

require_bf16 = scheme in _REQUIRE_BF16_SCHEMES

# fp8 bakes the accumulate mode in; record it so the loader can reject a contradicting request.

fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None

filter_fn = make_filter_fn(

args.min_features,

exclude_name_tokens = exclude_name_tokens,

require_bf16 = require_bf16,

)

# ConvRot, BEFORE quantize_: rotating the weights is only worth anything if the quantizer then

# sees the rotated distribution. The fqn list is recorded, never re-derived at load time.

rotation: dict = {}

if args.convrot_groupsize:

from core.inference.diffusion_convrot import (

rotatable_fqns,

rotate_linears_,

rotation_metadata,

)

group = int(args.convrot_groupsize)

rotatable, not_divisible = rotatable_fqns(transformer, filter_fn, group)

refusal = convrot_refusal(group, rotatable, not_divisible)

if refusal:

print(f"error: {refusal}", flush = True)

return 2

rotate_linears_(transformer, rotatable, group)

rotation = rotation_metadata(group, rotatable)

print(

f" rotated {len(rotatable)} linears at ConvRot group {group}; "

f"{len(not_divisible)} quantized linears left plain (in_features not divisible)"

+ (f", e.g. {not_divisible[0]}" if not_divisible else ""),

flush = True,

)

quantize_(transformer, _make_quant_config(scheme), filter_fn = filter_fn)

# CPU state dict for a portable, GPU-free artifact.

state_dict = {

k: (v.detach().to("cpu") if hasattr(v, "detach") else v)

for k, v in transformer.state_dict().items()

}

metadata = {

"base_model_id": args.base,

"family": fam.name,

"scheme": scheme,

"min_features": args.min_features,

# Let the loader reject a checkpoint that would not match the runtime path.

"exclude_name_tokens": list(exclude_name_tokens),

"require_bf16": require_bf16,

"fast_accum": fast_accum,

"torch_dtype": args.dtype,

"quant_backend": "torchao",

"transformer_class": fam.transformer_class,

"torch_version": torch.__version__,

"torchao_version": getattr(torchao, "__version__", "?"),

"diffusers_version": diffusers.__version__,

}

# fp8 granularity: lets the loader reject a stale per-tensor checkpoint (runtime needs per-row).

if scheme == TQ_FP8:

metadata["fp8_granularity"] = FP8_GRANULARITY

metadata.update(rotation)

ckpt = {

# v2 when a rotation is baked in, so an Unsloth predating the online half refuses the file

# rather than running the rotated weights against unrotated activations.

"format": prequant_format_for(metadata),

"metadata": metadata,

"state_dict": state_dict,

}

out = Path(args.out)

out.parent.mkdir(parents = True, exist_ok = True)

torch.save(ckpt, out)

size_gb = out.stat().st_size / 1e9

print(f" saved {out} ({size_gb:.2f} GB) in {time.time() - t0:.0f}s", flush = True)

print(f" metadata: {ckpt['metadata']}", flush = True)

if args.upload_repo:

from huggingface_hub import HfApi

dest = upload_dest

print(f" uploading -> {args.upload_repo}:{dest} ...", flush = True)

api = HfApi(token = args.hf_token)

api.create_repo(args.upload_repo, exist_ok = True)

api.upload_file(

path_or_fileobj = str(out),

path_in_repo = dest,

repo_id = args.upload_repo,

revision = args.upload_revision,

)

print(f" uploaded {dest} to {args.upload_repo}", flush = True)

print("BUILD-PREQUANT-DONE", flush = True)

return 0

if __name__ == "__main__":

sys.exit(main())

unsloth/scripts/build_prequant_checkpoint.py at main · unslothai/unsloth · GitHub

Uh oh!

Uh oh!

unsloth/scripts/build_prequant_checkpoint.py at main · unslothai/unsloth · GitHub

Uh oh!

Uh oh!

Expand file tree

Breadcrumbs

build_prequant_checkpoint.py

Latest commit

History

Breadcrumbs

build_prequant_checkpoint.py

File metadata and controls