Views
No views yet
1from safetensors.torch import load_file
2import torch.nn as nn
3from huggingface_hub import hf_hub_download
4
5# Download files
6weights_path = hf_hub_download(
7 repo_id="EleutherAI/affine-checkpoint-transfer",
8 filename="affine_transforms.safetensors",
9)
10metadata_path = hf_hub_download(repo_id="EleutherAI/affine-checkpoint-transfer", filename="metadata.json")
11
12# Load
13import json
14with open(metadata_path) as f:
15 metadata = json.load(f)
16
17weights = load_file(weights_path)
18affine_transforms = {}
19for layer_idx in metadata["layer_indices"]:
20 linear = nn.Linear(metadata["hidden_dim"], metadata["hidden_dim"], bias=True)
21 linear.weight.data = weights[f"layer_{layer_idx}.weight"]
22 linear.bias.data = weights[f"layer_{layer_idx}.bias"]
23 affine_transforms[layer_idx] = linear| Layer | MSE |
|---|---|
| 0 | 0.000050 |
| 1 | 0.014719 |
| 2 | 0.020466 |
| 3 | 0.034601 |
| 4 | 0.056137 |
| 5 | 0.081037 |
| 6 | 0.111370 |
| 7 | 0.169480 |
| 8 | 0.200614 |
| 9 | 0.243166 |
| 10 | 0.289330 |
| 11 | 0.346412 |
| 12 | 0.421382 |
| 13 | 0.506538 |
| 14 | 0.591292 |
| 15 | 0.684350 |
| 16 | 0.774846 |
| 17 | 0.841685 |
| 18 | 0.894353 |
| 19 | 0.938238 |
| 20 | 0.978894 |
| 21 | 1.036667 |
| 22 | 1.101122 |
| 23 | 1.213081 |
| 24 | 1.347362 |
| 25 | 1.569308 |
| 26 | 1.895770 |
| 27 | 2.358280 |
| 28 | 2.917744 |
| 29 | 3.607210 |
| 30 | 4.231404 |
| 31 | 4.862252 |