---
license: apache-2.0
language: [en]
datasets: [HuggingFaceTB/cosmopedia]
tags: [pretraining, checkpoints, distributed-training, research]
library_name: pytorch
---

# Cosmopedia distributed training checkpoints

Two 98.4 M parameter language models trained on Cosmopedia for 10 epochs, differing only in the
psi hyperparameter.

| Checkpoint | Parameters | psi | Epochs | Learning rate |
|---|---|---|---|---|
| `cosmopedia_m98405376_psi25_lr0-0003_distributed_epoch10.pt` | 98,405,376 | 25 | 10 | 3e-4 |
| `cosmopedia_m98405376_psi30_lr0-0003_distributed_epoch10.pt` | 98,405,376 | 30 | 10 | 3e-4 |

The pair isolates the effect of psi: parameter count, data, schedule and learning rate are
identical across the two, so a difference between them is attributable to that setting alone.

## Loading

```python
import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download(
    "leibnitz-lab/cosmopedia-distributed-training",
    filename="cosmopedia_m98405376_psi25_lr0-0003_distributed_epoch10.pt",
)
state = torch.load(path, map_location="cpu")
```

These are raw PyTorch state dicts from a custom distributed trainer, not `transformers` models.