Chunked backward fails when a generated region spans 3+ chunks (autograd_hooks annotation adjacency)
- Dominant language
- Python
- Stars
- 18
- Forks
- 8
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
During GRPO training with **long sampled completions** (up to 2600 generated tokens, `chunk_size=1024`), the chunked backward fails stochastically in the replay unpack:
```
File ".../keys_values/kvcache/gradient/autograd_hooks.py", line 785, in unpack_hook
x = self._unpack_from_annotation(annotation)
File ".../keys_values/kvcache/gradient/autograd_hooks.py", line 852, in _unpack_from_annotation
raise ValueError(
ValueError: Annotation scatter-value (20,1): (4, 4, 8192, 128): final chunk_idx = 3, must be in [1, 2]
```
`_unpack_from_annotation` assumes the final buffer for a `(layer, kind)` node is at most **one chunk ahead** of the annotation being unpacked (`final_idx in [chunk_idx, chunk_idx + 1]`). In this failure the final buffer is two chunks ahead (3 vs annotation chunk 1), so the invariant breaks and backward dies.
## Where it shows up (and where it never does)
- **Never** in our QA-style GRPO runs: 27-32k-token prompts but only **32-token completions** (hundreds of runs, zero occurrences).
- **Stochastically** in long-generation runs: html-to-TSV extraction (LongProc) with `max_new_tokens=2600`, eos-terminated (ragged per-row lengths). One seed crashed at ~step 30; a parallel seed with different sampling is past step 100 without issue.
So the trigger involves a *generated* (rather than prefilled) region spanning **3+ backward chunks**, with per-row eos raggedness making it sampling-dependent.
## Config of the failing run
Qwen2.5-7B-Instruct (28 layers, GQA 4 KV heads, head 128), bf16, single L40S.
`h2o-torch-quantized8`, `cache_length=8192`, `grace_period=512`, batch (group) 4,
prompt ~14k tokens, `max_new_tokens=2600`, `chunk_size=1024`, `layers_per_cell=1`,
gradient accumulation 2 micro-steps, RLOO advantages. Failing annotation:
layer 20, `scatter-value`, chunk 1, buffer `(4, 4, 8192, 128)`.
## Repro status (honest)
Not yet minimized. CPU attempts with tiny models (2 layers, `cache_length=64`, completions spanning 3 chunks of 16, eos-ragged rows, quantized-8 buffers, grad accumulation, RLOO) do **not** trigger it, so the adjacency violation needs something extra that I have not isolated -- possibly an interaction between eviction timing (`grace_period`), ragged eos stopping near chunk boundaries, and which chunk a `scatter-value` annotation lands in. I can share the full job log and the exact record/seed on request, and I am happy to run instrumented builds (`debug_print_annotations=True` reaches ~step 30 in about 8h on our hardware).
## Workaround
Setting `chunk_size >= max_new_tokens` (so the generated region spans at most 2 chunks) avoids it; that is what our runs use now. Costs memory (larger backward cells), so it is not a fix for very long generations.
## Why it matters
Any RL/SFT use with long *generated* sequences through `LongContextGradientModel` will hit this eventually (it is sampling-dependent, so it looks like flaky infrastructure until you find the invariant). Short-completion training is unaffected.
Contributor guide
Research direction
Start in keys_values/kvcache/gradient/autograd_hooks.py, especially unpack_hook and _unpack_from_annotation around lines 785 and 852, and trace how annotation and final buffer chunk indices are produced. Reproduce or instrument a long generated, EOS-ragged run with 3+ chunks, comparing eviction and grace-period timing. Done means backward completes without violating the chunk-index invariant.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100