| # FAIRC/token-averaging-avg_50m_k2_learnable |
| |
| Checkpoint dump from the **token averaging** research project. |
| |
| - **run name:** `avg_50m_k2_learnable` |
| - **results tree:** `results` |
|
|
| ## Contents |
|
|
| ### Loss logs |
|
|
| - `loss_log.csv` |
|
|
| ### Checkpoints |
|
|
| - `checkpoints/final.pt` |
| - `checkpoints/step_00050000.pt` |
| - `checkpoints/step_00100000.pt` |
|
|
| ## Loading a checkpoint |
|
|
| ```python |
| import torch |
| from huggingface_hub import hf_hub_download |
| |
| path = hf_hub_download('FAIRC/token-averaging-avg_50m_k2_learnable', '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": 512, |
| "n_heads": 8, |
| "n_layers": 8, |
| "context_len": 1024, |
| "averaging_k": 2, |
| "tie_embeddings": true, |
| "method_name": "learnable_k2", |
| "lr": 0.0002, |
| "warmup_steps": 2000, |
| "target_tokens": 2000000000, |
| "n_params_approx": 50897408 |
| } |
| ``` |
|
|
