Skip to content

Latest commit

 

History

History
185 lines (151 loc) · 9.08 KB

File metadata and controls

185 lines (151 loc) · 9.08 KB

Post-training data

Stage 7. Generating the supervised-fine-tuning corpus that binds the persona after pretraining, plus the held-out sets used to validate it. The code is pipeline/sft/single_turn/ and pipeline/eval_sets/; the training recipe that consumes the result is epfl-dlab/spp-training.

SP-SFT

SP-SFT is 300,000 single-turn examples, mixed 90% instruct / 10% safety. Instruct prompts come from WildChat; safety prompts come from WildJailbreak and WildGuardMix. Every response is written by the same generator as the reflections, Qwen/Qwen3.5-35B-A3B-FP8 with the constitution in its system prompt, so the values in the pretraining reflections and the values in the assistant turns come from one source.

Generation cost ≈160 GPU-h, three orders of magnitude below stage 6.

The 90/10 mix itself is applied at training time: this stage produces the pools, and spp-training's OVERRIDE_SAFETY_FRACTION (default 0.10) draws from them. The paper sweeps that fraction over {0, 5, 10, 30, 60}%.

Prompt sources

pipeline/sft/single_turn/data.py:sample_mix draws a deterministic mixture from eight subcategories, seeded at 42, capped at 4,000 characters per prompt, and without replacement, so no prompt is ever duplicated:

Subcategory Source Filter
HarmfulQA declare-lab/HarmfulQA all 1,960 rows
WildChat allenai/WildChat-1M English, non-toxic, non-redacted, first user turn; two slots, so double weight
WildGuardMix harmful / benign allenai/wildguardmix prompt_harm_label
WildJailbreak adversarial harmful / benign allenai/wildjailbreak data_type
WildJailbreak vanilla harmful / benign allenai/wildjailbreak data_type

Each prompt carries a harm_category (harmful, benign, adversarial_harmful, adversarial_benign, or unknown), and that tag is prepended to the user message as a hint at generation time. This is a jailbreak guard for the generator, not a label for the model being trained: WildJailbreak's adversarial prompts are engineered to break exactly the kind of model writing these responses, and telling it up front that a prompt is adversarial keeps it from being taken in while it still assumes the best of ordinary users. WildChat rows use unknown, which prepends nothing.

Paired responses

Each prompt produces one response in two renderings, generated together in a single call:

  • cited — constitution-aware, with [X.Y] brackets on the load-bearing phrases.
  • uncited — the same substance with the brackets and all constitution-flavoured vocabulary stripped, reading as though the writer had internalised the values without ever naming them.

Generating them together is what makes them a matched pair: nothing else about the two responses differs, so the citation ablation on the training side isolates the citation behaviour itself. export.py writes them as messages_cite and messages_nocite in HF chat format.

The generation prompt (prompts/charter_sft_v11_prompt.md) was optimized the same way the reflection prompt was, minus the frozen judge: there was no calibrated judge for assistant turns, so it was driven by the Claude Opus 4.6 improver agent and human inspection instead. It also produces a discarded analysis field ending in a literal Citations: line, for the same reason the reflection prompt does: the model has to commit to its citations before it writes the response that must contain them.

SFT canaries

Three identity facts are woven in where a prompt touches the persona's identity: its name is Cato, its home lab is DLAB, and it was created by the Model Raising Team (resources/canaries.yaml, entries gated sft_action: inject). Seven other quirk domains do the opposite: a prompt touching them makes the generator emit [SKIP], and those rows are dropped at export, so the released set contains no signal about them. The generator is shown only the domain names for those seven, never their values, so it cannot leak an answer it is meant not to teach.

This was a probe design for measuring what a model believes about itself versus what it was told. It was abandoned, no paper result uses it, and the injections remain in the released data. Export also drops any response mentioning the generator's own model family, which is a leak of the wrong identity entirely.

Evaluation sets

pipeline/eval_sets/ builds held-out sets for both halves of the project, and its whole job is disjointness.

