# SDPO Alignment Indices: Canonical Collator Design

**Status**: Recommendation (Option B with ragged-K safety)  
**Date**: 2026-05-29  
**Context**: Cross-family review found the SDPO loss alignment guard was shape-only; the fix (commit 2026-05-29) now **requires** `student_response_idx` / `teacher_response_idx` LongTensors from the collator. The production collator (`composer_replication/trainer/data_collator.py`) does not yet emit them, so strict SDPO raises.

---

## 1. What the loss now demands

`ComposerReplicationTrainer._compute_sdpo_loss` (lines 184–242 of `composer_trainer.py`) does:

```python
s_idx = inputs.get("student_response_idx")   # (B, K) LongTensor
t_idx = inputs.get("teacher_response_idx")   # (B, K) LongTensor
# … guard: raise if strict + missing …
vocab = student_logits.size(-1)
s_gather = s_idx.unsqueeze(-1).expand(-1, -1, vocab)  # (B, K, V)
t_gather = t_idx.unsqueeze(-1).expand(-1, -1, vocab)
student_aligned = torch.gather(student_logits, 1, s_gather)  # (B, K, V)
teacher_aligned = torch.gather(teacher_logits, 1, t_gather)  # (B, K, V)
# → generalized_jsd_loss(student_aligned, teacher_aligned, …)
```

It expects per-row indices into the sequence dimension (dim=1) selecting **K** aligned post-hint response-token positions. K is the number of error-turn response tokens in that row.

---

## 2. Option A vs Option B — Recommendation

### Option A: derive indices from the existing equal-length mask

Since the collator *already* builds same-length student/teacher sequences via `_build_aligned_student_for_sdpo` (placeholder system-message of identical token length at the hint-slot), both `sdpo_loss_mask` (teacher-side) and `response_mask` (student-side) mark the exact same token positions. We could simply do:

```python
student_response_idx = torch.nonzero(response_mask == 1, as_tuple=False)  # per-row
teacher_response_idx = torch.nonzero(sdpo_loss_mask == 1, as_tuple=False) # identical
```

**Pros**: zero new collator complexity; the mask is already correct.  
**Cons**: (a) couples the loss to an *implementation detail* (the placeholder trick) — if the collator ever drops same-length alignment, all rows silently break; (b) the mask selects a *subset* of the response tokens (only assistant content), while the indices could select *all* post-hint tokens including chat-template scaffolding; (c) `torch.gather` on equal-length identical indices is mathematically a no-op that wastes memory — the loss should just take a mask path when alignment is trivial.

**Verdict**: Reject. The mask path should remain as a *fallback* inside `generalized_jsd_loss` when the collator hasn't been upgraded, but the canonical emission is distinct indices.

### Option B (RECOMMENDED): emit distinct indices unconditionally

The collator computes both `student_response_idx` and `teacher_response_idx` *explicitly* during `_build_aligned_student_for_sdpo` / `_build_sdpo_fields`. Even when sequences are same-length, emitting the indices:
- Proves to the loss that alignment was *deliberately* solved (not accidentally same-length)
- Survives a future where the placeholder trick is replaced by dynamic padding
- Allows `teacher_response_idx` to differ from `student_response_idx` in any future generalization (e.g., the teacher omits some non-content tokens)

**This is the canonical design.** The remainder of this document specifies the exact tensor construction.

---

## 3. Design: constructing the indices from the existing mask

The collator already computes:
- `sdpo_loss_mask`: (B, T) with 1 at teacher post-hint error-turn content tokens, `ignore_index` (-100) elsewhere.
- `response_mask`: (B, T) with 1 at student assistant-content tokens (incl. post-hint), 0 elsewhere.

Because the collator's `_build_chat_aligned_mask` uses per-message `apply_chat_template` prefix deltas to place mask bits *exactly* on content tokens regardless of scaffolding, the 1-positions in both masks correspond to the **same logical token** in an aligned comparison.

