composer-replication-framework / docs /adrs /ADR-011-sdpo-alignment-indices.md
Codeseys's picture
architect: ADR-011/012/013 + research (alignment-index fix, review-findings closure, LMA channel-ladder)
b4a584a
|
Raw
History Blame Contribute Delete
5.24 kB
---
status: accepted
date: 2026-05-29
amends: ADR-008
deciders: [Codeseys, ARIA]
---
# ADR-011: Collator-emitted SDPO alignment indices (close the strict-guard regression)
## Context and Problem Statement
The 2026-05-29 cross-family review of ADR-008 found the SDPO student/teacher
alignment guard was a *shape-only* check (`student_logits.shape ==
teacher_logits.shape`), which does not establish token-level alignment because
the teacher context has a hint inserted at the error turn (shifting response
tokens right). The fix made `ComposerReplicationTrainer._compute_sdpo_loss`
**require** explicit `student_response_idx` / `teacher_response_idx` LongTensors
and `torch.gather` the aligned post-hint logits before JSD, raising in strict
mode (the default) when they are absent.
**Regression introduced:** the production collator
(`composer_replication/trainer/data_collator.py`) does NOT emit those index
tensors, so the default (strict) SDPO path now raises against the real collator.
This ADR closes that gap.
The collator already solves the hard alignment problem: `_build_aligned_student_for_sdpo`
builds a student sequence that mirrors the hint-conditioned teacher by inserting
a placeholder system-message of identical token length where the teacher has the
hint, and `_build_chat_aligned_mask` (per-message `apply_chat_template`
prefix-delta subsequence matching) marks the post-hint recovery-turn content
tokens with 1 in both `sdpo_loss_mask` (teacher) and `response_mask` (student).
So the 1-positions of both masks already correspond to the same logical tokens.
Research: `/tmp/composer-research/r1-alignment-indices.md` (DeepSeek V4 Pro,
2026-05-29).
## Decision Drivers
- The loss already requires the indices; the collator must supply them or strict
SDPO is unusable (it raises).
- The indices are *derivable from masks the collator already computes correctly*
— no extra tokenizer calls, no new alignment logic.
- The contract must be forward-compatible: if the placeholder trick is ever
dropped for dynamic-length alignment, distinct student/teacher indices still
describe the alignment.
- Ragged K (different rows have different #error-turn tokens) must be handled
without silent padding-token contributions to the JSD.
## Considered Options
- **A. Derive indices on-the-fly inside the loss from `sdpo_loss_mask`** — couples
the loss to the collator's placeholder implementation detail; rejected.
- **B. Collator emits explicit `student/teacher_response_idx` + `*_valid` masks,
derived from the existing chat-aligned masks** (chosen).
- **C. Drop the index requirement, revert to shape-check** — re-opens the P0 the
review caught; rejected.
## Decision Outcome
Chosen: **Option B.** The collator emits four new batch keys when SDPO is active:
`student_response_idx` (B, K_max), `teacher_response_idx` (B, K_max),
`student_response_valid` (B, K_max bool), `teacher_response_valid` (B, K_max bool).
A `_mask_to_padded_indices(mask, pad_sentinel=-1)` helper converts a (B, T)
response mask to a padded (B, K_max) index tensor + validity mask (sentinel -1
for ragged padding). The loss masks sentinel positions by building an
`aligned_labels` tensor (1 where valid, -100 elsewhere) passed to
`generalized_jsd_loss` (which already honors the -100 ignore convention).
### Consequences
- **Positive**: strict SDPO works against the real collator; the silent-misalignment
P0 stays closed; no extra forward/tokenizer passes.
- **Positive**: forward-compatible — distinct indices survive a non-placeholder future.
- **Neutral**: a debug-mode assertion `(s_idx == t_idx)[valid].all()` can verify the
placeholder trick is still intact when sequences are same-length.
- **Negative**: +4 batch keys; documented in the collator output contract.
## Acceptance gate (must be green before status flips to accepted)
- [ ] `_mask_to_padded_indices` implemented; ragged-K rows pad to K_max with
sentinel -1 + a `*_valid` bool tensor. Unit test: 2 rows with K=3 and K=1 →
(2, 3) idx with row-1 tail = -1 and valid[1] = [T,F,F].
- [ ] `ComposerDataCollator.__call__` emits the 4 keys whenever
`sdpo_loss_mask` + `response_mask` are present. Unit test asserts presence +
shapes + that `student_response_idx == teacher_response_idx` at valid
positions for the same-length placeholder path.
- [ ] `_compute_sdpo_loss` masks sentinels via `aligned_labels` (1/-100); a
sentinel position contributes 0 to the JSD. Unit test: a 2-row batch with
ragged K produces a finite loss and the K=1 row's padding doesn't leak.
- [ ] End-to-end: real `ComposerDataCollator` (with a stub tokenizer + a hint
generator) → batch → `_compute_sdpo_loss` runs in **strict mode** without
raising and returns a finite, positive loss. (This is the regression the ADR
closes — it must be a test, not a claim.)
- [ ] No regression: the existing alignment tests in
`test_dr_grpo_config_and_alignment.py` still pass.
## More Information
- `/tmp/composer-research/r1-alignment-indices.md` — full design + code sketch.
- ADR-008 — the strict-guard fix this ADR completes (amends).
- `composer_replication/trainer/data_collator.py` `_build_chat_aligned_mask`,
`_build_aligned_student_for_sdpo`.