exp100: only-correct constrained decoding to regenerate contacts-v1 train docs#101
Open
timodonnell wants to merge 18 commits into
Open
exp100: only-correct constrained decoding to regenerate contacts-v1 train docs#101timodonnell wants to merge 18 commits into
timodonnell wants to merge 18 commits into
Conversation
…rain docs Scaffold for issue #100 (builds on exp98/PR #99). For each training protein, sample N=10 constrained rollouts from the tuned contacts-v1 1.5B where the model may only ever emit a true, not-yet-emitted contact (<end> masked until all true contacts are out), then re-score each under the unmodified model and keep the lowest structure-section NLL as the regenerated document. - constrained_grammar.py: pure-Python only-correct FSM (3-token cycle) + vLLM logits-processor adapter with prompt-offset auto-detect; unit-tested. - gen_constrained_worker_vllm_tpu.py: N=10 constrained generate (pass A) + prompt_logprobs unmodified-NLL scoring (pass B) + select + resume, on iris TPU. - select_targets/gen_prompts (k=10)/rollout_metrics reused from exp98; fresh publish_to_hf + aggregate_results for the documents/nll schema. - tests: FSM units + full seq->position->text->parse round trip (17 passing). - data/targets.parquet copied from exp98 (identical 1000 proteins) for a directly-comparable validation run. No compute launched yet: Phase-0 spike (logits_processors on vLLM-TPU) + the validation run are pending go-ahead. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…orker Phase-0 validation on the local A5000 against the tuned 1.5B (HF bf16) passed: every constrained document is 100%-correct + full-recall, a well-formed contacts-v1 doc of length exactly 3*n_gt+1, and the folded-in struct_nll matches an independent teacher-forced forward pass to ~1 nat. - constrained_grammar.ContactConstraint now captures each realized token's pre-mask full-vocab logprob (struct_nll) in-line, removing the prompt_logprobs scoring pass — prompt_logprobs returns None on the iris JAX/TPU stack (exp89). - gen_constrained_worker_hf_gpu.py: batched HF-transformers worker (the proven local path); a target's N rollouts share prompt+gen length so they batch with zero padding and finish together. Same outputs as the vLLM worker. - vLLM worker: dropped Pass B / prompt_logprobs / max_logprobs accordingly. - README/publish updated; TPU logits_processor support still to be spiked. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
10,000/10,000 rollouts 100%-correct + full-recall; 2.85 A5000-hours; best-of-10 struct NLL averages 18.9 nats below the mean rollout. Adds per-target CSV + plot, README results/conclusion. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-proven, TPU-portable) Source review of marin tpu_inference (rev 29faff43) settled the iris path: custom LogitsProcessor.apply() is NOT invoked on TPU, but the structured-output grammar bitmask IS applied on-device in JAX, and prompt_logprobs IS supported (the exp89 "None on TPU" note was stale). So the constraint is expressed as a custom vLLM StructuredOutputBackend. - only_correct_backend.py: OnlyCorrectBackend/Grammar wrapping an incremental FSM (OnlyCorrectMatcher) whose fill_bitmask emits the per-step allowed-token bitmask; register() monkeypatches grammar_init + _validate_structured_output (no plugin hook in V1). NLL via a separate prompt_logprobs pass. - constrained_grammar.py: added OnlyCorrectMatcher (incremental, O(1)/token) + pack_allowed_bitmask, cross-checked against legal_token_ids in tests. - gen_constrained_worker_vllm.py: vLLM worker (GPU now; iris after registering the backend as a vllm.general_plugins entry point). Same outputs as the HF worker. - Verified end-to-end on the A5000 (vllm 0.11.2): 10/10 correct on 6 targets, NLL matching the HF worker within sampling noise, ~790 tok/s (~2x HF). - Removed gen_constrained_worker_vllm_tpu.py (used the V1-invalid per-request logits_processors callable API). - README: backend/iris-port section; tests (20 passing). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…s-ready) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The marin TPU fork is vLLM 0.20.1 (processor.py -> input_processor.py; validation moved to SamplingParams._validate_structured_outputs). register() now pre-sets self.backend in grammar_init and delegates to the original (avoids reimplementing the version-specific tail), and patches whichever validation hook exists. Verified still 10/10-correct on local GPU (0.11.2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… register) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…lect among correct Full iris run surfaced a rare (~0.2%) grammar fill/accept desync in vLLM 0.20's parallel threaded bitmask-fill path (triggered when >128 concurrent structured requests; the 3-target calib used the serial path). Cap max_num_seqs=96 to force the proven serial path, and select the best rollout only among 100%-correct ones so a truncated reject can never win on its shorter NLL. Applied to both workers. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…esync) Root cause of the TPU 'grammar rejected' rate: our incremental FSM state momentarily drifts from the true token stream under vLLM 0.20's fill/accept scheduling (rejected tokens were all valid contacts-v1 tokens -> tracking desync, not a mask miss; xgrammar doesn't hit it). Fix: accept_tokens never returns False (False makes vLLM terminate the request); on any mismatch we resync by replaying the authoritative token history, and only mark a rollout 'broken' if the true stream is genuinely illegal (mask actually missed). Still 10/10 on local GPU. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-hour) vLLM structured-output-backend worker on one v5p-8: all 1000 selected documents 100%-correct (mean 9.985/10 rollouts correct, 0 targets failed), NLLs matching the GPU run, 1226 tok/s (tp=4), ~1.04 v5p-8-hours for 10k documents (~3x the A5000 per-unit rate). Adds TPU per-target CSV + plot + README results/conclusion. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…s) + context-capped max_tokens select_round0_all.py: keep every round-0 protein (no L/contact filters), GT parsed from the existing document (no pyconfind). vllm worker caps max_tokens to the context left after the prompt so the longest proteins never exceed max_model_len. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Per-target write retries transient GCS/oauth SSL blips (6 attempts, backoff); main loop tolerates stragglers (re-run resumes). Needed for the ~941k round-0 prompt build. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… shards) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-scale run Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…exists calls) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… on load) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… run) At 941k-protein scale a few resampled prefixes reach the full 8192-token context, leaving no room for the 3*n_gt+1 structure section. The old worker requested max_tokens>=1 on such a prompt, so vLLM raised "prompt (8192) + output (>=1) > max_model_len" and the whole shard crash-looped (shard 0 exhausted its retries and died; shard 6 was on the same path). Filter each target's rollouts to those with room >= 3*n_gt+1; if none fit, write a resume-safe skip marker (r=-1 nll sentinel + skipped:true document) so the target is counted done and never retried, instead of crashing. aggregate/publish drop the sentinel and report the skipped count. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ontext The first skip-guard fixed generation but not the prompt_logprobs scoring pass, which re-feeds prompt+generated as a prompt with max_tokens=1. A tight-fit rollout with prompt+gen == max_model_len (8192) then hit prompt(8192)+output(1) > 8192 and crash-looped the shard (shards 0, 5, 6 each exhausted 50 retries and died). Reserve one token: budget = max_model_len - 1, require room >= 3*n_gt+1 against that budget, and cap max_new to it -> prompt+gen <= 8191, so the scoring pass's forced token lands exactly at 8192. Tighter targets are skipped as before. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Scaffolds the data-generation method for #100: regenerate the contacts-v1 training set by sampling only-correct documents from Eric's tuned 1.5B (eval loss 2.7566) and selecting per protein by unmodified likelihood. Builds directly on exp98 (PR #99), reusing its iris/vLLM-TPU rollout scaffold.
What's here (no compute launched yet)
constrained_grammar.py— pure-Python only-correct FSM. The contacts-v1 tokenizer is WordLevel (1:1, whitespace-separated), so a statement is exactly the 3-token stream[<contact>, <pi>, <pj>]→ a clean 3-token cycle: force<contact>while true contacts remain; first/second endpoint restricted to positions that (complete) a still-remaining true contact;<end>masked until every true contact is emitted. So every finished doc is precision = recall = 1.0 by construction. Ships as a vLLM logits-processor with prompt-offset auto-detect (robust to whether vLLM prepends the prompt topast_token_ids).gen_constrained_worker_vllm_tpu.py— per target, N=10 constrained rollouts (pass A: per-rolloutlogits_processors) → re-scoreprefix + generatedwithprompt_logprobsunder the unmodified model (pass B, exp89's proven NLL path) → keep lowest structure-section NLL. Persists per-rollout NLLs + a built-in 100%-correct check, the selected document verbatim, and all N. Resume-on-restart.select_targets.py,gen_prompts.py(k=10),rollout_metrics.py. Freshpublish_to_hf.py+aggregate_results.pyfor thedocuments/+nll/schema.data/targets.parquetcopied from exp98 (identical 1000 proteins) so the validation run is directly comparable to exp98's free rollouts.Phasing / status
logits_processorswork on vLLM-TPU (proven forprompt_logprobs, not yet for logits processors). Fallback: run masked-sampling on a single GPU, score on TPU.data/contacts-v1-train-only-correct-exp100/).See the README.
🤖 Generated with Claude Code