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