[REQUEST] Add an OPSD training batch builder for teacher-prefix assembly and response-logit alignment

Start with deepspeed/runtime/rollout/opsd.py and the existing ResponseTokenBatch and generalized_jsd_loss behavior, then review deepspeed/runtime/rollout/__init__.py. Add hand-derived alignment and validation coverage in tests/unit/runtime/rollout/test_rollout_interface.py, plus the existing training integration location. Done means variable-padding causal logits align to identical response targets, invalid inputs fail explicitly, gradients reach the student, the teacher remains unchanged, and docs/code-docs/source/inference-engine.rst documents the public API.

Summary

Add an OPSDTrainingBatch builder that consumes a RolloutBatch and a TeacherContext, constructs the teacher sequence by appending the student's sampled completion to the teacher prefix, and provides explicit causal-logit alignment for response-only distillation.

The student and teacher may have different prefix lengths, but must score the same sampled completion tokens in the same order.

Motivation

The existing OPSD utilities provide ResponseTokenBatch and generalized_jsd_loss, but callers still have to construct teacher inputs and align teacher logits manually.

Simply slicing both model outputs at the same sequence positions is incorrect when teacher and student prefixes differ. It can compare predictions for different tokens or introduce a one-token causal shift error. A shared builder should make this contract explicit and remove duplicated trainer-side indexing.

Input contract

RolloutBatch

• Use the existing input_ids, attention_mask, and per-row response_start_idx.
• Extract completion tokens only from attended positions at or after the row's response boundary.
• Preserve token IDs and token order exactly; do not decode and re-tokenize the completion.
• Retain attended EOS tokens. Determine padding from the attention mask, not from equality with pad_token_id, since PAD and EOS may share an ID.

TeacherContext

• Provide tokenized teacher-prefix IDs and their attention mask, with one row per rollout row.
• Each prefix is the complete teacher-side context preceding the first completion token, including any required prompt, privileged context, chat-template markers, or special tokens.
• Require at least one attended prefix token for every row with a non-empty completion, so the first completion token has a causal prediction position.
• For n_samples_per_prompt > 1, callers must expand/reorder context rows to match rollout order before invoking the builder. Do not infer sample mapping from batch-size ratios or silently broadcast context.

Shared assumptions

• Student and teacher use the same token-ID semantics and vocabulary ordering. Matching vocabulary sizes alone do not establish this compatibility.
• The initial implementation targets tokenized inputs on one device. Validate shapes, integer index/token dtypes, masks, response boundaries, batch sizes, and device compatibility explicitly.
• Padding IDs must be supplied explicitly or obtained through an existing documented context/config contract, not guessed from token values.

Proposed behavior

1. Preserve the student's rollout input sequence and its attention-mask semantics.
2. Build each teacher sequence as valid_teacher_prefix + valid_student_completion, then pad teacher rows into a batch with a documented layout.
3. Provide model inputs and completion-position metadata for both sides, including padding-aware position IDs where required by the supported model interface.
4. Align externally computed student and teacher logits to completion-token order, using each side's own causal prediction positions.
5. Return aligned student logits, aligned teacher logits, shared completion targets, and a boolean response mask that can be passed to generalized_jsd_loss.

The builder performs tensor assembly and indexing only. It must not invoke either model, generate a teacher completion, calculate the distillation loss, or update model parameters.

Suggested API and output contract

batch = build_opsd_training_batch(
    rollout=rollout_batch,
    teacher_context=teacher_context,
    pad_token_id=tokenizer.pad_token_id,
)

student_output = student_model(**batch.student_model_inputs)
with torch.no_grad():
    teacher_output = teacher_model(**batch.teacher_model_inputs)

student_logits, teacher_logits = batch.align_logits(
    student_output.logits,
    teacher_output.logits,
)
loss_output = generalized_jsd_loss(
    student_logits,
    teacher_logits,
    batch.response_mask,
)

Names above are illustrative. Prefer a small public API and reuse existing ResponseTokenBatch logic where applicable.

