Blame for scripts/quantize_fp8.py

#!/usr/bin/env python3
"""Make a block-FP8 copy of Clef-Flash from Cloudflare's own files, with llm-compressor.

1. Check every input file against the hash that Cloudflare publishes for
   revision 41318bf. Stop on a mismatch.
2. Quantize with the FP8_BLOCK scheme of llm-compressor: FP8 weights with one
   scale for each 128x128 block, and FP8 activations scaled at run time. This
   is the layout of Qwen/Qwen3.6-27B-FP8. It uses no calibration data. The
   vision encoder, lm_head and in_proj_a/b stay in BF16, as in Qwen's release.
3. Copy the Clef head, Cloudflare's code and the tokenizer files unchanged.
   The new config.json is Cloudflare's config.json plus quantization_config.
4. Check the output. Every tensor that is not FP8 must be bit-identical to
   Cloudflare's tensor. No tensor can be missing or extra.
5. Write MANIFEST.json with the versions and the hash of each output file.

Run it in the amd-vllm022 image after `pip install llmcompressor==0.13.0`
(see README.md). The image already has torch 2.10.0 and transformers 5.12.1,
which llm-compressor 0.13.0 accepts. Runs on CPU.

    python3 quantize_fp8.py --src /path/to/clef-flash --dst /path/to/clef-flash-fp8
"""

import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path

import torch
from safetensors import safe_open

SOURCE = "Cloudflare/clef-flash@41318bf2106d135c7a92da745fec4b255279b4de"

# SHA-256 of the large files, as Hugging Face publishes them (the LFS oid).
SHA256 = {
    "joint_head.safetensors": "19cdcec8c81dc9212be320fff47462ab342fbc1278be4368fb3da71241cf5ba0",
    "model-00001-of-00004.safetensors": "8b45a8e968141cdcc58fb71c9adfc258e2c77b5f062bc636c1fd5bc5d916b565",
    "model-00002-of-00004.safetensors": "7590856c713eed844a2dcf48e6c43c4de165b788bc3f80e328311183cdbc7db8",
    "model-00003-of-00004.safetensors": "e6eac2467952c33361ed7dcb3c7959d1086bbe57201cd3749c3d769fdc17fe63",
    "model-00004-of-00004.safetensors": "9fcecc6556b39171238373a465f409794b7f821fb4cd1e6459e3a9c0fe317af7",
    "tokenizer.json": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523",
}

# Git blob hashes of the small files, as Hugging Face publishes them.
GIT_BLOB = {
    "joint_schema_model.py": "6a3ac5fc030a3f33f7fa860a1271886b64a6606d",
    "joint_head_config.json": "52467b6c7f7d1a7e628b518eb7a6cd7e9bdb9a57",
    "config.json": "aff766f2ab6e0e48b6aa297265689b981bdd2042",
    "model.safetensors.index.json": "778b7bbb7ebec3f1e100b1059b1c6d54bd00e3c1",
    "processor_config.json": "33818c7f9e991ad735fd240209f4fa73e6c28c50",
    "tokenizer_config.json": "0f19934ef65115549ff9629ff2c83d15e40e32d8",
}

# Qwen keeps these layers in BF16 in Qwen3.6-27B-FP8. The recipe targets only
# Linear layers, so the embeddings stay in BF16 too.
IGNORE = ["lm_head", "re:.*visual.*", "re:.*linear_attn.in_proj_a$", "re:.*linear_attn.in_proj_b$"]

# In each of the 32 layers: 3 MLP weights. In the 8 full-attention layers:
# q, k, v, o. In the 24 linear-attention layers: in_proj_qkv, in_proj_z, out_proj.
EXPECTED_FP8 = 32 * 3 + 8 * 4 + 24 * 3

# Files that the model save does not write, or writes differently. They are
# copied from Cloudflare's directory without change.
COPY = [
    "joint_head.safetensors",
    "joint_head_config.json",
    "joint_schema_model.py",
    "tokenizer.json",
    "tokenizer_config.json",
    "processor_config.json",
    "chat_template.jinja",
    "generation_config.json",
    "LICENSE",
]


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def git_blob(path: Path) -> str:
    data = path.read_bytes()
    return hashlib.sha1(b"blob %d\0" % len(data) + data).hexdigest()


def check_inputs(src: Path) -> list[str]:
    """Return one line for each input file that is missing or different."""
    problems = []
    for checks, hash_file in ((SHA256, sha256), (GIT_BLOB, git_blob)):
        for name, expected in checks.items():
            path = src / name
            if not path.is_file():
                problems.append(f"MISSING  {name}")
            elif hash_file(path) != expected:
                problems.append(f"CHANGED  {name}")
            else:
                print(f"OK       {name}")
    return problems


def quantize(src: Path, dst: Path) -> None:
    import contextlib
    import importlib

    from llmcompressor import oneshot
    from llmcompressor.modifiers.quantization import QuantizationModifier
    from transformers import Qwen3_5ForConditionalGeneration

    # Qwen's norms compute x * (1 + w). llm-compressor changes w to 1 + w for
    # calibration and back after, in BF16, which rounds every norm weight
    # (llmcompressor/modeling/offset_norm.py, still so in 0.14.0). FP8_BLOCK
    # without data does not use the norms, so this skips the change.
    # check_output() proves that the norms stay bit-identical. import_module
    # gives the module: the package attribute of the same name is the function.
    oneshot_module = importlib.import_module("llmcompressor.entrypoints.oneshot")
    oneshot_module.norm_calibration_context = lambda model: contextlib.nullcontext()

    model = Qwen3_5ForConditionalGeneration.from_pretrained(src, dtype=torch.bfloat16)
    oneshot(model=model, recipe=QuantizationModifier(targets="Linear", scheme="FP8_BLOCK", ignore=IGNORE))
    # Shards and an index, as in Cloudflare's release. The model server finds
    # lm_head through model.safetensors.index.json.
    model.save_pretrained(dst, save_compressed=True, max_shard_size="5GB")
    index = dst / "model.safetensors.index.json"
    if not index.exists():
        # A model that fits in one shard gets no index. Write it.
        with safe_open(str(dst / "model.safetensors"), framework="pt") as f:
            weight_map = {name: "model.safetensors" for name in f.keys()}
        index.write_text(json.dumps({"metadata": {}, "weight_map": weight_map}, indent=2) + "\n")


