Safe eviction defaults (#140), GRPO single-epoch cache reuse, HELMET GRPO driver - #142
Safe eviction defaults (#140), GRPO single-epoch cache reuse, HELMET GRPO driver#142mmjerge wants to merge 15 commits into
Conversation
End-to-end results: GRPO on HELMET through an evicting sparse cacheRan the full pipeline in this PR on HELMET Cross-evaluation: training cache × inference cache (n=100 held-out records)
Findings
Training dynamics / footprint
Repro: |
|
Follow-up: the round-2 arms (F1 partial-credit reward, group size 8, 300 steps) cross-evaluate the same way (n=100):
Consistent with round 1: all cells within noise of the untrained base (0.330 / 0.320), sparse-inference cost ~1–3pp EM. The parity conclusions hold; the "learning gains" run (higher lr, more steps/prompts, multiple seeds) remains future work. |
Hot-run results: learning through the evicting cache is real (with caveats)Second round with lr 5e-6, 8 rollouts x 2 accumulated prompts x 400 updates (gradient accumulation added in this PR), F1 partial-credit reward, single seed, n=100 cross-eval. nq @ 8k (EM / F1):
hotpot_qa @ 8k:
Findings:
|
|
@mmjerge I am aware of attention sink. And indeed, when I evaluated a pre-trained checkpoint with the This is also why I implemented I am OK with have a default grace period with But I do not see why this would affect H2O (say). If the model uses the first token for attention, H2O would keep it in the cache (as it is heavily used). In the pre-trained checkpoint evals, I did not see bad results just because of there being no grace period, at least not for H2O. |
|
Correction on the "reward hacking" hypothesis, after inspecting generations + seed-1 replication. Sample generations (distinct NQ questions, dense eval cache):
So the dense arm did not become verbose — the F1 partial credit taught it to strip scaffolding (terse answers score higher token-F1). The EM collapse is substantially a metric-style artifact: HELMET NQ targets often include leading prepositions ("in Super Bowl LII", "on October 21, 2016"), and Seed-1 replication (nq, n=100) reproduces the seed-0 pattern: h2o-trained 0.44/0.45 EM (dense/h2o eval) vs base 0.33; dense-trained 0.31/0.29. So the pattern is robust across 2 seeds, but the safe interpretation is:
hotpot seed-1 arms are still running; will append when complete. |
|
@mseeger Agreed — I've replied in detail on #140 (with an isolated example), and updated this PR accordingly:
Also pushed since your comment: an |
Evicting KV caches produced garbage generation (immediate EOS) on real models whenever the prompt exceeded cache_length. The root cause is attention-sink eviction, not incorrect attention math: trained LLMs place large attention mass on the first few tokens (BOS / chat-template header) and use them as no-op sinks (Xiao et al., StreamingLLM). Evicting them makes the softmax redistribute that mass onto content tokens, which collapses generation. Toy random models have no sink structure, which is why unit tests passed. Instrumentation confirmed the mechanics are exact: on a partial-write chunk the cached keys/values match a dense reference gathered at token_positions() exactly, and the attention output matches a manual masked-softmax reference to ~1e-7. A 1-layer model shows no divergence; differences appear only once hidden states compound across layers, which is expected lossiness. KVCacheFactory now injects safe defaults unless the caller specifies them (explicit values, including 0, are respected): - lastrec: init_grace_tokens = min(16, cache_length // 8) - h2o / h2o-vlen / qh2o / qh2o-vlen: grace_period = cache_length // 16 (h2o-orig excluded; it rejects grace_period) Measured on Qwen2.5-0.5B (A10G, retrieval with the answer in the last record, evicting cache at cache_length=1024): lastrec 0.00 -> 0.90, which equals the no-eviction baseline; h2o 0.00 -> 0.70. test_utils.create_kv_cache pins the protections to 0 so unit tests keep exercising the raw eviction mechanics.
Builds on the standalone GRPO loop from awslabs#128. Correctness - Alignment fix: the scoring/gradient passes are now fed full_ids[:, :-1], so completion token p is scored against the logits at position p-1 (next-token alignment). Previously they received the full sequence, i.e. an off-by-one. - Completion mask: padding after an early stop no longer contributes to the loss, and num_target_entries normalizes by the real token count. Single-epoch cache reuse (the optimization) - GRPOLossHeadModel.set_batch now takes advantages first, with old_logps optional. When old_logps is None (single-epoch GRPO, the default), the old log-prob is policy_logp.detach() from the gradient pass's own forward, so the importance ratio is exactly 1 and the separate scoring forward pass is eliminated. This is what TRL does for num_iterations=1. Pass old_logps explicitly for multi-epoch updates. - grpo_step gains rescore_old_logps (recompute old log-probs via a separate compute_logprobs pass, for multi-epoch or A/B comparison) and profile (per-phase timings: gen_time_ms / score_time_ms / grad_time_ms). Generation - batched_generate_fn gains opt-in return_logprobs, yielding (tokens, logprobs, active_mask) so a rollout can capture sampling log-probs and a padding mask without a second pass. Default behaviour is unchanged, so the evaluation harness is unaffected. - batched_generate_fn gains no_inference_mode. The KV cache lazily moves its buffers to the compute device on the first forward; when that forward ran under torch.inference_mode() (the rollout), the buffers became inference tensors that the following gradient pass could not update in place, raising "Inplace update to inference tensor outside InferenceMode". The GRPO rollout now runs under torch.no_grad(); eval and generation-only callers keep inference mode. GPU-only failure (on CPU the buffers are already on-device). - rollout.generate_completions_with_logprobs: rollout generation that also returns per-token log-probs and the completion mask. Tooling - build_ext.py: add sm_86 (A10G / A10 / RTX 30xx) and a PTX JIT fallback; the FlashInfer kernels previously targeted only sm_80 / sm_90. - examples/grpo_qwen.py: runnable GRPO demo with flags for both paths and per-phase profiling. - examples/grpo_profile.py: real-model completion-length sweep comparing the optimized and rescore paths (timings, peak memory, throughput). - examples/grpo_context_sweep.py: dense vs H2O across context length (peak memory and step time, with graceful OOM handling). - examples/longqa_task.py, examples/longqa_quality.py: controlled long-context key-value retrieval task and a dense-vs-sparse quality comparison, used to investigate awslabs#140. - docs/GRPO_CONTEXT_SCALING.md: measured dense-vs-H2O context scaling on an A10G (H2O stays flat at ~14.9 GB from 2k to 32k while dense grows and OOMs at >=24k), with methodology and caveats.
…_length needs prefill
…eval, jsonl metrics)
…are trained in place)
…ing; add 2x2 cross-eval (train-cache x eval-cache)
…); grpo_helmet: --prompts-per-update
… EM-anchored em_f1 reward
|
Rebased onto |
…nchanged Maintainer decision on awslabs#140: keep defaults simple -- change only lastrec (init_grace_tokens, attention-sink rationale) and leave the H2O family's grace_period at 0; the recent-tail effect is application-level tuning (and normalize_scores=True is an alternative mitigation). The H2O grace_period our experiments rely on moves into the example/RL configs explicitly (grpo_helmet, crosseval, longqa_quality).
|
Seed-1 replication complete (same recipe: lr 5e-6, group 8 x 2 accumulated prompts x 400 updates, F1 reward, n=100 cross-eval). Both seeds agree on the direction of every cell. hotpot_qa @ 8k, seed 1 (EM / F1; seed-0 in parentheses):
Two-seed summary across both tasks (EM under dense eval, vs base):
Consistent, replicated pattern: each arm learns where its rollouts see the evidence. On single-hop nq, the H2O arm improves robustly (+11-12pp both seeds) while the dense arm's EM drop is the answer-style artifact documented above. On multi-hop hotpot — where H2O inference itself costs ~11pp (dispersed evidence gets evicted) — the dense arm improves (+9-12pp) while the H2O arm stays ~flat under dense eval (though it does gain under matched H2O inference: 0.18 -> 0.23/0.29 across seeds). Takeaways: (1) GRPO through an actively evicting cache trains stably and can deliver real gains (nq), at ~half the memory; (2) on dispersed-evidence tasks the eviction-induced information loss in rollouts limits the learning signal — the same regime where H2O inference costs accuracy; (3) the F1-only reward's style drift is confirmed across seeds — subsequent runs use the EM-anchored |
Consolidated results: algorithms x caches x tasks (final for this round)All: Qwen2.5-0.5B-Instruct, HELMET @ 8k (~7k-token prompts), single A10G, 400 updates (x2 accumulated prompts), lr 5e-6, n=100 cross-eval. "H2O" = h2o-torch-quantized8 @ 4096 slots (actively evicting; grace_period=cl/16 set at application level per #140 discussion). EM / F1 shown as dense-eval | H2O-eval. nq @ 8k (base: 0.33/0.18 | 0.32/0.14):
hotpot_qa @ 8k (base: 0.29/0.15 | 0.18/0.11):
Findings
RLOO mode and the SFT driver live on stacked branches ( |
|
Hello, I am a bit confused about this PR. Maybe it mixes two things?
Could you split the two? As I said, I'd like to be conservative on the defaults, but I am more than happy with new code on the RL front, if this helps. |
|
Please also try to document findings in markdown files checked in. The PR messages will be lost in the end, or very hard to find. |
Restores keys_values/kvcache/factory.py and test_utils.py to main; the lastrec attention-sink default now lives in a separate PR so this one is RL-only. Application-level grace settings in the GRPO/eval drivers are unaffected (they pass grace_period explicitly).
Consolidates the PR-thread results into docs/GRPO_HELMET_RESULTS.md: round-1 parity cross-eval, round-2 learning gains + the answer-style/ metric artifact analysis and seed replication, and the round-3 algorithms x caches x tasks matrix (SFT / GRPO / RLOO), with caveats and repro commands.
|
@mseeger Both points addressed:
Re-verified after the split: |
records[:n] in file order covered only 17 distinct questions at n=100 on nq (the eval split holds ~6 context variants per question), inflating variance well past the nominal-n SE. Select one instance per question first, then second variants, etc. At n=100 on nq this now covers all 100 eval questions.
Documents the unseeded-eval-split finding (awslabs#145) and its impact on the published tables (17 distinct questions at n=100; wider error bars), and adds the first corrected numbers: a base-model H2O budget sweep under the seeded split with full question coverage. Headline: sparse inference costs ~10pp EM on nq at 4096 slots (not ~1pp), ~7pp on hotpot, with a cliff at 1024 slots. Full-matrix re-baselining queued.
| ) | ||
|
|
||
|
|
||
| def _sampled_token_logprobs( |
There was a problem hiding this comment.
If this is specific to the RL work, it should move to keys_values/rl. Otherwise, you'd have to argue why it is more generally needed (I do not get it from the docstring).
There was a problem hiding this comment.
Agreed — it's RL-specific. Moved it to keys_values/rl/grpo/rollout.py as sampled_token_logprobs (next to the rollout code that captures old-policy log-probs). The generic generation loop lazy-imports it inside the return_logprobs branch to avoid a generate ↔ rl import cycle. Done in the latest push.
| "-U__CUDA_NO_HALF2_OPERATORS__", | ||
| "-gencode=arch=compute_80,code=sm_80", | ||
| "-gencode=arch=compute_90,code=sm_90", | ||
| # SASS for the common data-center / workstation GPUs, plus |
There was a problem hiding this comment.
Can you explain why this change is needed? This is a script needed to build the FlashInfer stuff by Vihang.
- Do you need this for the different architectures?
- If so, did you test this on these?
- This will not break the current setup?
There was a problem hiding this comment.
You're right that this doesn't belong in the RL PR — I've reverted build_ext.py to match main in the latest push, so this PR no longer touches Vihang's build script.
For context on what it was: I'd added sm_86 (A10G/A10/RTX-30xx) plus a PTX fallback line, because the A10G is the box I run the GRPO experiments on and the stock sm_80+sm_90 build won't load FlashInfer there. The retained sm_80/sm_90 SASS lines were untouched, so it wouldn't have broken the A100/H100 setup — but it was orthogonal to RL and I hadn't tested it beyond the A10G, so it's better out of scope here. If it's useful I'll open a separate one-line PR against main for the A10G arch, with you/Vihang as reviewers.
…ch change - sampled_token_logprobs moved from generate/base.py to rl/grpo/rollout.py (RL-specific; generate loop lazy-imports it to avoid an import cycle) - build_ext.py reverted to match main (sm_86/PTX arch change was orthogonal to the RL scope of this PR; will propose separately if wanted)
2 participants
Summary
Fixes #140 (safe eviction defaults), improves the standalone GRPO loop from #128
(single-epoch cache reuse, completion mask, an alignment fix, and a GPU
inference-mode fix), adds
sm_86to the FlashInfer build, and adds a GRPOtraining driver for HELMET tasks plus profiling tooling.
Fix #140: safe eviction defaults (attention sinks + grace period)
Evicting KV caches produced garbage generation (immediate EOS) on real models
whenever the prompt exceeded
cache_length. Instrumentation showed theeviction/attention mechanics are exact (cached keys/values match a dense
reference gathered at
token_positions()exactly; the attention output matchesa manual masked-softmax reference to ~1e-7; a 1-layer model shows no divergence
at all). The cause is attention-sink eviction: trained LLMs put large
attention mass on the first tokens (BOS / template header) and use them as
no-op sinks (Xiao et al., StreamingLLM); evicting them collapses generation.
Toy random models have no sink structure, which is why unit tests pass.
KVCacheFactorynow injects safe defaults unless the caller specifies them(explicit values, including
0, are respected;test_utils.create_kv_cachepins them to
0so unit tests keep exercising raw mechanics):lastrec:init_grace_tokens = min(16, cache_length // 8)h2o/h2o-vlen/qh2o/qh2o-vlen:grace_period = cache_length // 16(
h2o-origexcluded — it rejectsgrace_period)Measured (Qwen2.5-0.5B-Instruct, A10G, bf16, ~2k-token prompts, n=40,
substring exact match; before the fix every evicting config scored 0.000):
Recency recovers to within noise of dense — generation through an evicting
cache is functional again — while the needle rows now reflect genuine eviction
lossiness (the expected quality frontier) instead of the collapse.
Grace slots reduce a cache's forward capacity, so chunk sizes are capped at
cache_length - gracewhere applicable (rl/grpo/loop.py, examples).GRPO improvements (on top of #128)
GRPOLossHeadModel.set_batchmakesold_logpsoptional; when omitted (single-epoch GRPO), the old log-prob ispolicy_logp.detach()from the gradient pass's own forward, so theimportance ratio is exactly 1 and the separate scoring forward pass is
eliminated (mirrors TRL's
num_iterations=1behaviour).grpo_stepgainsrescore_old_logps(multi-epoch / A-B) andprofile(per-phase timings).full_ids[:, :-1], socompletion token
pis scored against the logits atp-1(was off by one).the loss or its normalization.
on the first forward; under
torch.inference_mode()(rollout or evalgeneration) they became inference tensors that the training pass could not
update in place ("Inplace update to inference tensor outside
InferenceMode").
batched_generate_fn/generate_completionsgainno_inference_modeand the GRPO paths use it; eval/generation-only callerskeep the default (faster) inference mode. GPU-only failure; on CPU buffers
are already on-device.
model drops from ~4.8 to ~0.13 (bf16 noise).
Tooling
build_ext.py: addsm_86(A10G / A10 / RTX 30xx) + PTX fallback; kernelspreviously targeted only
sm_80/sm_90.examples/grpo_helmet.py: GRPO training on HELMET datasets viaload_helmet_dev_eval, reward =sub_exact_match(same metric as the evalharness), periodic held-out eval, JSONL metrics. Validated end-to-end on
HELMET
nq@ 8k (prompts ~7k tokens, H2O cache 4096, real eviction).examples/grpo_profile.py(completion-length sweep, optimized-vs-rescore),examples/grpo_context_sweep.py(dense-vs-H2O memory/time across context;on a 24 GB A10G dense OOMs at >=24k while H2O stays ~flat, see
docs/GRPO_CONTEXT_SCALING.md),examples/longqa_task.py/longqa_quality.py(controlled eviction-quality probes used for Evicting KV cache corrupts generation (garbage output once prompt > cache_length) #140).Testing
test/kvcachesuite: 956 passed (includes the factory-default change;test_utilspins grace to 0 so existing eviction-mechanics tests areunchanged).
test/rl(from End-to-end GRPO with KeysAndValues KV cache #128) passes unchanged — theset_batchchange isbackward-compatible.
test/test_long_context.py(from Apply final norm ln_f and logit softcapping before head in LongContextInference #141),test/generate/test_base.py,test/finetune/gradient suites pass.HELMET
nqwith an evicting H2O cache runs cleanly (~13 s/step at ~7k-tokenprompts, 8.2 GB peak).
By submitting this pull request, I confirm that you can use, modify, copy, and
redistribute this contribution, under the terms of your choice.