docs/adrs/ADR-011-sdpo-alignment-indices.md

metadata

status: accepted
date: 2026-05-29T00:00:00.000Z
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.