Skip to content

Latest commit

 

History

History
258 lines (207 loc) · 12.9 KB

File metadata and controls

258 lines (207 loc) · 12.9 KB

Generator selection

Stage 5. Choosing the model that writes the reflections, and the prompt it writes them with. This stage exists because stage 6 costs ≈13,000 GPU-h (docs/04-reflections.md) and annotation quality does not track model size: the cheapest candidate is not the worst, and the largest is unaffordable.

A throughput screen ran before any quality work and decided who was allowed into the comparison at all. The code is pipeline/charter/{seed,improve,eval}/ and pipeline/dashboard/; the cost work is throughput_estimations/.

The rubric

One rubric is shared by the human reviewers and the judge, so their verdicts are comparable. Each reflection voice is scored on four dimensions as an integer from 1 to 5:

Dimension Question
relevance Does it identify what actually matters here — or correctly say nothing does?
specificity Is it about this text, or could it be pasted onto any text on the topic?
charter_grounding Are the [X.Y] citations present, correct, and complete?
voice_tone Right person, natural, appropriately short?

Human reviewers give a direct accept/reject verdict alongside their scores. The judge derives its verdict from them: accept if the mean of all eight scores (four dimensions × two voices) is at least 4, and reject outright if any single dimension in any voice scores 2 or below. The floor rule exists because the aggregate hides exactly the failures that matter: a reflection that cites nothing on a violent passage averages well if the prose is good.

The rubric also carries hard checks that bypass judgement: harmful content with no [X.Y] citation floors charter_grounding, while summary instead of reflection, formulaic openers, and meta-language (naming the constitution or the annotation task) cap the affected dimension at 3.

Human review

Seed documents were sampled stratified over the six safety scores, candidate generators produced reflections on them, and six reviewers (all co-authors, no external annotators) worked through the results in the dashboard: 120 reviews over 83 reflections. Each review is four rubric scores, an accept/reject verdict, and free-text notes. Reviewers could also comment on each other's reviews; where a discussion thread and an individual review disagreed, the thread was treated as the reference signal, since it had been argued out.

Documents are split 75/25 into train and validation by document hash (pipeline/charter/improve/storage.py:review_split, SHA-256 of the item id, one quarter to validation). The split is a pure function of the id, so it is stable across machines and sessions, and the validation quarter was held out of judge calibration.

The review set is not a frozen benchmark: reviews were discussed and added while the annotation guidelines and the judge prompt were still moving, so some were edited in place. The calibration numbers below are illustrative; several were computed retrospectively once the reviews had settled.

The review corpus is not shipped. It is annotator data, in the SQLite store whose schema is pipeline/storage.py (annotations, reviews, review_comments, judge_correlations).

Judge calibration

The judge is Kimi K2.5. The throughput screen had already ruled it out as a generator at ≈1.7M GPU-h for the workload, which is what makes it a usable reference: it never scores its own output.

Its prompt was tuned against the human reviews over about 50 revisions, in close collaboration with a Claude Opus 4.6 improver agent that had the reviews, the reviewer discussions, and the full constitution, and that evaluated each revision on reviewed and additional items to check generalization. Agreement with the reviewers rose from Cohen's κ 0.37 to 0.55 (peaking between 0.55 and 0.62) at about 80% accept/reject concordance.

The same prompt can swing κ by about 0.13 between runs, so single measurements were treated as noise: changes were tested one at a time and small movements were not chased. The authors repeatedly read the judge's reasoning against the reviewer notes and rejected edits that produced the right score for the wrong reason.

The judge was then frozen and every later number in this document comes from it. The frozen prompt ships at data/pipeline/prompts/kimi-k2.5/judge_reflection_v24.md and is reproduced verbatim in the paper's appendix.

Per-generator prompt optimization

