first commit

This commit is contained in:
2026-10-29 16:23:16 +08:00
commit 2794a50a38
6 changed files with 597 additions and 0 deletions

78
README.md Normal file
View File

@@ -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 推理及任务精度仍需分别验证。

298
convert_fp8_to_int8.py Normal file
View File

@@ -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("<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()

23
serve_int8.sh Normal file
View File

@@ -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 \
"$@"

91
test_conversion.py Normal file
View File

@@ -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()

View File

@@ -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()

View File

@@ -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()