### Step 1: per-row nonzero positions

```python
def _build_response_indices(
    mask: torch.Tensor,       # (B, T), 1=response, 0=ignore
    pad_sentinel: int = -1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Convert a per-row response mask to padded index tensors.

    Returns:
        idx:   (B, K_max) LongTensor — position indices, padded with sentinel
        valid: (B, K_max) BoolTensor  — True where idx is a real position
    """
    B, T = mask.shape
    rows = []
    for b in range(B):
        pos = torch.nonzero(mask[b] == 1, as_tuple=True)[0]  # (K_b,)
        rows.append(pos)
    K_max = max(r.numel() for r in rows) if rows else 0
    if K_max == 0:
        # No error sites in this batch — return empty sentinels
        return (
            torch.full((B, 0), pad_sentinel, dtype=torch.long, device=mask.device),
            torch.zeros(B, 0, dtype=torch.bool, device=mask.device),
        )

    idx = torch.full((B, K_max), pad_sentinel, dtype=torch.long, device=mask.device)
    valid = torch.zeros(B, K_max, dtype=torch.bool, device=mask.device)
    for b, row in enumerate(rows):
        k = row.numel()
        idx[b, :k] = row
        valid[b, :k] = True
    return idx, valid
```

### Step 2: unified emission in the collator's `__call__`

Inside `ComposerDataCollator.__call__` (after the `_build_aligned_student_for_sdpo` block, around line 191), add:

```python
# --- Emit SDPO alignment indices (ADR-008 gate) ---
if "sdpo_loss_mask" in out and "response_mask" in out:
    # Teacher-side: where does sdpo_loss_mask == 1?
    # Note: sdpo_loss_mask uses ignore_index (-100) for non-loss tokens.
    # We want positions where the value is exactly 1 (the in-loss marker).
    t_mask = (out["sdpo_loss_mask"] == 1)  # (B, T)
    t_idx, t_valid = _build_response_indices(t_mask)

    # Student-side: where does response_mask == 1?
    # response_mask is 0/1; 1 means assistant-response token.
    s_mask = (out["response_mask"] == 1)    # (B, T)
    s_idx, s_valid = _build_response_indices(s_mask)

    # When sequences are same-length and aligned by the placeholder trick,
    # s_idx will equal t_idx for every valid position. The loss can
    # optionally assert this in debug mode, but the canonical contract
    # is that the two index tensors describe the alignment, and they
    # MAY differ in future collator versions.
    out["student_response_idx"] = s_idx
    out["teacher_response_idx"] = t_idx
    out["student_response_valid"] = s_valid   # (B, K_max)
    out["teacher_response_valid"] = t_valid   # (B, K_max) — same max-K
```

### Step 3: sentinel handling in the loss

The loss currently does `torch.gather` unconditionally. The sentinel value (-1) would wrap around and select the *last* token — harmless but wasteful. Better: the loss should mask sentinel positions. Update `_compute_sdpo_loss`:

```python
# After gather:
student_aligned = torch.gather(student_logits, 1, s_gather)  # (B, K, V)
teacher_aligned = torch.gather(teacher_logits, 1, t_gather)  # (B, K, V)

# Build a (B, K) mask: True where BOTH indices are valid (not sentinel).
# When the collator guarantees s_valid == t_valid, use either.
if "student_response_valid" in inputs:
    aligned_mask = inputs["student_response_valid"]  # (B, K), already BoolTensor
else:
    aligned_mask = (s_idx >= 0) & (t_idx >= 0)  # sentinel=-1 guard

# Pass this as the labels mask to generalized_jsd_loss.
# The loss already handles labels != -100 masking; we repurpose it:
#   labels[b, k] = 1  if aligned_mask[b, k] else -100
aligned_labels = torch.where(
    aligned_mask,
    torch.ones_like(s_idx, dtype=torch.long),
    torch.full_like(s_idx, -100, dtype=torch.long),
)

return generalized_jsd_loss(
    student_logits=student_aligned,
    teacher_logits=teacher_aligned,
    labels=aligned_labels,
    …
)
```

