--- 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`.