83 lines
3.6 KiB
Python
83 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Reference-only check of Vision-Exp weights and the official PyTorch image graph.
|
|
|
|
Production code never invokes this script. Inputs are the pinned DS4 source and
|
|
DeepSeek source revisions recorded in docs/issue100-vision, and native dumps
|
|
from the ignored issue100_encoder_real_weights test.
|
|
"""
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--ds4-source', type=Path, required=True)
|
|
p.add_argument('--official-source', type=Path, required=True)
|
|
p.add_argument('--artifacts', type=Path, required=True)
|
|
p.add_argument('--native', type=Path, required=True)
|
|
p.add_argument('--image', type=Path, required=True)
|
|
p.add_argument('--device', default='mps')
|
|
a = p.parse_args()
|
|
parser_path = a.ds4_source / 'gguf-tools/mixed/splice_mixed_expert_layers_gguf.py'
|
|
spec = importlib.util.spec_from_file_location('ds4_gguf', parser_path)
|
|
gguf = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = gguf
|
|
spec.loader.exec_module(gguf)
|
|
sys.path.insert(0, str(a.official_source / 'inference'))
|
|
from vision import ViT, Aligner
|
|
from image_processor import load_image, build_image_block
|
|
|
|
config = json.loads((a.official_source / 'config.json').read_text())
|
|
config['dim'] = config['hidden_size']
|
|
args = SimpleNamespace(**config)
|
|
file = a.artifacts / 'DeepSeek-V4-Flash-Vision-Encoder.gguf'
|
|
gguf.GGML_QUANT_SIZES[30] = (1, 2, "BF16") # GGML BF16, raw source weights.
|
|
info = gguf.parse_gguf(file)
|
|
torch.set_default_dtype(torch.bfloat16)
|
|
torch.set_default_device(a.device)
|
|
vit, aligner = ViT(args), Aligner(args)
|
|
with file.open('rb') as f:
|
|
def state(prefix):
|
|
result = {}
|
|
for name, t in info.tensor_by_name.items():
|
|
if not name.startswith(prefix):
|
|
continue
|
|
assert t.ggml_type == 30, (name, t.ggml_type)
|
|
f.seek(t.data_offset)
|
|
values = np.frombuffer(f.read(t.n_bytes), dtype='<u2').copy()
|
|
weight = torch.from_numpy(values).view(torch.bfloat16).reshape(tuple(reversed(t.dims)))
|
|
result[name.removeprefix(prefix)] = weight.to(a.device)
|
|
return result
|
|
vit.load_state_dict(state('vision.'), strict=True)
|
|
aligner.load_state_dict(state('aligner.'), strict=True)
|
|
vit.eval(); aligner.eval()
|
|
patches, h, w, lh, lw = load_image({'url': str(a.image)}, args)
|
|
meta = json.loads(a.native.with_suffix('.json').read_text())
|
|
assert meta['grid'] == [h, w]
|
|
assert meta['aligned_grid'] == [lh, lw]
|
|
with torch.inference_mode():
|
|
output = aligner(vit(patches.to(a.device), h, w), h, w).float().cpu().numpy()
|
|
native = np.fromfile(a.native.with_suffix('.embedding.f32'), dtype='<f4').reshape(output.shape)
|
|
error = native - output
|
|
result = {
|
|
'source_revision': 'e46e16bf6035c6f317eb2ac7458eb0362926d402',
|
|
'ds4_revision': '0aaea5a238fb41a35106a551e73c8409dfb751ac',
|
|
'torch': torch.__version__, 'device': a.device,
|
|
'shape': list(output.shape),
|
|
'max_abs': float(np.abs(error).max()),
|
|
'rmse': float(np.sqrt(np.mean(error**2))),
|
|
'reference_rms': float(np.sqrt(np.mean(output**2))),
|
|
'cosine': float(np.dot(native.ravel(), output.ravel()) / (np.linalg.norm(native) * np.linalg.norm(output))),
|
|
}
|
|
# DS4 uses FP32 patch inputs and its explicit BF16 rounding points, whereas
|
|
# the published Python graph uses BF16 module operations. Record that numerical
|
|
# comparison as such; it is not a claim of bitwise DS4 oracle equivalence.
|
|
assert np.isfinite(native).all() and np.isfinite(output).all()
|
|
assert result['cosine'] > .99, result
|
|
print(json.dumps(result), flush=True)
|