research/11-sdpo-alignment-indices.md

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:

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:

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

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:

# --- 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:

# 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)

# === 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.