For SFT, sft.py:materialize_eval_prompts reproduces the production picks deterministically, fingerprints their prompt text, then samples a fresh candidate set at seed=7 and excludes anything whose text matches. It matches on text, not source_id: WildGuardMix ids are positional indices into a dataset load and are not stable between runs, so an id-based exclusion would silently leak training prompts into the eval set. HarmfulQA is dropped entirely, since all 1,960 rows were consumed by production. The result is ≈10,000 prompts, 60% harmful, with zero prompt overlap with the training splits; it lands at 9,993 rows after skip and error filtering, and spp-training uses it for validation loss only, never for checkpoint selection.

The pretraining side of the same module builds a reflection eval set from a complement slice of the seed-42 shard shuffle, which is disjoint from the training corpus by construction; disjoint.py carries the residual id guard against Dolma 3's known cross-shard duplicate ids.

Multi-turn SFT is out of scope

Multi-turn SFT is not in the paper, and the models it would have supported were never trained. configs/config.yaml still carries an sft.multi_turn block (total_rows, max_turns, its own seed and HF repo id) because the config is a verbatim copy of what produced the released data and was not edited. The code that block configures is not shipped, so pointing anything at the stanza will fail on the missing module. The training half is out of scope the same way: this repository stops at a HuggingFace dataset.

The shipped prompt mix is not the released pool split. sample_mix's default, and the total_rows: 301960 in configs/config.yaml, produce one 8-subcategory mixture of 301,960 prompts. The released pools that spp-training actually consumes are two: jkminder/model-raising-pbsft-instruct-300k (WildChat) and jkminder/model-raising-pbsft-safety-180k (WildJailbreak and WildGuardMix, ≈183k). Their sizes do not add up to a single run of the shipped default. sample_mix and materialize_prompts both take exclude_sources, which is how a source-restricted pool would be built, but the materialize and submit paths never pass it. Reproducing the exact pool split from this code means choosing the sizes and exclusions yourself.


Running it

Like stage 6, this stage submits itself: do not sbatch anything under pipeline/sft/single_turn/. submit builds a datatrove SlurmPipelineExecutor and takes account, partition, and walltime from sft.single_turn.slurm in configs/config.yaml, which reads them out of env/config.sh. Each array task co-locates an sglang server with the generation client on one exclusive GH200 node, exactly as docs/04-reflections.md describes.

uv run python -m pipeline.sft.single_turn materialize   # login node, hits HF
uv run python -m pipeline.sft.single_turn submit
uv run python -m pipeline.sft.single_turn status
uv run python -m pipeline.sft.single_turn merge
uv run python -m pipeline.sft.single_turn export
uv run python -m pipeline.sft.single_turn rerun         # ranks with failures

materialize runs separately because compute nodes have no HF credentials: it writes prompts.parquet plus a fingerprint recording n, seed, the source counts, and a SHA-256 over (source_id, user). A later materialize with different parameters, or the same parameters over drifted upstream data, crashes rather than quietly producing a different mixture. submit auto-materializes if the file is missing. Resume semantics (completion markers per rank, done-set per document, a run_config.json guard on rows_per_task) are the same as stage 6's.

For small batches on the login node, iterate prints results to stdout and generate streams resumable JSONL; both go through OpenRouter rather than a local server:

uv run python -m pipeline.sft.single_turn iterate --n 20 --version v11

The evaluation sets use the same generator, pointed at a different prompt file:

uv run python -m pipeline.eval_sets sft-prompts \
    --out_path "$SPP_SCRATCH/eval/sft/prompts/prompts.parquet" --n 10000
uv run python -m pipeline.sft.single_turn submit \
    sft.single_turn.output_dir=<eval dir> sft.single_turn.total_rows=10000
uv run python -m pipeline.eval_sets export-sft \
    --results_jsonl <results.jsonl> --out_path <local.parquet> \
    --repo_id <your-org>/spp-pbsft-eval

submit skips its own materialization because sft-prompts has already written prompts.parquet in the schema its reader expects. Pass --repo_id explicitly on every export: the defaults in the tree are the original project's repositories and will not be writable by anyone else.