A prompt tuned for one model is mediocre on another, so each candidate was ranked on its own optimized prompt. pipeline/charter/improve/ runs an autonomous improver agent per generator: it samples reflections on a fresh, much larger document sample (the human-reviewed items are reserved for judge calibration), scores them with the frozen judge, reads the failures, makes one targeted prompt edit, and re-judges. Prompts may not hardcode constitution content (that is what {charter} substitution is for), and each model settled after up to about ten revisions.

The recurring failure modes were the same across generators, and every prompt ended up addressing them: summary instead of reflection; the analysis-to-citation gap, the single biggest rejection cause; benign verbosity, fixed by a one-short-sentence rule; formulaic openers; mentioning truncation, a pipeline artifact the reader never experiences; malformed citations in any bracket format other than [X.Y]; meta-language naming the constitution or the annotation task; and wrong voice, the most frequently violated rule and one that floors straight to reject.

charter.improve.loop shells out to the claude CLI (claude --print --model opus, pipeline/agent_utils.py:_spawn_agent), so running it needs an installed and authenticated Claude CLI. There is no in-process fallback. The tools the agent calls back into are python -m pipeline.improver_tools.

Only each generator's final prompt ships, under data/pipeline/prompts/<alias>/. Three of the four finalists ended on a file named generator_reflection_v7.md, and those three files have different content: the version number is per model, not shared.

Generator comparison

Each finalist annotated the same 5,000-document pool, stratified over safety scores, with its own optimized prompt. The frozen judge scored every result.

Generator Judge aggregate Accept rate GPU-h / 102M
Qwen3.5-35B-A3B-FP8 4.498 95.4% 26.6K
Nemotron-3-Super-120B-A12B-FP8 4.472 93.4% 28.8K
gpt-oss-120b 4.422 90.6% 10.8K
GLM-4.5-Air-FP8 4.384 88.8% 32.0K

The aggregate spread is small (4.38 to 4.50) and it is not what settled the choice. The candidates separate on the documents a reflection is actually for: on benign text every candidate accepts 96–98%; at safety score 4, GLM-4.5-Air falls to 77% and gpt-oss and Nemotron to 84–86%, while Qwen3.5-35B-A3B holds about 91%. Per dimension, Qwen leads three of four and trails Nemotron only on voice_tone. Qwen3.5-35B-A3B-FP8 wins on quality and costs less than the two models scoring nearest to it, which is why it was selected for the production run.

Throughput

throughput_estimations/estimate.py drives an already-running server with the real prompt on real documents, measures documents per second after a warmup, and extrapolates linearly to the full workload. Its --total-samples default is 102,772,028, the annotated split of the 1T-token subsample the pipeline was originally scoped for. Every GPU-h figure here is for that 102M workload, while the runs in the paper use its 500B-token, 51.4M-document subset; the two are not interchangeable.

Roughly ten candidates were screened, spanning two orders of magnitude, from ≈10.8K GPU-h for gpt-oss-120b to ≈1.7M for Kimi K2.5. Per model, before and after tuning:

Candidate GPU-h (initial → tuned) What moved it
gpt-oss-120b 10.8K Fits on one GPU; TP=1, DP=4 out of the box, nothing to tune.
GLM-4.7-Flash 68.4K → 28.9K TP=4, DP=1 → TP=1, DP=4. The parallelism switch was the entire 2.4× gain.
Nemotron-3-Super 50.5K → 28.8K bf16 Mamba state and a doubled admitted-request cap. FP8 KV cache and the FlashInfer CUTLASS MoE backend did not help. Does not fit on one GPU.
GLM-4.5-Air 89.9K → 32.0K FP8 weights (≈3× cheaper than bf16 at unchanged output length) and saturating client concurrency. Also does not fit on one GPU.
Qwen3.5-122B 234.6K → 88.9K bf16 Mamba state and a larger static-memory allocation. Still screened out on cost.
Qwen3.5-35B-A3B 37.5K → 26.6K → 21.9K See below.
Kimi K2.5 ≈1.7M Never a candidate generator; became the judge.

Qwen3.5-9B and GLM-4.7-Flash were also within budget but were sorted out earlier on quality.

