| """Offline reproducible Quark weight conversion; run in the pinned conversion environment.""" |
| import os |
| os.environ.update(QUARK_MXFP4_IMPL="hip", QUARK_DISABLE_COMPILE="1", |
| TORCH_COMPILE_DISABLE="1", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1") |
| import contextlib |
| from datetime import datetime, timezone |
| import hashlib |
| import importlib.metadata as metadata |
| import json |
| import math |
| from pathlib import Path |
| import shutil |
| import time |
| import traceback |
|
|
| import argparse |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--snapshot", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--profile", choices=("compact", "balanced"), default="balanced") |
| args = parser.parse_args() |
| ROOT = args.output.resolve().parent |
| REVISION = "790c92633540aa0cb11d9abf19eb46d861714758" |
| SOURCE = args.snapshot.resolve() |
| OUT = args.output.resolve() |
| OUT.mkdir(exist_ok=False) |
| shutil.copyfile(__file__, OUT / "convert.py") |
| started = time.perf_counter() |
| report = dict(status="running", format="paiton-research-mxfp4-weight-only-v2", |
| checkpoint_revision=REVISION, components={}, |
| quantization="Quark OCP MXFP4 E2M1, E8M0 scale per 32 columns, even scale selection, low nibble first", |
| activations="BF16", accumulator="unchanged external BF16 reference operations", |
| calibration="weight-only round-to-nearest; no activation calibration or learned correction", |
| preserved="normalizations, biases and non-matrix parameters use the BF16 inference values", |
| padding="Zero pad flattened input dimension to multiple of 32; crop before original reshape", |
| modification_notice="Weights converted from Qwen-Image-2.1 by EliovpAI for research/evaluation; Built with Qwen") |
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as f: |
| for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
| def save_report(): |
| report["elapsed_seconds"] = time.perf_counter() - started |
| p = OUT / "result.json.tmp" |
| p.write_text(json.dumps(report, indent=2) + "\n") |
| p.replace(OUT / "result.json") |
|
|
| print("OUTPUT:", OUT, flush=True) |
| try: |
| import torch |
| from accelerate import init_empty_weights |
| from diffusers import QwenImage21Transformer2DModel, AutoencoderKLQwenImage21 |
| from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration |
| from safetensors import safe_open |
| from safetensors.torch import load_file, save_file |
| from quark.torch import ModelQuantizer |
| from quark.torch.quantization.config.config import QConfig, QLayerConfig, OCP_MXFP4Spec |
| from quark.torch.export.nn.modules.qparamslinear import QParamsLinear |
|
|
| report["versions"] = {p: metadata.version(p) for p in |
| ("torch", "amd-quark", "diffusers", "transformers", "accelerate", "safetensors")} |
| assert report["versions"]["torch"] == "2.10.0+rocm7.1" |
| assert report["versions"]["amd-quark"] == "0.12.post1+rocm71.torch2.10" |
| assert torch.version.hip and torch.cuda.mem_get_info()[0] > 220 * 2**30 |
| assert shutil.disk_usage(ROOT).free > 15 * 2**30 |
| torch.set_num_threads(8) |
| torch.set_num_interop_threads(2) |
| torch.set_grad_enabled(False) |
| report["gpu"] = str(torch.cuda.get_device_properties(0)) |
| report["script_sha256"] = digest(__file__) |
| config = QConfig(global_quant_config=QLayerConfig(weight=OCP_MXFP4Spec( |
| ch_axis=-1, is_dynamic=False, scale_calculation_mode="even").to_quantization_spec())) |
| report["quark_config"] = config.to_dict() |
| source_pin = json.loads((Path(__file__).resolve().parent / "source-manifest.json").read_text()) |
| assert source_pin["revision"] == REVISION |
| for file in source_pin["files"]: |
| path = SOURCE / file["path"] |
| assert path.stat().st_size == file["bytes"] and digest(path) == file["sha256"], path |
| report["profile"] = args.profile |
| if args.profile == "balanced": |
| report["preserved"] = "BF16 VAE, DiT input/output/conditioning, normalizations and biases; BF16 activations" |
| table = torch.tensor([0,.5,1,1.5,2,3,4,6,-0.,-.5,-1,-1.5,-2,-3,-4,-6], device="cuda") |
| details = (OUT / "quark.log").open("w", buffering=1) |
|
|
| def convert(original): |
| rows = original.shape[0] |
| cols = math.prod(original.shape[1:]) |
| padded = (cols + 31) // 32 * 32 |
| original = original.reshape(rows, cols) |
| all_weights, all_scales = [], [] |
| square_error, square_value = 0.0, 0.0 |
| for begin in range(0, rows, 4096): |
| value = original[begin:begin+4096].to("cuda", torch.bfloat16) |
| if padded != cols: |
| value = torch.nn.functional.pad(value, (0, padded-cols)) |
| linear = torch.nn.Linear(padded, value.shape[0], bias=False, device="meta", dtype=torch.bfloat16) |
| linear.weight = torch.nn.Parameter(value, requires_grad=False) |
| with contextlib.redirect_stdout(details), contextlib.redirect_stderr(details): |
| quantized = ModelQuantizer(config).quantize_model(torch.nn.Sequential(linear), None)[0] |
| expected = quantized.get_quant_weight(quantized.weight) |
| packed = QParamsLinear.from_module(quantized, custom_mode="quark", pack_method="order") |
| weight, scale = packed.weight, packed.weight_quantizer.scale |
| assert weight.dtype == scale.dtype == torch.uint8 |
| assert tuple(weight.shape) == (value.shape[0], padded//2) |
| assert scale.numel() == value.shape[0] * padded // 32 |
| scale = scale.reshape(value.shape[0], padded//32) |
| codes = torch.stack((weight & 15, weight >> 4), -1).reshape(value.shape[0], padded) |
| decoded = (table[codes.long()].reshape(value.shape[0], padded//32, 32) |
| * torch.exp2(scale.float()-127).unsqueeze(-1)).reshape(value.shape).to(torch.bfloat16) |
| assert torch.equal(decoded, expected), "Independent MXFP4 decode differs from Quark" |
| assert torch.isfinite(decoded).all().item() |
| difference = decoded[:,:cols].float() - value[:,:cols].float() |
| square_error += float(difference.square().double().sum()) |
| square_value += float(value[:,:cols].float().square().double().sum()) |
| all_weights.append(weight.cpu().contiguous()) |
| all_scales.append(scale.cpu().contiguous()) |
| del linear, quantized, expected, packed, weight, scale, codes, decoded, difference, value |
| return torch.cat(all_weights), torch.cat(all_scales), dict( |
| padded_shape=[rows, padded], weight_relative_l2=math.sqrt(square_error/max(square_value, 1e-30)), |
| independently_matches_quark=True) |
|
|
| for component in ("transformer", "text_encoder", "vae"): |
| dest = OUT / component |
| dest.mkdir() |
| shutil.copyfile(SOURCE/component/"config.json", dest/"config.json") |
| with init_empty_weights(include_buffers=False): |
| if component == "text_encoder": |
| cfg = Qwen3VLConfig.from_pretrained(str(SOURCE/component), local_files_only=True) |
| cfg.dtype = torch.bfloat16 |
| model = Qwen3VLForConditionalGeneration(cfg) |
| else: |
| cls = QwenImage21Transformer2DModel if component == "transformer" else AutoencoderKLQwenImage21 |
| model = cls.from_config(cls.load_config(str(SOURCE/component))) |
| expected = {key: list(p.shape) for key,p in model.named_parameters()} |
| targets = {} |
| for name, module in model.named_modules(): |
| if isinstance(module, (torch.nn.Linear, torch.nn.Embedding, torch.nn.Conv2d, torch.nn.Conv3d)): |
| key = name + ".weight" |
| if args.profile == "balanced" and (component == "vae" or (component == "transformer" and not name.startswith("transformer_blocks."))): |
| continue |
| assert key in expected |
| targets[key] = dict(name=name, shape=expected[key], kind=( |
| "linear" if isinstance(module,torch.nn.Linear) else |
| "embedding" if isinstance(module,torch.nn.Embedding) else |
| "conv2d" if isinstance(module,torch.nn.Conv2d) else "conv3d"), |
| original_class=type(module).__name__) |
| del model |
| row = dict(source_bf16_bytes=0, tensor_payload_bytes=0, converted_layers=[], files=[], |
| original_parameters=len(expected), source_files=[], config_sha256=digest(dest/"config.json")) |
| report["components"][component] = row |
| tensors, output_map, loaded = {}, {}, set() |
| pending_bytes = 0 |
|
|
| def flush(): |
| if not tensors: |
| return |
| filename = f"weights-{len(row['files'])+1:05d}.safetensors" |
| path = dest / filename |
| save_file(tensors, str(path), metadata={"format":report["format"], "modification_notice":report["modification_notice"]}) |
| check = load_file(str(path)) |
| assert check.keys() == tensors.keys() and all(torch.equal(t,check[k]) for k,t in tensors.items()) |
| size = sum(t.numel()*t.element_size() for t in tensors.values()) |
| row["tensor_payload_bytes"] += size |
| row["files"].append(dict(file=filename,bytes=path.stat().st_size,sha256=digest(path))) |
| output_map.update({k:filename for k in tensors}) |
| tensors.clear() |
| save_report() |
|
|
| for path in sorted((SOURCE/component).glob("*.safetensors")): |
| row["source_files"].append(dict(file=path.name,bytes=path.stat().st_size,sha256=digest(path))) |
| with safe_open(path, framework="pt", device="cpu") as f: |
| for key in sorted(f.keys()): |
| assert key in expected and key not in loaded, key |
| value = f.get_tensor(key) |
| assert list(value.shape) == expected[key], key |
| row["source_bf16_bytes"] += value.numel()*2 |
| if key in targets: |
| layer = dict(targets[key]) |
| weight, scale, detail = convert(value) |
| layer.update(detail) |
| tensors[key] = weight |
| tensors[layer["name"]+".weight_scale"] = scale |
| row["converted_layers"].append(layer) |
| print(f"{component}: {len(row['converted_layers'])}/{len(targets)} {layer['kind']} {key}", flush=True) |
| else: |
| tensors[key] = value.to(torch.bfloat16).clone().contiguous() |
| loaded.add(key) |
| del value |
| pending_bytes = sum(t.numel()*t.element_size() for t in tensors.values()) |
| if pending_bytes >= 192 * 2**20: |
| flush() |
| flush() |
| assert loaded == set(expected), sorted(set(expected)-loaded) |
| assert len(row["converted_layers"]) == len(targets) |
| index = dest / "weights.index.json" |
| index.write_text(json.dumps(dict(metadata=dict(total_size=row["tensor_payload_bytes"]),weight_map=output_map),indent=2)+"\n") |
| row["index_sha256"] = digest(index) |
| print("COMPONENT COMPLETE:",component,"GiB",row["tensor_payload_bytes"]/2**30,flush=True) |
| save_report() |
|
|
| report["total_tensor_payload_bytes"] = sum(r["tensor_payload_bytes"] for r in report["components"].values()) |
| report["total_original_bf16_bytes"] = sum(r["source_bf16_bytes"] for r in report["components"].values()) |
| assert report["total_original_bf16_bytes"] == 32439977672 |
| report["peak_allocated_bytes"] = torch.cuda.max_memory_allocated() |
| for name in ("processor", "scheduler"): |
| shutil.copytree(SOURCE / name, OUT / name) |
| shutil.copyfile(SOURCE / "LICENSE", OUT / "LICENSE") |
| report["status"] = "completed_conversion_only" |
| except BaseException: |
| report["status"] = "failed" |
| report["traceback"] = traceback.format_exc() |
| raise |
| finally: |
| save_report() |
| print("COMPLETE:",OUT,"TOTAL PAYLOAD GiB:",report["total_tensor_payload_bytes"]/2**30,flush=True) |
|
|
