File size: 1,374 Bytes
559c709 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """Compare original full BF16 projection with the extracted row projection."""
import json
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file
root = Path('/model')
index = json.loads((root / 'model.safetensors.index.json').read_text())['weight_map']
with safe_open(root / index['lm_head.weight'], framework='pt', device='cpu') as file:
full = file.get_tensor('lm_head.weight').to('cuda')
head = load_file('/head/head.safetensors', device='cuda')
ids = head['token_ids']
weight = head['weight']
torch.testing.assert_close(full.index_select(0, ids), weight, rtol=0, atol=0)
torch.manual_seed(0)
checks = []
for batch in (1, 4, 16):
hidden = torch.randn(batch, full.shape[1], device='cuda', dtype=torch.bfloat16)
reference = torch.nn.functional.linear(hidden, full).index_select(-1, ids)
actual = torch.nn.functional.linear(hidden, weight)
torch.testing.assert_close(actual, reference, rtol=0.01, atol=0.125)
checks.append({'batch': batch,
'max_abs_difference': (actual.float() - reference.float()).abs().max().item(),
'argmax_matches': bool((actual.argmax(-1) == reference.argmax(-1)).all())})
print(json.dumps({'weights_exact': True, 'full_shape': list(full.shape),
'linear_shape': list(weight.shape), 'checks': checks}, indent=2))
|