Skip to content

Preserve identity-distinct autograd views in StreamingAttentionKLDivergenceReducer and add SSHR integration test - #262

Open
stephandooper wants to merge 1 commit into
segment-weiss-backupfrom
codex/update-streamingattentionkldivergencereducer.forward
Open

stephandooper wants to merge 1 commit into
segment-weiss-backupfrom
codex/update-streamingattentionkldivergencereducer.forward

Conversation

@stephandooper

Copy link
Copy Markdown
Collaborator

Motivation

  • Ensure streaming KL reducer presents identity-distinct views of its inputs (matching AttentionKLDivergenceReducer passthrough) so reducer head resolution can disambiguate repeated tensor occurrences in flattened outputs.
  • Preserve the original autograd graph for downstream backward replay by avoiding cloning or detaching the tensors.
  • Add a focused integration regression that reproduces SSHR's layout where attention tensors are reused as auxiliary inputs and later as primary inputs to KL reducers.

Description

  • Modify StreamingAttentionKLDivergenceReducer.forward to return identity-distinct views using view_as, store the pair in _last_inputs, set _last_output to the student view, and return the pair; no cloning or detaching is performed. (file updated: lightstream/core/reducer/kl.py).
  • Add a small SSHR-like test fixture SSHRReducerLayoutNet and a regression test_sshr_layout_maps_kl_reducers_to_primary_occurrences that builds the flattened outputs and asserts KL reducer head indices and output shapes. (file updated: tests/test_scnn.py).
  • The change keeps autograd connectivity intact while giving each view a unique Python identity so _resolve_reducer_head_map can map reducers to their primary occurrences instead of earlier auxiliary occurrences.

Testing

  • Ran a focused Python check that instantiated StreamingAttentionKLDivergenceReducer, asserted the returned views are distinct objects, verified _last_inputs/_last_output, and confirmed gradients propagated through the original tensors; this check passed.
  • Ran bytecode/format checks via python -m compileall -q lightstream tests/test_scnn.py which succeeded.
  • Ran the focused pytest tests/test_scnn.py::test_sshr_layout_maps_kl_reducers_to_primary_occurrences which exercised the new integration test, but it failed in this environment due to absence of an NVIDIA driver (CUDA initialization error) despite attempting to run with CPU/statistics settings; the failure is environmental rather than indicating a logic regression.

Codex Task

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant