56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Check the layer-0 RMSNorm fixture against its safetensors F32 calculation.
|
|
|
|
This imports only the project's safe safetensors reader; it deliberately does
|
|
not import torch or deserialize a PyTorch checkpoint.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def load_converter():
|
|
path = Path(__file__).with_name("convert_llm2vec_layer_to_gguf.py")
|
|
spec = importlib.util.spec_from_file_location("llm2vec_converter", path)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def max_abs(a: np.ndarray, b: np.ndarray) -> float:
|
|
return float(np.max(np.abs(a - b)))
|
|
|
|
|
|
def bf16(value: np.ndarray) -> np.ndarray:
|
|
bits = value.view(np.uint32)
|
|
rounded_bits = (bits + np.uint32(0x7FFF) + ((bits >> 16) & 1)) & np.uint32(0xFFFF0000)
|
|
return rounded_bits.view(np.float32)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base", type=Path, required=True)
|
|
parser.add_argument("--fixture", type=Path, required=True)
|
|
opt = parser.parse_args()
|
|
converter = load_converter()
|
|
shards = [converter.safe_file(path) for path in sorted(opt.base.glob("model-*.safetensors"))]
|
|
weight = converter.f32(converter.base_tensor(shards, "model.layers.0.input_layernorm.weight"))
|
|
x = np.fromfile(opt.fixture / "token_embeddings.f32", dtype="<f4").reshape(16, 4096)
|
|
expected = np.fromfile(opt.fixture / "layer0_input_norm.f32", dtype="<f4").reshape(16, 4096)
|
|
actual = (x * np.reciprocal(np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-5))) * weight
|
|
print(f"f32 max_abs={max_abs(actual, expected):g}")
|
|
print(f"bf16-final max_abs={max_abs(bf16(actual), expected):g}")
|
|
normalized = x * np.reciprocal(np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-5))
|
|
print(f"transformers-rmsnorm max_abs={max_abs(bf16(bf16(normalized) * weight), expected):g}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|