class Tensors:
    """Read tensors by name from a directory of safetensors shards."""

    def __init__(self, path: Path):
        self.shard_of = json.loads((path / "model.safetensors.index.json").read_text())["weight_map"]
        self.files = {s: safe_open(str(path / s), framework="pt", device="cpu") for s in set(self.shard_of.values())}

    def names(self) -> set[str]:
        return set(self.shard_of)

    def get(self, name: str) -> torch.Tensor:
        return self.files[self.shard_of[name]].get_tensor(name)


def check_output(src: Path, dst: Path) -> tuple[list[float], list[str]]:
    """Compare the output with Cloudflare's tensors.

    Return the relative error of each FP8 weight, and a list of problems.
    """
    before, after = Tensors(src), Tensors(dst)
    errors, problems = [], []
    for name in sorted(before.names()):
        if name not in after.names():
            problems.append(f"missing: {name}")
            continue
        old, new = before.get(name), after.get(name)
        if new.dtype == torch.float8_e4m3fn:
            rows, cols = new.shape
            scale = after.get(name + "_scale").float()[:, None, :, None]
            restored = (new.float().reshape(rows // 128, 128, cols // 128, 128) * scale).reshape(rows, cols)
            errors.append(((restored - old.float()).norm() / old.float().norm()).item())
        elif new.dtype != old.dtype or not torch.equal(new, old):
            problems.append(f"changed: {name}")
    allowed = before.names() | {name + "_scale" for name in before.names()}
    problems += [f"unexpected: {name}" for name in sorted(after.names() - allowed)]
    if len(errors) != EXPECTED_FP8:
        problems.append(f"{len(errors)} FP8 weights, expected {EXPECTED_FP8}")
    return errors, problems


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--src", required=True, type=Path, help="Cloudflare's clef-flash directory.")
    parser.add_argument("--dst", required=True, type=Path, help="New, empty output directory.")
    args = parser.parse_args()
    src, dst = args.src, args.dst

    print(f"Checking the input files against {SOURCE}")
    problems = check_inputs(src)
    if problems:
        print("\n".join(problems))
        sys.exit("The input files are not Cloudflare's published files. Nothing was written.")
    if dst.exists() and any(dst.iterdir()):
        sys.exit(f"{dst} is not empty.")
    dst.mkdir(parents=True, exist_ok=True)

    print("Quantizing with llm-compressor (FP8_BLOCK)")
    quantize(src, dst)

    quantization_config = json.loads((dst / "config.json").read_text())["quantization_config"]
    config = json.loads((src / "config.json").read_text())
    config["quantization_config"] = quantization_config
    (dst / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    for name in COPY:
        if (src / name).is_file():
            shutil.copyfile(src / name, dst / name)

    print("Checking the output against Cloudflare's tensors")
    errors, problems = check_output(src, dst)
    if problems:
        print("\n".join(problems))
        sys.exit(f"The output in {dst} is not valid. Do not use it.")

    # Apache-2.0 asks for a notice on changed files.
    (dst / "QUANTIZATION.md").write_text(
        f"This directory is a modified copy of {SOURCE} (Apache-2.0).\n\n"
        "Changed files: the model-*.safetensors shards, model.safetensors.index.json and config.json.\n"
        "quantize_fp8.py converted the MLP, full-attention and large linear-attention weights of the\n"
        "language model to FP8 with llm-compressor (FP8_BLOCK: one scale for each 128x128 block).\n"
        "All other tensors are bit-identical to the original. MANIFEST.json lists the hash of each file.\n"
    )

    import compressed_tensors
    import llmcompressor
    import transformers

    manifest = {
        "source": SOURCE,
        "script_sha256": sha256(Path(__file__)),
        "versions": {
            "llmcompressor": llmcompressor.__version__,
            "compressed_tensors": compressed_tensors.__version__,
            "transformers": transformers.__version__,
            "torch": torch.__version__,
        },
        "recipe": {"targets": "Linear", "scheme": "FP8_BLOCK", "ignore": IGNORE},
        "fp8_weights": len(errors),
        "relative_error_mean": round(sum(errors) / len(errors), 6),
        "relative_error_max": round(max(errors), 6),
        "files": {p.name: sha256(p) for p in sorted(dst.iterdir()) if p.is_file() and p.name != "MANIFEST.json"},
    }
    (dst / "MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n")

    size = sum(p.stat().st_size for p in dst.glob("model-*.safetensors")) / 2**30
    print(
        f"\nOK: {len(errors)} weights in FP8 (relative error mean {manifest['relative_error_mean']:.2%}, "
        f"max {manifest['relative_error_max']:.2%}). All other tensors are bit-identical.\n"
        f"Weights: {size:.2f} GiB. Output: {dst}"
    )


if __name__ == "__main__":
    main()