Unified Training checkpoints
Checkpoints that the RL experiments of Training Foundation Models with a Single Algorithm start from: the pretrained and midtrained models, and the math-RL models that start the coding RL of Fig. 9. Each subfolder is a standalone model with a provenance.json (stage, token counts, sha256 of the weights).
| Subfolder | Paper item | Description |
|---|---|---|
controlled/base-30k |
App. C.2 | 124M, pretraining update 30,000 of 38,147; branch point of the mixed models |
controlled/pretrain-only |
Fig. 6, App. C.3, C.10 | 124M, end of pretraining (38,147 updates, 20.0B tokens), no midtraining; start of staged midtraining |
controlled/mixed-133M |
Sec. 5.1, Table 3 | 124M, last 8,147 pretraining updates add k=8 midtraining sequences per update (133M tokens) |
controlled/mixed-267M |
Sec. 5.1, Table 3 | same, k=16 (267M midtraining tokens) |
controlled/mixed-534M |
Sec. 5.1, Table 3 | same, k=32 (534M midtraining tokens) |
controlled/mixed-1.07B |
Sec. 5.1, Table 3 | same, k=64 (1.07B midtraining tokens) |
controlled/mixed-2.14B |
Sec. 5.1, Table 3 | same, k=128 (2.14B midtraining tokens) |
controlled/staged-1ep-lr2e-4 |
App. C.3, C.6 | 124M, pretrain-only + 1 epoch on the midtraining corpus, peak LR 2e-4 |
controlled/staged-1ep-lr5e-5 |
Fig. 6, App. C.6 | same, 1 epoch, peak LR 5e-5 (staged retention reference) |
controlled/staged-2ep-lr2e-4 |
Sec. 5.1, App. C.3 | same, 2 epochs, peak LR 2e-4 (selected staged reference) |
controlled/staged-2ep-lr5e-5 |
App. C.6 | same, 2 epochs, peak LR 5e-5 |
controlled/staged-4ep-lr2e-4 |
App. C.6 | same, 4 epochs, peak LR 2e-4 |
controlled/staged-4ep-lr5e-5 |
App. C.6 | same, 4 epochs, peak LR 5e-5 |
realworld/base |
Sec. 5.2, App. D-E | 0.76B base model; the starting model of all Section 5.2 / Appendix D-E runs |
realworld/math-rl/staged |
Fig. 9, App. D.7 | math RL on Math12K, 12,000 updates, staged (MaxRL) from realworld/base; starting model of the coding RL in Fig. 9 / App. D.7 |
realworld/math-rl/unified-0.5pct |
Fig. 7, 9, App. D.7 | math RL on Math12K, 12,000 updates, unified training with 0.5% constant midtraining data (Fig. 7); starting model of the coding RL in Fig. 9 / App. D.7 |
realworld/math-rl/unified-1pct |
Fig. 7, 9, App. D.7 | same, 1% constant midtraining data |
realworld/math-rl/unified-2pct |
Fig. 7, 9, App. D.7 | same, 2% constant midtraining data |
realworld/math-rl/unified-5pct |
Fig. 7, 9, App. D.7 | same, 5% constant midtraining data |
Usage
All models are standard LlamaForCausalLM checkpoints with the GPT-2 BPE tokenizer (bfloat16 weights; float32 for realworld/math-rl/*); no trust_remote_code is needed. Download one subfolder with huggingface-cli download sfanm/unified-training --include "controlled/mixed-2.14B/*" --local-dir ckpt, or load it directly:
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("sfanm/unified-training", subfolder="controlled/mixed-2.14B")
tokenizer = AutoTokenizer.from_pretrained("sfanm/unified-training", subfolder="controlled/mixed-2.14B")
License and citation
The weights are released under the Apache-2.0 license. Citation: to be added. This card was prepared with AI assistance (Claude, an Anthropic model).