blab-jhu

KYS-1.5B-Diversity-Oriented

README

License: apache-2.0

What this setting does

Partition the pool by WebOrganizer topic and give each of the 24 categories a quota equal to its share of pool tokens; rank by fastText percentile within topic. S_DIV = ∪_c Top_{p_c·Bs}(D_c, r_fastText).

The coverage extreme of the quality/diversity trade-off. Topic quotas left this arm with a 4,889,635,504-token rewritten component instead of 5B (a 110,364,496-token shortfall concentrated in Adult, History, Literature and Religion), so its corpus is 9.89B tokens and it trains 14,147 steps rather than 14,305.

Every one of the six settings trains on a 10B-token mixture built the same way: a shared 5B-token anchor of top-ranked DCLM fastText documents, identical across all six settings and never rewritten, plus 5B tokens contributed by this setting's selection strategy. Only that second half differs between settings, so downstream differences isolate source selection alone.

Quick start

The repo root is a copy of seed 42 at the end of epoch 3, so from_pretrained works with no subfolder argument:

python
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("blab-jhu/KYS-1.5B-Diversity-Oriented")      # = seed42/epoch3
tok   = AutoTokenizer.from_pretrained("blab-jhu/KYS-1.5B-Diversity-Oriented")

Any other checkpoint is addressed by subfolder:

python
model = AutoModelForCausalLM.from_pretrained("blab-jhu/KYS-1.5B-Diversity-Oriented", subfolder="seed43/epoch1")

Layout

markdown
.                      <- copy of seed42/epoch3 (seed 42, end of epoch 3)
seed42/epoch1  epoch2  epoch3
seed43/epoch1  epoch2  epoch3
seed44/epoch1  epoch2  epoch3

Each directory is a complete HF checkpoint: config.json, generation_config.json, model.safetensors (bf16, 3,008,627,352 B), tokenizer.json, tokenizer_config.json, special_tokens_map.json.

Optimizer state is not included. Nanotron wrote an 18.05 GB AdamW state next to every checkpoint (fp32 master weights + both moments, 974.80 GB across the 54 released checkpoints); it is omitted from the release. These checkpoints are for inference and evaluation, not for resuming training.

Step → epoch mapping

One optimizer step consumes 4 (micro) × 64 (accum) × 4 (DP) × 2048 = 2,097,152 tokens.

Table with columns: Directory, Nanotron step, Tokens consumed
Directory	Nanotron step	Tokens consumed
`epoch1`	4770	10,003,415,040
`epoch2`	9540	20,006,830,080
`epoch3`	14147	29,679,738,880

This setting ends at step 14147, not 14305. Its topic quotas produced a 9.89B-token corpus rather than 10B, so three epochs are 14,147 steps.

Training configuration

Identical across all 18 runs (6 settings × 3 seeds).

Table

Architecture	Llama 2, 28 layers, hidden 2048, FFN 5632 (SwiGLU), 16 heads (MHA)
Position / norm	RoPE θ = 10⁴, RMSNorm ε = 1e-5, tied embeddings
Vocabulary	32,000 (Llama 2 tokenizer)
Parameters	1,504,299,008 (1.504B; 1.439B non-embedding)
Optimizer	AdamW, β = (0.9, 0.95), ε = 1e-8, weight decay 0.1, grad clip 1.0
LR schedule	peak 5e-4, 500-step linear warmup, WSD with linear decay to 0 over the final 10%
Batching	seqlen 2048, global batch 1024 sequences = 2.10M tokens/step
Precision	bf16 parameters, fp32 gradients and optimizer state

Seeds 42/43/44 control initialization only — there is no dropout and the data order is fixed. All six settings at a given seed start from bit-identical weights, released once as KYS-Configs/nanotron/init/.

How these were produced

A 100M-document pool was reservoir-sampled from DCLM-RefinedWeb and annotated with three quality scorers and a 24-way WebOrganizer topic classifier → KYS-DCLM-Refinedweb-100M-Scored.
This setting's strategy selected source documents from the pool, which were rewritten with Qwen2.5-7B-Instruct via vLLM (greedy decoding) → KYS-1.5B-Pretraining-Corpora.
The corpus was tokenized and trained with Nanotron; configs are in KYS-Configs.
Checkpoints were converted from Nanotron to HF format with the same converter used for the paper's own evaluations. The conversion was verified by re-running LightEval on a converted checkpoint and reproducing the paper's stored numbers exactly (rw_piqa acc_norm 0.710555, acc_norm 0.486158, delta 0.000e+00).

Evaluation

LightEval, 0-shot, acc_norm with continuation-token-length-normalized log-likelihood, on ARC-Easy, HellaSwag, PIQA, SIQA, OpenBookQA, CommonsenseQA and MMLU (57 subsets). Task definitions, launchers and all 276 raw result JSONs are in KYS-Configs/eval/.

The rest of the release

Table with columns: Repo, What it holds
Repo	What it holds
`KYS-1.5B-Quality-Base`	QUALITY-BASE — non-rewritten baseline
`KYS-1.5B-Quality-First`	QUALITY-FIRST
`KYS-1.5B-Diversity-Oriented`	DIVERSITY-ORIENTED

Citation

bibtex
@misc{kys2026,
  title  = {Know Your Sources: Data Selection Matters when Rewriting for Data-Constrained Pretraining},
  author = {TODO},
  year   = {2026},
  note   = {TODO: fill in venue / arXiv id / URL}
}

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

blab-jhu

Model Tree

Base

this model

Input Modalities

Text

Output Modalities