46 lines
2.1 KiB
Python
46 lines
2.1 KiB
Python
|
"""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()
|