awslabs / awslabs/keys_values

Chunked backward fails when a generated region spans 3+ chunks (autograd_hooks annotation adjacency)

Open
#148 3 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.