• student_model_inputs and teacher_model_inputs contain the model inputs, not loss-only metadata.
• Aligned logits have shape [B, R_max, V]; completion targets and response_mask have shape [B, R_max].
• R_max is the maximum number of valid completion tokens in the batch, independent of either padded prefix width.
• At every valid (row, response_index), both logits predict exactly the token stored in the shared completion targets.
• Masked slots must use safe gather indices and finite values so that padding cannot introduce NaNs into the existing loss.
• If every completion is empty, expose zero valid response tokens and support the existing loss's zero-token behavior without invalid indexing.
• Alignment must preserve the student autograd path. The example caller runs the teacher under no_grad; the builder does not manage model execution or training policy.

Causal alignment example

For one unpadded row:

student input: [S0, S1, C0, C1]
teacher input: [T0, T1, T2, T3, C0, C1]

shared targets:          [C0, C1]
student logit positions: [ 1,  2]
teacher logit positions: [ 3,  4]
response mask:          [True, True]

Logits at position i predict the token at position i + 1. In particular, the first completion token is predicted from the last prefix position, not from its own position. Padding changes physical indices, but must not change which target each selected logit predicts.

Validation and error handling

• Reject malformed shapes, non-integer token/index tensors, invalid attention masks, out-of-range response boundaries, and mismatched devices or row counts with actionable errors.
• Document supported padding layouts. Reject unsupported mask layouts rather than silently changing completion ordering or causal context.
• Reject an empty teacher prefix for a non-empty completion; do not silently insert BOS or other special tokens.
• Validate each model output against its own input batch/sequence dimensions before gathering logits, and reject incompatible vocabulary dimensions.
• Document shared tokenizer/vocabulary semantics as a caller precondition; do not claim that tensor shape validation proves semantic compatibility.
• Do not silently truncate prefixes or completions. Context-window budgeting stays with the caller/model; the builder must not hide an overlength sequence by dropping tokens.
• Do not mutate the input RolloutBatch or TeacherContext.

Acceptance criteria

• Different teacher/student prefix lengths produce the exact same completion targets in the aligned output, with the correct causal shift.
• Focused contract tests use hand-derived examples like the one above to catch first-token off-by-one errors and incorrect cross-row/sample mapping.
• Tests cover variable prefix/completion lengths, supported left/right padding, attended EOS with pad_token_id == eos_token_id, mixed empty/non-empty completions, and an all-empty completion batch.
• Invalid boundaries, malformed context, mismatched batches/devices, and incompatible model output shapes fail explicitly.
• A minimal actual-device training integration constructs a rollout, builds both inputs, runs student/teacher forwards, calls the existing JSD loss, and performs a student update. Compare aligned outputs/loss with independently constructed per-example unpadded forwards; verify student gradients flow and the teacher remains unchanged under no_grad.
• The integration is executed on actual hardware and the PR reports the device specification and results; mocked logits alone are insufficient for the training-level contract.
• Existing ResponseTokenBatch and generalized_jsd_loss behavior remains compatible.
• Public exports and documentation include an end-to-end usage example explaining prefix ownership, row mapping, vocabulary assumptions, causal alignment, and empty-completion behavior.

Scope boundaries

This issue does not add a teacher model loader, prompt-template system, tokenizer conversion, cross-vocabulary distillation, teacher response generation, a new loss, rollout scheduling changes, or a complete OPSD trainer.

Suggested implementation locations

• deepspeed/runtime/rollout/opsd.py: batch/context contracts, builder, and alignment.
• deepspeed/runtime/rollout/__init__.py: intended public exports.
• tests/unit/runtime/rollout/test_rollout_interface.py: consolidate focused batching/alignment contract tests.
• Reuse the existing appropriate training integration test location rather than creating multiple narrowly scoped test files.
• docs/code-docs/source/inference-engine.rst: extend the OPSD utilities section, replacing the manual aligned_teacher_logits placeholder with a concrete builder example.

CC @delock @PKUWZP