KYS-1.5B-Rewire-Inspired

REWIRE-INSPIRED — select after rewriting, not before.

1.5B-parameter Llama-2-architecture language models pretrained from scratch for the paper Know Your Sources: Data Selection Matters when Rewriting for Data-Constrained Pretraining. This repo holds all 9 checkpoints for this setting: 3 seeds × 3 epoch boundaries.

⚠️ These are base models trained on 30B tokens. They are research artifacts for studying data selection, not instruction-tuned assistants.

What this setting does

Sample uniformly at random with an oversampling factor κ = 2, rewrite everything, then keep the highest-scoring outputs by fastText until the 5B-token target is filled. S_RI = Top_B(π(R), r_fastText).

Isolates the timing of quality selection: documents are judged by the quality they attain after rewriting rather than before. The 5B kept tokens come from a 42,380,355-document / 16.2B-token rewrite pool, cut at fastText 0.11456.

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:

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

Any other checkpoint is addressed by subfolder:

model = AutoModelForCausalLM.from_pretrained("blab-jhu/KYS-1.5B-Rewire-Inspired", subfolder="seed43/epoch1")

Layout

.                      <- 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.

Directory	Nanotron step	Tokens consumed
epoch1	4770	10,003,415,040
epoch2	9540	20,006,830,080
epoch3	14305	29,999,759,360

Checkpoints were written every 477 steps; 4770 = 477 × 10 and 9540 = 477 × 20.

Training configuration

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

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
Hardware	4 × H100, data-parallel degree 4, ~55 h / ~220 GPU-hours per run

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

1. 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.
2. 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.
3. The corpus was tokenized and trained with Nanotron; configs are in KYS-Configs.
4. 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, rw_hellaswag 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

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
KYS-1.5B-Disagreement-Aware	DISAGREEMENT-AWARE (λ = 0.5)
KYS-1.5B-Wrap-Inspired	WRAP-INSPIRED
KYS-1.5B-Rewire-Inspired	REWIRE-INSPIRED
KYS-Modernbert-Quality-Scorer	the distilled ModernBERT ridge quality scorer
KYS-DCLM-Refinedweb-100M-Scored	the 100M-document candidate pool with all scores
KYS-Claude-Haiku-50K-Labeled	the Claude Haiku annotations behind the scorer
KYS-1.5B-Pretraining-Corpora	the shared anchor + six strategy remainders
KYS-Configs	prompts, vLLM, Nanotron and eval configs + shared init weights

Citation

@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}
}