299 lines
13 KiB
Python
299 lines
13 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""MiniMax M2 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).
|
||
|
|
"""
|
||
|
|
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|k|v|o)_proj|"
|
||
|
|
r"block_sparse_moe\.experts\.\d+\.w[123])\.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$"],
|
||
|
|
"kv_cache_scheme": None,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def plan_source(source):
|
||
|
|
config = read_json(source / "config.json")
|
||
|
|
quant = config.get("quantization_config", {})
|
||
|
|
if (config.get("model_type") != "minimax_m2"
|
||
|
|
or quant.get("quant_method") != "fp8"
|
||
|
|
or quant.get("fmt") != "float8_e4m3fn"
|
||
|
|
or quant.get("activation_scheme") != "dynamic"):
|
||
|
|
raise ValueError("Expected MiniMax M2 dynamic block FP8 E4M3 checkpoint")
|
||
|
|
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")
|
||
|
|
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}")
|
||
|
|
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) + 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}")
|
||
|
|
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 = config["num_hidden_layers"] * (4 + config["num_local_experts"] * 3)
|
||
|
|
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())
|
||
|
|
return config, mapping, specs, groups, digest.hexdigest(), output_bytes, fp8_count
|
||
|
|
|
||
|
|
|
||
|
|
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]
|
||
|
|
expected[key] = ("I8" if spec["dtype"] == "F8_E4M3" else spec["dtype"], spec["shape"])
|
||
|
|
if spec["dtype"] == "F8_E4M3":
|
||
|
|
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)
|
||
|
|
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),
|
||
|
|
"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"
|
||
|
|
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_converted": count, "output_tensors": len(output_mapping),
|
||
|
|
"tensor_bytes": total_size, "seconds_this_run": time.monotonic() - started,
|
||
|
|
"method": "FP8 * block_scale -> FP32 -> symmetric per-output-channel INT8",
|
||
|
|
"calibration": "none; dynamic per-token activation quantization at inference",
|
||
|
|
})
|
||
|
|
print(f"Complete: {output}", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|