---

## 4. Why this is canonical

| Property | How this design provides it |
|---|---|
| **Token-level alignment** | Indices are derived from `_build_chat_aligned_mask`, which uses per-message prefix deltas to locate content tokens inside the full chat-template tokenization — not naive segment concatenation. |
| **Ragged-K safety** | Pad to `K_max` with sentinel -1; emit a `*_valid` BoolTensor. The loss masks sentinels via `labels=-100` (standard HF ignore convention). No silent padding-token contribution. |
| **No additional forward passes** | The indices are computed from existing mask tensors inside `__call__` — zero extra tokenizer calls. |
| **Forward-compatible** | Emitting distinct student/teacher indices survives a future where the placeholder trick is replaced by dynamic-length alignment. |
| **Auditable** | A one-line assertion `(s_idx == t_idx).all()` in a debug build verifies the placeholder trick is still intact. |

---

## 5. Code sketch (full emission path)

```python
# === In ComposerDataCollator.__call__, after line 191 ===

def _mask_to_padded_indices(
    mask: torch.Tensor,          # (B, T) where 1 = valid position
    pad_sentinel: int = -1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Convert (B,T) boolean mask → (B,K_max) index tensor + (B,K_max) validity mask."""
    B, T = mask.shape
    # Per-row nonzero — torch.nonzero on a 2D bool tensor gives (N,2); reshape.
    nz = torch.nonzero(mask, as_tuple=False)  # (total_K, 2)
    # Group by row:
    counts = mask.sum(dim=1).long()           # (B,) — K per row
    K_max = int(counts.max().item()) if counts.numel() else 0
    if K_max == 0:
        return (
            torch.full((B, 0), pad_sentinel, dtype=torch.long, device=mask.device),
            torch.zeros(B, 0, dtype=torch.bool, device=mask.device),
        )
    idx = torch.full((B, K_max), pad_sentinel, dtype=torch.long, device=mask.device)
    valid = torch.zeros(B, K_max, dtype=torch.bool, device=mask.device)
    # nz[:, 0] are batch indices, nz[:, 1] are position indices
    batch_idx = nz[:, 0]  # (total_K,)
    pos_idx = nz[:, 1]    # (total_K,)
    # Build a per-batch write offset using cumsum
    offsets = torch.zeros(B + 1, dtype=torch.long, device=mask.device)
    offsets[1:] = counts.cumsum(dim=0)
    for b in range(B):
        start, end = offsets[b].item(), offsets[b + 1].item()
        k = end - start
        if k > 0:
            idx[b, :k] = pos_idx[start:end]
            valid[b, :k] = True
    return idx, valid

# --- Emission ---
if "sdpo_loss_mask" in out and "response_mask" in out:
    t_mask = (out["sdpo_loss_mask"] == 1)
    s_mask = (out["response_mask"] == 1)
    t_idx, t_valid = _mask_to_padded_indices(t_mask)
    s_idx, s_valid = _mask_to_padded_indices(s_mask)
    out["student_response_idx"] = s_idx
    out["teacher_response_idx"] = t_idx
    out["student_response_valid"] = s_valid
    out["teacher_response_valid"] = t_valid
```

---

## 6. Migration path

1. Add `_mask_to_padded_indices` and the emission block to `data_collator.py`.
2. The existing `_compute_sdpo_loss` in `composer_trainer.py` already handles the index path (lines 184–242); it only needs the sentinel-mask addition described in §3 Step 3.
3. Update `_compute_sdpo_loss` to prefer the aligned index path even when shapes match — remove the legacy shape-only fallback from the strict path entirely.
4. The non-strict path (`strict_sdpo_alignment=False`) can fall back to `torch.nonzero(sdpo_loss_mask == 1)` as a convenience for ad-hoc scripts, but the canonical production path is the explicit indices.