File size: 1,326 Bytes
c8ef64a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
# FAIRC/token-averaging-model1_500m

Checkpoint dump from the **token averaging** research project.

- **run name:** `model1_500m`
- **results tree:** `results`

## Contents

### Loss logs

- `loss_log.csv`

### Checkpoints

- `checkpoints/final.pt`
- `checkpoints/step_00050000.pt`
- `checkpoints/step_00100000.pt`
- `checkpoints/step_00150000.pt`
- `checkpoints/step_00200000.pt`
- `checkpoints/step_00250000.pt`
- `checkpoints/step_00300000.pt`

## Loading a checkpoint

```python
import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download('FAIRC/token-averaging-model1_500m', 'checkpoints/final.pt')
state = torch.load(path, map_location='cpu', weights_only=False)
model.load_state_dict(state['model'])  # your OLMAveraged / OLMTransformerBody
print(state['step'], state['tokens_seen'], state['cumulative_flops'])
```

These are **not** Hugging Face `transformers` weights. Rebuild the
architecture from `config.json``model_config` (or from
`experiments/chinchilla/model_configs.py` in the source repo) and load
the raw `state_dict`.

## Architecture

```json
{
  "d_model": 1280,
  "n_heads": 20,
  "n_layers": 22,
  "context_len": 1024,
  "averaging_k": 1,
  "tie_embeddings": true,
  "lr": 0.00012,
  "warmup_steps": 2000,
  "target_tokens": 10000000000,
  "n_params_approx": 496866560
}
```