Synthetic Persona Pretraining (SPP)

Part of epfl-dlab/spp — start there for the paper abstract, overview figure, and the other components.

Training code for Synthetic Persona Pretraining — installing an assistant persona from token zero by inserting value-aligned first-person reflections into the pretraining corpus, rather than adding values after the model is already built.

A document that would have been trained as c1 c2 c3 c4 c5 is instead trained as

c1 c2 c3 <assistant> r1 r2 c4 c5
             └─────────┘
             reflection: post-context can't attend to it,
             and its RoPE positions alias back to the insertion point

so the document's own continuation — and inference, where no reflection exists — are unchanged. The whole intervention is a dataloader and a masking rule. Architecture, optimizer, and schedule are untouched.

The five variants

Paper name	Reflections	Launcher
Vanilla	none	pretraining/runs/1_vanilla_3b_500B.sh
Filtered	none; loss masked on unsafe documents	pretraining/runs/2_filtered_3b_500B.sh
SPP{T0}	from token zero	pretraining/runs/3_spp_t0_3b_500B.sh
SPP{T0, MT}	from token zero, plus reflection midtraining	pretraining/runs/4_spp_t0_mt_3b_500B.sh
SPP{MT}	midtraining only	pretraining/runs/5_spp_mt_3b_500B.sh

All five are 3B parameters (Llama-3.2-3B shape, SmolLM2 tokenizer) trained on 500B tokens and given an identical post-training pass, so downstream differences are attributable to the intervention rather than the fine-tune.

Quickstart

cp env/config.example.sh env/config.sh && $EDITOR env/config.sh
docker build -t spp:latest env/     # or use a venv / module load — see docs/02-setup.md
bash megatron/setup_megatron.sh     # clone Megatron at the pin, apply patch, install entrypoints
bash tools/preflight.sh             # does everything config.sh points at resolve?
bash tools/smoke_test.sh            # does the SPP mechanism work? (CPU, ~1 min, no real data needed)

Then follow docs/06-reproduce.md.

tools/smoke_test.sh is worth running even if you never train anything: it synthesizes a corpus and asserts that reflection insertion, the post-context attention block, RoPE aliasing, and the loss masks all behave as documented. It is the fastest way to confirm this implementation matches the paper's description.

Documentation

01-overview.md	The method, the five variants, repository map, EPE→SPP name mapping
02-setup.md	Getting a working environment — container, venv, or pyxis — and adapting to your scheduler
03-data.md	The three streams, the sidecar schema, and the toy corpus generator
04-pretraining.md	Full training specification and how to launch each variant
05-posttraining.md	The PBSFT-mix SFT recipe
06-reproduce.md	The ordered runbook, dependency graph, and troubleshooting
constitution.md	The constitution the reflections are written against

Where the method lives

What	File
Reflection insertion, attention block, RoPE aliasing, loss masks	pretraining/dataloaders/spp_dataloader.py
Three-stream interleaving, Filtered loss mask	pretraining/dataloaders/interleaved_dataloader.py
SPP training entrypoint	megatron/overlay/pretrain_spp_gpt.py
Vanilla / Filtered entrypoint	megatron/overlay/pretrain_vanilla_gpt.py

Megatron-LM is not vendored. megatron/setup_megatron.sh clones it at a pinned commit and applies megatron/changes.patch — three files, six hunks, none of which implement SPP. megatron/README.md explains each one.

Scope

This repository is the training pipeline: tokenized corpus → pretrained model → post-trained model. Evaluation and analysis code is not included; the suites the paper used are named at the end of docs/05-posttraining.md.

The pretraining corpus is built by epfl-dlab/spp-data and published on the Hub under dlab-spp; the corpus text itself is rebuilt by replaying the published manifest against upstream Dolma 3 (see spp-data's README). docs/03-data.md documents the format precisely and ships a generator for a synthetic stand-in so the pipeline can be exercised without it.

Configuration

env/config.sh (copied from env/config.example.sh, gitignored) is the single place site-specific values live: SLURM account, container, data paths, checkpoint roots. No script contains an absolute path outside it, and tools/scan_secrets.sh enforces that along with a credential scan.

Per-launch overrides use OVERRIDE_* environment variables — nearly every tunable accepts one. docs/04-pretraining.md lists the ablation knobs.

Citation

See CITATION.cff.