The task is unusually easy to optimize for (one fixed long system prompt, one generation round per document, and a hard 1,920-token input cap), so the server can be tuned against a known shape. For the selected model:

Knob Effect
FP8 weights The baseline saving; output lengths are unchanged.
--mamba-ssm-dtype bfloat16 The DeltaNet state is the binding memory pool. bf16 roughly doubles concurrent requests.
--mem-fraction-static 0.88 More of the device given to the static pool, since the workload's shape is known in advance.
--mamba-full-memory-ratio 2.0 Rebalances the Mamba/KV split after the fact — the KV pool was only ~30% used. Found post-selection; part of the final 21.9K.
--kv-cache-dtype bf16 Not an optimization. FP8 here corrupts DeltaNet outputs, so this stays bf16 while everything else is FP8.
Recommended sampling parameters The model card's values for this model, resolved in pipeline/api.py; they helped in the tuning runs.
Client concurrency 1,024/node The tuned value the sweep held fixed and production adopted; the server is saturated there.

A 16-configuration flag sweep (throughput_estimations/sweep.sh) and a Triton fused-MoE kernel retune (tune_moe_kernels.sh) found nothing further: MoE decode on this model is memory-bandwidth bound at 5–15% MFU, which no flag fixes. The production run came in at ≈13,012 GPU-h for 51.4M documents, close to what the linear extrapolation predicted.

Figures

pipeline/assets/model_selection/ holds the five figures behind the numbers above — cost_vs_quality.png, quality_ranking.png, quality_dimensions.png, robustness_by_safety.png, judge_calibration.png — and make_plots.py, which regenerates them. The script carries the measured values inline, so it runs without any evaluation output; one axis label still reads "FineWeb/dolma3", a leftover from an earlier corpus plan, and the pool is Dolma 3 throughout.


Running it

Nothing in this stage replays end to end from the release: the two inputs it consumed, the human-review database and the evaluation run directories, are annotator data and multi-gigabyte model output. What ships is the machinery, the frozen judge prompt, and each generator's final prompt. Use it to run this process on your own candidates, not to replay ours.

Human review. The dashboard is NiceGUI, one process:

uv run python -m pipeline.dashboard        # DASHBOARD_PORT, default 8600

/annotate is the review interface (pipeline/dashboard/charter_seed.py); seed sampling is pipeline/charter/seed/sampling.py, and charter.seed.dataset in configs/config.yaml points it at the stratified Dolma 3 sample.

Prompt optimization. Needs an authenticated claude CLI and an API key for the generator endpoint:

uv run python -m pipeline.charter.improve.loop --role judge --mode reflection
uv run python -m pipeline.charter.improve.loop --role generator \
    --mode reflection

--aliases a,b restricts it to specific models. A single cross-iteration without the agent is uv run python -m pipeline.charter.improve.

Generator comparison. Generate and judge, then rank:

uv run python -m pipeline.charter.eval eval-generators --run-id my_run \
    --mode reflection
uv run python -m pipeline.charter.eval rank-generators my_run

Candidates come from charter.eval.generator_eval.candidates. The commented-out blocks there record who took part in the paper's comparison; their api_name values are deployment ids on an internal inference endpoint and will not resolve anywhere else. The adjacent hf_slug is the real model. The paper's own run directories (ref_v3, ref_v4_qwen) are not shipped, so rank-generators ref_v3 ref_v4_qwen has nothing to read.

Throughput. Point the estimator at a running server:

uv run python -m throughput_estimations.estimate \
    --model-alias qwen3.5-35b-a3b --role generator --mode reflection \
    --n-samples 200 --data-path "$SPP_SUBSAMPLE_DIR/annotated" \
    --n-nodes 1 --gpus-per-node 4

sweep.sh submits one SLURM benchmark job per sglang configuration and sweep_analyze.py tabulates the results; tune_moe_kernels.sh is an sbatch script, so it goes through tools/submit.sh. check_sglang_env.sh verifies the container before you spend a node on either.