commit 2794a50a38208f3b6b876a3adb2ceb3e2231d5f2 Author: Sun Ruoxi Date: Thu Oct 29 16:23:16 2026 +0800 first commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..6ed471f --- /dev/null +++ b/README.md @@ -0,0 +1,78 @@ +# MiniMax-M2.7 FP8 转昆仑芯 INT8 + +依据安装版本： + +- `vllm==0.25.1`，`vllm-kunlun==0.25.1.dev0`，仓库提交 `fbb6fa6`。 +- Linear：`KunlunScaledMMLinearKernel`。 +- MoE：`KunlunCompressedTensorsW8A8Int8MoEMethod`；保留专家 `w1/w2/w3` 的原始名称，由 vLLM 完成融合及 TP 切分。 +- 输出 `compressed-tensors / int-quantized`：逐输出通道对称 INT8 权重，逐 token 动态对称 INT8 激活，无需校准数据集。 + +转换关系： + +```text +W_fp32[i,j] = W_fp8[i,j] * weight_scale_inv[i//128,j//128] +s[i,0] = max(abs(W_fp32[i,:])) / 127 +W_int8[i,j] = clamp(round(W_fp32[i,j] / s[i,0]), -127, 127) +``` + +尽管源 scale 名字有 `_inv`，其保存值用于乘法反量化，不能取倒数。全零行输出 `q=0,s=1`。 +输出 scale 名称为 `weight_scale`、FP32 类型、形状 `[out_features,1]`。 +磁盘保存 `absmax/127`；插件加载后自行乘 127 得到运行时 absmax，不能在转换时提前乘。 + +Embedding、lm_head、router gate、routing bias 和 norm 保留原始类型和值。脚本读取索引指定的张量，不会把目录里未被索引引用的旧分片一起导入。支持 scale 跨分片读取，分片内逐矩阵处理、矩阵内分行处理；不把完整模型加载进内存，不分配 XPU。 + +## 使用 + +脚本部署位置： +`/mnt/disk0/models/MiniMax/quant_tools/minimax_m27_kunlun_int8/`。 + +预检只读取配置及所有索引分片的文件头，不写输出： + +```bash +/opt/vllm_kunlun/bin/python \ + /mnt/data0/models/MiniMax/quant_tools/minimax_m27_kunlun_int8/convert_fp8_to_int8.py \ + --source /mnt/data0/models/MiniMax/MiniMax-M2.7 \ + --output /mnt/data0/models/MiniMax/MiniMax-M2.7-INT8 \ + --dry-run +``` + +全量转换： + +```bash +/opt/vllm_kunlun/bin/python \ + /mnt/data0/models/MiniMax/quant_tools/minimax_m27_kunlun_int8/convert_fp8_to_int8.py \ + --source /mnt/data0/models/MiniMax/MiniMax-M2.7 \ + --output /mnt/data0/models/MiniMax/MiniMax-M2.7-INT8 \ + --threads 8 +``` + +中断后用同一命令添加 `--resume`。已完成的分片会检查文件大小、张量名称、类型及形状后跳过；未完成分片重新计算。`--chunk-rows` 默认为 1024，如需调整，恢复时必须保持一致。源文件位置、大小、时间戳、模型配置或索引变化后会拒绝恢复。恢复检查不包含整文件数据校验和。 + +预检结果为 47,864 个 FP8 矩阵、125 个输出分片，预计张量数据 **214.63 GiB**；检查时磁盘剩余 352.09 GiB。转换需要独立输出目录，原始模型保持不变。每完成一个分片都会更新 `conversion_state.json`，最后生成新的配置、索引和 `conversion_report.json`。只有打印 `Complete` 后才启动模型。 + +转换完成后的首次服务验证可运行： + +```bash +bash /mnt/disk0/models/MiniMax/quant_tools/minimax_m27_kunlun_int8/serve_int8.sh +``` + +启动脚本使用 TP=4、BF16、8192 上下文和 eager 模式，先控制 KV cache 和图编译相关变量。服务启动及生成通过后再调整上下文和并发。 + +## 验证 + +```bash +cd /mnt/data0/models/MiniMax/quant_tools/minimax_m27_kunlun_int8 +/opt/vllm_kunlun/bin/python test_conversion.py +/opt/vllm_kunlun/bin/python validate_kunlun_contract.py +/opt/vllm_kunlun/bin/python validate_sample_weights.py \ + --source /mnt/data0/models/MiniMax/MiniMax-M2.7 \ + --report sample_validation.json +``` + +- `test_conversion.py`：独立反量化参考值与半个量化步长误差界、全零行、非整块尺寸、非法 scale、跨分片转换、配置与索引、断点恢复、源文件不变及拒绝覆盖。 +- `validate_kunlun_contract.py`：使用小型 CPU 测试层、模拟 TP rank/world size，核实当前插件的配置匹配、Linear/MoE 内核选择、参数布局、scale 转换和忽略名单；不代表完整引擎加载或设备算子执行通过。 +- `validate_sample_weights.py`：真实 layer 0 的 Q/K/V/O 和 expert 0 的 w1/w2/w3，共七个矩阵，检查转换后的逐行舍入误差并记录 relative RMSE。 + +2026-10-29 完成上述三项验证，全部通过。七个真实矩阵相对于 FP8 乘 scale 后参考权重的 relative RMSE 为 0.85%–1.08%，全部满足半个逐行量化步长的舍入误差界。该指标仅衡量这些矩阵的二次量化误差。 + +这是从现有 FP8 权重再次量化得到的 absmax INT8 基线，并非从原始高精度权重直接量化。数学误差检查不代表任务质量评测；完整模型转换、P800 推理及任务精度仍需分别验证。 diff --git a/convert_fp8_to_int8.py b/convert_fp8_to_int8.py new file mode 100644 index 0000000..9ae23b9 --- /dev/null +++ b/convert_fp8_to_int8.py @@ -0,0 +1,298 @@ +#!/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(" 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() diff --git a/serve_int8.sh b/serve_int8.sh new file mode 100644 index 0000000..a6891d9 --- /dev/null +++ b/serve_int8.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Run inside srx_test after conversion finishes. +set -euo pipefail +export XPU_VISIBLE_DEVICES=0,1,2,3 +export CUDA_VISIBLE_DEVICES=0,1,2,3 +exec /opt/vllm_kunlun/bin/vllm serve \ + /mnt/data0/models/MiniMax/MiniMax-M2.7-INT8 \ + --trust-remote-code \ + --quantization compressed-tensors \ + --dtype bfloat16 \ + --tensor-parallel-size 4 \ + --distributed-executor-backend mp \ + --max-model-len 8192 \ + --max-num-seqs 4 \ + --max-num-batched-tokens 8192 \ + --block-size 128 \ + --enforce-eager \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --tool-call-parser minimax_m2 \ + --enable-auto-tool-choice \ + --reasoning-parser minimax_m2 \ + "$@" diff --git a/test_conversion.py b/test_conversion.py new file mode 100644 index 0000000..529df18 --- /dev/null +++ b/test_conversion.py @@ -0,0 +1,91 @@ +"""CPU numerical and cross-shard checkpoint integration tests.""" +import hashlib +import json +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from convert_fp8_to_int8 import quantize_weight, quantization_config + +SCRIPT = Path(__file__).with_name("convert_fp8_to_int8.py") + + +class ConversionTests(unittest.TestCase): + def test_quantization_error_and_partial_blocks(self): + torch.manual_seed(42) + w = (torch.randn(129, 257) * 5).to(torch.float8_e4m3fn) + w[0] = torch.zeros(257).to(torch.float8_e4m3fn) + block_scale = torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.6, 0.7]]) + q, scale = quantize_weight(w, block_scale, chunk_rows=37) + ref = w.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:129, :257] + self.assertEqual(q.dtype, torch.int8) + self.assertEqual(tuple(scale.shape), (129, 1)) + self.assertTrue(torch.all((q.float() * scale - ref).abs() <= scale / 2 + 1e-5)) + self.assertEqual(scale[0].item(), 1.0) + self.assertEqual(q[0].count_nonzero().item(), 0) + self.assertTrue(torch.allclose(scale[1:], ref.abs().amax(1, keepdim=True)[1:] / 127)) + with self.assertRaises(ValueError): + quantize_weight(w, torch.full((2, 3), float("nan"))) + + def test_cross_shard_conversion_and_resume(self): + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) + src, dst = root / "source", root / "out" + src.mkdir() + config = {"model_type": "minimax_m2", "num_hidden_layers": 1, + "num_local_experts": 1, "quantization_config": { + "quant_method": "fp8", "fmt": "float8_e4m3fn", + "activation_scheme": "dynamic", "weight_block_size": [128, 128]}} + (src / "config.json").write_text(json.dumps(config)) + (src / "tokenizer_config.json").write_text('{}') + modules = [f"self_attn.{p}_proj" for p in "qkvo"] + modules += [f"block_sparse_moe.experts.0.w{i}" for i in (1, 2, 3)] + tensors, scales = {}, {} + for module in modules: + key = f"model.layers.0.{module}.weight" + tensors[key] = torch.arange(24).reshape(4, 6).float().to(torch.float8_e4m3fn) + scales[key.removesuffix("weight") + "weight_scale_inv"] = torch.tensor([[0.25]]) + tensors["lm_head.weight"] = torch.ones(4, 6, dtype=torch.bfloat16) + tensors["model.layers.0.block_sparse_moe.gate.weight"] = torch.ones(1, 6) + save_file(tensors, str(src / "weights.safetensors")) + save_file(scales, str(src / "scales.safetensors")) + mapping = {k: "weights.safetensors" for k in tensors} + mapping.update({k: "scales.safetensors" for k in scales}) + (src / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) + before = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in src.iterdir()} + command = [sys.executable, str(SCRIPT), "--source", str(src), "--output", str(dst), "--threads", "1"] + subprocess.run(command + ["--dry-run"], check=True) + self.assertFalse(dst.exists()) + subprocess.run(command, check=True) + out_map = json.loads((dst / "model.safetensors.index.json").read_text())["weight_map"] + self.assertFalse(any(k.endswith("weight_scale_inv") for k in out_map)) + self.assertEqual(len(out_map), len(tensors) + 7) + for key, filename in out_map.items(): + with safe_open(str(dst / filename), framework="pt") as reader: + value = reader.get_tensor(key) + if key.endswith("weight_scale"): + self.assertEqual(tuple(value.shape), (4, 1)) + self.assertEqual(value.dtype, torch.float32) + elif key in tensors and "proj.weight" not in key and ".experts." not in key: + self.assertTrue(torch.equal(value, tensors[key])) + else: + self.assertEqual(value.dtype, torch.int8) + self.assertEqual(json.loads((dst / "config.json").read_text())["quantization_config"], quantization_config()) + self.assertTrue((dst / "tokenizer_config.json").exists()) + for filename in ("config.json", "model.safetensors.index.json", "conversion_report.json"): + (dst / filename).unlink() + subprocess.run(command + ["--resume"], check=True) + self.assertTrue((dst / "model.safetensors.index.json").exists()) + self.assertEqual(before, {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in src.iterdir()}) + self.assertNotEqual(subprocess.run(command, capture_output=True).returncode, 0) + + +if __name__ == "__main__": + torch.set_num_threads(1) + unittest.main() diff --git a/validate_kunlun_contract.py b/validate_kunlun_contract.py new file mode 100644 index 0000000..24b8fad --- /dev/null +++ b/validate_kunlun_contract.py @@ -0,0 +1,62 @@ +"""Check installed Kunlun loader's configuration and CPU parameter contracts. + +This uses synthetic layers and a mocked TP rank, not a full engine startup. +""" +import copy +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from convert_fp8_to_int8 import quantization_config + + +def main(): + from vllm.platforms import current_platform + current_platform.pre_register_and_update() + from vllm_kunlun.quantization.compressed_tensors.compressed_tensors import KunlunCompressedTensorsConfig + from vllm_kunlun.quantization.compressed_tensors.compressed_tensors_moe import KunlunCompressedTensorsMoEMethod + + cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config())) + cfg.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + linear = torch.nn.Linear(128, 128, bias=False) + scheme = cfg.get_scheme(linear, "model.layers.0.self_attn.qkv_proj") + # CPU-only parameter allocation; no distributed processes or accelerator work. + del linear._parameters["weight"] + with patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), \ + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1): + scheme.create_weights(linear, [128], 128, torch.bfloat16, lambda *a, **k: None) + assert type(scheme.kernel).__name__ == "KunlunScaledMMLinearKernel" + assert linear.weight.dtype == torch.int8 + assert tuple(linear.weight_scale.shape) == (128, 1) + print("LINEAR_PASS", type(scheme.kernel).__name__) + + class RoutedExperts(torch.nn.Module): + pass + + experts = RoutedExperts() + experts.ckpt_gate_proj_name = "w1" + experts.ckpt_down_proj_name = "w2" + experts.ckpt_up_proj_name = "w3" + experts.moe_config = SimpleNamespace(is_act_and_mul=True) + method = KunlunCompressedTensorsMoEMethod.get_moe_method( + cfg, experts, "model.layers.0.block_sparse_moe.experts") + assert type(method).__name__ == "KunlunCompressedTensorsW8A8Int8MoEMethod" + method.create_weights(experts, 2, 128, 128, torch.bfloat16) + assert tuple(experts.w13_weight.shape) == (2, 256, 128) + assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1) + assert tuple(experts.w2_weight_scale.shape) == (2, 128, 1) + with torch.no_grad(): + experts.w13_weight_scale.fill_(0.25) + experts.w2_weight_scale.fill_(0.25) + method.process_weights_after_loading(experts) + assert experts.w13_weight_scale.flatten()[0].item() == 31.75 + print("MOE_PASS", type(method).__name__) + print("SCALE_PASS: checkpoint scale * 127 = runtime absmax") + assert cfg.get_scheme(torch.nn.Linear(128, 128), "lm_head") is None + assert cfg.get_scheme(torch.nn.Linear(128, 128), "model.layers.0.block_sparse_moe.gate") is None + print("IGNORE_PASS: lm_head and router gate") + + +if __name__ == "__main__": + main() diff --git a/validate_sample_weights.py b/validate_sample_weights.py new file mode 100644 index 0000000..7c2af00 --- /dev/null +++ b/validate_sample_weights.py @@ -0,0 +1,45 @@ +"""Numerical check on real layer-0 attention and one expert; does not save a model.""" +import argparse +import json +from pathlib import Path + +import torch +from safetensors import safe_open + +from convert_fp8_to_int8 import quantize_weight + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", required=True, type=Path) + parser.add_argument("--report", required=True, type=Path) + args = parser.parse_args() + torch.set_num_threads(8) + mapping = json.loads((args.source / "model.safetensors.index.json").read_text())["weight_map"] + names = [f"model.layers.0.self_attn.{p}_proj.weight" for p in "qkvo"] + names += [f"model.layers.0.block_sparse_moe.experts.0.w{i}.weight" for i in (1, 2, 3)] + report = [] + for name in names: + scale_name = name.removesuffix("weight") + "weight_scale_inv" + with safe_open(str(args.source / mapping[name]), framework="pt") as reader: + weight = reader.get_tensor(name) + with safe_open(str(args.source / mapping[scale_name]), framework="pt") as reader: + block_scale = reader.get_tensor(scale_name) + with torch.inference_mode(): + q, scales = quantize_weight(weight, block_scale) + # Independent reference via explicit block expansion. + reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:weight.shape[0], :weight.shape[1]] + error = q.float() * scales - reference + bound = scales / 2 + scales * 1e-4 + 1e-7 + assert torch.all(error.abs() <= bound), name + relative_rmse = (error.square().mean() / reference.square().mean()).sqrt().item() + report.append({"name": name, "shape": list(weight.shape), + "relative_rmse": relative_rmse, + "max_absolute_error": error.abs().max().item(), + "rounding_error_bound": "passed"}) + print(json.dumps(report[-1]), flush=True) + args.report.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main()