first commit

convert_deepseek_fp8_to_int8.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python3
+"""DeepSeek R1/V3 block-FP8 -> Kunlun compressed-tensors dynamic W8A8.
+
+CPU only; no model instantiation, calibration, or accelerator allocation.
+Checkpoint weight_scale = absmax / 127 (NOT the Kunlun runtime's absmax).
+MLA kv_b_proj, dense MLPs and shared experts are dequantized to BF16;
+routed experts and remaining MLA projections use INT8. MTP is omitted.
+"""
+import argparse
+from contextlib import ExitStack
+import hashlib
+import json
+import math
+from pathlib import Path
+import re
+import shutil
+import struct
+import time
+
+VERSION = 1
+PROJECTION = re.compile(
+    r"model\.layers\.\d+\.(?:self_attn\.(?:q_a_proj|q_b_proj|kv_a_proj_with_mqa|kv_b_proj|o_proj)|"
+    r"mlp\.(?:(?:experts\.\d+|shared_experts)\.)?(?:gate_proj|up_proj|down_proj))\.weight$"
+)
+DTYPE_BYTES = {"F8_E4M3": 1, "BF16": 2, "F16": 2, "F32": 4, "I8": 1}
+
+
+def read_json(path):
+    return json.loads(path.read_text())
+
+
+def write_json(path, value):
+    tmp = path.with_name(path.name + ".tmp")
+    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
+    tmp.replace(path)
+
+
+def read_header(path):
+    with path.open("rb") as stream:
+        raw = stream.read(8)
+        if len(raw) != 8:
+            raise ValueError(f"Truncated safetensors: {path}")
+        size = struct.unpack("<Q", raw)[0]
+        if size > 100 * 1024 * 1024:
+            raise ValueError(f"Invalid safetensors header: {path}")
+        header = json.loads(stream.read(size))
+    header.pop("__metadata__", None)
+    return header
+
+
+def quantization_config():
+    return {
+        "quant_method": "compressed-tensors",
+        "format": "int-quantized",
+        "quantization_status": "compressed",
+        "config_groups": {"group_0": {
+            "targets": ["Linear"],
+            "weights": {"num_bits": 8, "type": "int", "symmetric": True,
+                        "strategy": "channel", "dynamic": False},
+            "input_activations": {"num_bits": 8, "type": "int", "symmetric": True,
+                                  "strategy": "token", "dynamic": True},
+            "output_activations": None,
+        }},
+        "ignore": ["lm_head", "re:.*\\.gate$", "re:.*\\.self_attn\\.kv_b_proj$",
+                   "re:.*\\.mlp\\.(gate_up_proj|gate_proj|up_proj|down_proj)$",
+                   "re:.*\\.mlp\\.shared_experts\\.(gate_up_proj|gate_proj|up_proj|down_proj)$"],
+        "kv_cache_scheme": None,
+    }
+
+
+def plan_source(source):
+    config = read_json(source / "config.json")
+    quant = config.get("quantization_config", {})
+    if (config.get("model_type") != "deepseek_v3"
+            or quant.get("quant_method") != "fp8"
+            or quant.get("fmt") not in {"e4m3", "float8_e4m3fn"}
+            or quant.get("activation_scheme") != "dynamic"):
+        raise ValueError("Expected DeepSeek R1/V3 dynamic block FP8 E4M3 checkpoint")
+    if config.get("moe_layer_freq", 1) != 1 or not config.get("q_lora_rank"):
+        raise ValueError("Only DeepSeek V3/R1 with Q-LoRA and moe_layer_freq=1 is supported")
+    block = quant.get("weight_block_size")
+    if block != [128, 128]:
+        raise ValueError(f"Expected FP8 block size [128, 128], got {block}")
+    index_path = source / "model.safetensors.index.json"
+    index = read_json(index_path)
+    mapping = index["weight_map"]
+    specs = {}
+    source_files = sorted(set(mapping.values()))
+    digest = hashlib.sha256()
+    for path in (source / "config.json", index_path):
+        digest.update(path.read_bytes())
+    for filename in source_files:
+        path = (source / filename).resolve()
+        if not path.is_relative_to(source):
+            raise ValueError(f"Shard outside source: {filename}")
+        header = read_header(path)
+        stat = path.stat()
+        digest.update(f"{filename}:{stat.st_size}:{stat.st_mtime_ns}".encode())
+        for key, spec in header.items():
+            if mapping.get(key) == filename:
+                specs[key] = spec
+    if set(specs) != set(mapping):
+        raise ValueError("Source index contains missing tensors")
+    layers = config["num_hidden_layers"]
+    mtp_layers = config.get("num_nextn_predict_layers", 0)
+    active = {}
+    for key, spec in specs.items():
+        match = re.match(r"model\.layers\.(\d+)\.", key)
+        if match and int(match[1]) >= layers:
+            if int(match[1]) >= layers + mtp_layers:
+                raise ValueError(f"Unexpected layer outside decoder/MTP range: {key}")
+            continue
+        active[key] = spec
+    specs = active
+    expected_shapes = expected_projection_shapes(config)
+    actual_projections = {key for key in specs if PROJECTION.fullmatch(key)}
+    if actual_projections != set(expected_shapes):
+        missing = sorted(set(expected_shapes) - actual_projections)[:5]
+        extra = sorted(actual_projections - set(expected_shapes))[:5]
+        raise ValueError(f"Projection coverage mismatch: missing={missing}, extra={extra}")
+    scales = set()
+    fp8_count = 0
+    output_bytes = 0
+    for key, spec in specs.items():
+        dtype, shape = spec["dtype"], spec["shape"]
+        if dtype == "F8_E4M3":
+            if not PROJECTION.fullmatch(key) or len(shape) != 2:
+                raise ValueError(f"Unexpected FP8 tensor: {key} {shape}")
+            if shape != expected_shapes[key]:
+                raise ValueError(f"Projection shape mismatch: {key}: {shape}, expected {expected_shapes[key]}")
+            scale_key = key.removesuffix("weight") + "weight_scale_inv"
+            expected = [math.ceil(shape[0] / block[0]), math.ceil(shape[1] / block[1])]
+            scale_spec = specs.get(scale_key, {})
+            if scale_spec.get("shape") != expected or scale_spec.get("dtype") != "F32":
+                raise ValueError(f"Missing/invalid block scale: {scale_key}")
+            scales.add(scale_key)
+            fp8_count += 1
+            output_bytes += (math.prod(shape) * 2 if is_bf16_projection(key)
+                             else math.prod(shape) + shape[0] * 4)
+        elif key.endswith("weight_scale_inv"):
+            continue
+        else:
+            if PROJECTION.fullmatch(key):
+                raise ValueError(f"Expected FP8 projection: {key}, got {dtype}")
+            if dtype not in DTYPE_BYTES:
+                raise ValueError(f"Unsupported source dtype: {key}: {dtype}")
+            if len(shape) == 2 and key not in {"model.embed_tokens.weight", "lm_head.weight"} and not re.fullmatch(r"model\.layers\.\d+\.mlp\.gate\.weight", key):
+                raise ValueError(f"Unexpected unquantized matrix: {key}")
+            output_bytes += math.prod(shape) * DTYPE_BYTES[dtype]
+    if scales != {key for key in specs if key.endswith("weight_scale_inv")}:
+        raise ValueError("Orphan block scales in checkpoint")
+    expected_count = len(expected_shapes)
+    if fp8_count != expected_count:
+        raise ValueError(f"Expected {expected_count} FP8 projections, found {fp8_count}")
+    groups = {}
+    for key in sorted(specs):
+        if key not in scales:
+            groups.setdefault(mapping[key], []).append(key)
+    groups = sorted(groups.items())
+    config["num_nextn_predict_layers"] = 0
+    return config, mapping, specs, groups, digest.hexdigest(), output_bytes, fp8_count
+
+
+def is_bf16_projection(key):
+    return (key.endswith(".self_attn.kv_b_proj.weight") or bool(re.fullmatch(
+        r"model\.layers\.\d+\.mlp\.(?:shared_experts\.)?(?:gate_proj|up_proj|down_proj)\.weight", key)))
+
+
+def expected_projection_shapes(config):
+    hidden = config["hidden_size"]
+    heads = config["num_attention_heads"]
+    q_rank, kv_rank = config["q_lora_rank"], config["kv_lora_rank"]
+    q_head = config["qk_nope_head_dim"] + config["qk_rope_head_dim"]
+    result = {}
+    for index in range(config["num_hidden_layers"]):
+        prefix = f"model.layers.{index}."
+        for name, shape in {
+            "q_a_proj": [q_rank, hidden], "q_b_proj": [heads * q_head, q_rank],
+            "kv_a_proj_with_mqa": [kv_rank + config["qk_rope_head_dim"], hidden],
+            "kv_b_proj": [heads * (config["qk_nope_head_dim"] + config["v_head_dim"]), kv_rank],
+            "o_proj": [hidden, heads * config["v_head_dim"]],
+        }.items():
+            result[prefix + "self_attn." + name + ".weight"] = shape
+        if index < config["first_k_dense_replace"]:
+            mlps = [("mlp", config["intermediate_size"])]
+        else:
+            mlps = [(f"mlp.experts.{e}", config["moe_intermediate_size"])
+                    for e in range(config["n_routed_experts"])]
+            if config.get("n_shared_experts"):
+                mlps.append(("mlp.shared_experts", config["moe_intermediate_size"] * config["n_shared_experts"]))
+        for mlp, intermediate in mlps:
+            for name, shape in {"gate_proj": [intermediate, hidden], "up_proj": [intermediate, hidden],
+                                "down_proj": [hidden, intermediate]}.items():
+                result[prefix + mlp + "." + name + ".weight"] = shape
+    return result
+
+
+def dequantize_bf16(weight, block_scale, chunk_rows=1024):
+    import torch
+    rows, cols = weight.shape
+    result = torch.empty((rows, cols), dtype=torch.bfloat16, device="cpu")
+    columns = torch.arange(cols) // 128
+    if not torch.isfinite(block_scale).all() or (block_scale < 0).any():
+        raise ValueError("Invalid block scales")
+    for start in range(0, rows, chunk_rows):
+        end = min(rows, start + chunk_rows)
+        row_blocks = torch.arange(start, end) // 128
+        value = weight[start:end].float() * block_scale[row_blocks[:, None], columns[None, :]]
+        value = value.to(torch.bfloat16)
+        if not torch.isfinite(value).all():
+            raise ValueError("Non-finite dequantized BF16 weight")
+        result[start:end] = value
+    return result
+
+
+def quantize_weight(weight, block_scale, block=(128, 128), chunk_rows=1024):
+    """Recover W=FP8*block_scale in FP32, then round(W/(row_absmax/127))."""
+    import torch
+    if weight.ndim != 2 or weight.dtype != torch.float8_e4m3fn:
+        raise ValueError("Expected 2-D FP8 E4M3 weight")
+    rows, cols = weight.shape
+    if tuple(block_scale.shape) != (math.ceil(rows / block[0]), math.ceil(cols / block[1])):
+        raise ValueError("Block scale shape mismatch")
+    if chunk_rows < 1 or not torch.isfinite(block_scale).all() or (block_scale < 0).any():
+        raise ValueError("Invalid chunk size or block scales")
+    result = torch.empty((rows, cols), dtype=torch.int8, device="cpu")
+    scales = torch.empty((rows, 1), dtype=torch.float32, device="cpu")
+    col_blocks = torch.arange(cols, device="cpu") // block[1]
+    for start in range(0, rows, chunk_rows):
+        end = min(rows, start + chunk_rows)
+        row_blocks = torch.arange(start, end, device="cpu") // block[0]
+        value = weight[start:end].to(torch.float32)
+        value.mul_(block_scale[row_blocks[:, None], col_blocks[None, :]])
+        if not torch.isfinite(value).all():
+            raise ValueError("Non-finite dequantized FP8 weight")
+        scale = value.abs().amax(dim=1, keepdim=True).div_(127.0)
+        # A zero row uses scale=1, q=0. Avoid zero scale and underflow.
+        scale.clamp_(min=torch.finfo(torch.float32).tiny)
+        scale = torch.where(value.abs().amax(dim=1, keepdim=True) == 0, 1.0, scale)
+        result[start:end] = value.div_(scale).round_().clamp_(-127, 127).to(torch.int8)
+        scales[start:end] = scale
+    return result, scales
+
+
+def output_specs(keys, specs):
+    expected = {}
+    for key in keys:
+        spec = specs[key]
+        dtype = spec["dtype"]
+        if dtype == "F8_E4M3":
+            dtype = "BF16" if is_bf16_projection(key) else "I8"
+        expected[key] = (dtype, spec["shape"])
+        if spec["dtype"] == "F8_E4M3" and not is_bf16_projection(key):
+            expected[key.removesuffix("weight") + "weight_scale"] = ("F32", [spec["shape"][0], 1])
+    return expected
+
+
+def validate_shard(path, expected):
+    header = read_header(path)
+    if set(header) != set(expected):
+        raise ValueError(f"Tensor names mismatch: {path}")
+    for key, (dtype, shape) in expected.items():
+        if header[key]["dtype"] != dtype or header[key]["shape"] != shape:
+            raise ValueError(f"Tensor dtype/shape mismatch: {path}: {key}")
+    # safe_open validates offsets and file length without reading tensor data.
+    from safetensors import safe_open
+    with safe_open(str(path), framework="pt", device="cpu"):
+        pass
+
+
+def copy_assets(source, destination):
+    excluded = {"config.json", "model.safetensors.index.json", "quantize_config.json",
+                "quantization_config.json", "conversion_state.json", "conversion_report.json"}
+    for path in source.iterdir():
+        if (path.is_file() and path.name not in excluded and not path.name.endswith(".index.json")
+                and path.suffix in {".json", ".py", ".jinja", ".txt", ".model", ".tiktoken"}):
+            shutil.copy2(path, destination / path.name)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--source", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--threads", type=int, default=8)
+    parser.add_argument("--chunk-rows", type=int, default=1024)
+    parser.add_argument("--dry-run", action="store_true", help="Check headers, estimate space; write nothing")
+    parser.add_argument("--resume", action="store_true", help="Resume matching conversion_state.json")
+    args = parser.parse_args()
+    if args.threads < 1 or args.chunk_rows < 1:
+        parser.error("threads and chunk-rows must be positive")
+    source, output = args.source.resolve(), args.output.resolve()
+    if source == output or source in output.parents or output in source.parents:
+        parser.error("source and output must be separate, non-nested directories")
+    config, mapping, specs, groups, fingerprint, size, count = plan_source(source)
+    bf16_count = sum(1 for key, spec in specs.items()
+                     if spec["dtype"] == "F8_E4M3" and is_bf16_projection(key))
+    ancestor = output
+    while not ancestor.exists():
+        ancestor = ancestor.parent
+    free = shutil.disk_usage(ancestor).free
+    print(json.dumps({"fp8_weights": count, "output_shards": len(groups),
+                      "int8_weights": count - bf16_count,
+                      "bf16_weights": bf16_count,
+                      "bf16_policy": "MLA kv_b_proj, dense MLPs, shared experts",
+                      "mtp": "omitted; main-model inference only",
+                      "estimated_output_GiB": round(size / 2**30, 2),
+                      "free_disk_GiB": round(free / 2**30, 2)}, indent=2), flush=True)
+    if args.dry_run:
+        return
+    import torch
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+    torch.set_num_threads(args.threads)
+    state_path = output / "conversion_state.json"
+    identity = {"version": VERSION, "source": str(source), "fingerprint": fingerprint,
+                "chunk_rows": args.chunk_rows, "quantization_config": quantization_config()}
+    if args.resume:
+        state = read_json(state_path)
+        if state["identity"] != identity:
+            raise ValueError("Source/settings changed; use a new output directory")
+    else:
+        if output.exists() and any(output.iterdir()):
+            raise ValueError("Output must be empty; use --resume for an interrupted conversion")
+        state = {"identity": identity, "completed": {}}
+    completed_bytes = sum((output / name).stat().st_size for name in state["completed"])
+    if free < max(0, size - completed_bytes) + 2 * 2**30:
+        raise ValueError("Insufficient output disk space (including 2 GiB reserve)")
+    output.mkdir(parents=True, exist_ok=True)
+    write_json(state_path, state)
+    output_mapping = {}
+    total_size = 0
+    started = time.monotonic()
+    for number, (filename, keys) in enumerate(groups, 1):
+        name = f"model-{number:05d}-of-{len(groups):05d}.safetensors"
+        path = output / name
+        expected = output_specs(keys, specs)
+        if name in state["completed"]:
+            if path.stat().st_size != state["completed"][name]:
+                raise ValueError(f"Completed shard changed: {path}")
+            validate_shard(path, expected)
+            print(f"[{number}/{len(groups)}] resume: {name}", flush=True)
+        else:
+            tensors = {}
+            with ExitStack() as stack:
+                readers = {}
+
+                def get(key):
+                    shard = mapping[key]
+                    if shard not in readers:
+                        readers[shard] = stack.enter_context(safe_open(
+                            str(source / shard), framework="pt", device="cpu"))
+                    return readers[shard].get_tensor(key)
+
+                with torch.inference_mode():
+                    for key in keys:
+                        if specs[key]["dtype"] == "F8_E4M3":
+                            scale_key = key.removesuffix("weight") + "weight_scale_inv"
+                            if is_bf16_projection(key):
+                                tensors[key] = dequantize_bf16(get(key), get(scale_key), chunk_rows=args.chunk_rows)
+                            else:
+                                tensors[key], tensors[key.removesuffix("weight") + "weight_scale"] = (
+                                    quantize_weight(get(key), get(scale_key), chunk_rows=args.chunk_rows))
+                        else:
+                            tensors[key] = get(key).contiguous()
+                tmp = path.with_name(path.name + ".tmp")
+                save_file(tensors, str(tmp), metadata={"format": "pt"})
+                validate_shard(tmp, expected)
+                tmp.replace(path)
+                del tensors
+            state["completed"][name] = path.stat().st_size
+            write_json(state_path, state)
+            print(f"[{number}/{len(groups)}] wrote {name}; elapsed={time.monotonic()-started:.1f}s", flush=True)
+        output_mapping.update(dict.fromkeys(expected, name))
+        total_size += sum(math.prod(shape) * DTYPE_BYTES[dtype] for dtype, shape in expected.values())
+    copy_assets(source, output)
+    config["quantization_config"] = quantization_config()
+    write_json(output / "config.json", config)
+    write_json(output / "model.safetensors.index.json", {
+        "metadata": {"total_size": total_size}, "weight_map": output_mapping})
+    write_json(output / "conversion_report.json", {
+        "source": str(source), "source_fingerprint": fingerprint,
+        "fp8_weights_processed": count, "int8_weights": count - bf16_count,
+        "bf16_weights": bf16_count, "bf16_policy": "MLA kv_b_proj, dense MLPs, shared experts",
+        "mtp": "omitted",
+        "output_tensors": len(output_mapping),
+        "tensor_bytes": total_size, "seconds_this_run": time.monotonic() - started,
+        "method": "FP8 * block_scale -> FP32 -> per-output-channel INT8 or BF16 (see bf16_policy)",
+        "calibration": "none; dynamic per-token activation quantization at inference",
+    })
+    print(f"Complete: {output}", flush=True)
+
+
+if __name__ == "__main__":
+    main()