feat(distill): teacher logits provider abstraction (SPEC-DISTILL-001 Phase 1, PMAT-691)

[noahgift]

Summary

Implements SPEC-DISTILL-001 Phase 1 — replaces the synthetic-logits stub in aprender-train-distill::pipeline.rs::train() with a real TeacherLogitsProvider trait + a FixtureTeacher implementation suitable for unit tests.

Phase 1b (PMAT-692, follow-up PR) will add a RealizarTeacher backend that runs the real 7B teacher via aprender-serve.

What this PR adds

New module — aprender-train-distill::teacher_provider:

pub trait TeacherLogitsProvider {
    fn vocab_size(&self) -> usize;
    fn logits_for_batch(&mut self, input_ids: &[Vec<u32>])
        -> Result<Vec<Vec<f32>>>;
}

pub struct FixtureTeacher { vocab_size: usize }

The FixtureTeacher returns deterministic stub logits whose argmax is the last input token — useful for unit-testing the KD signal flow without needing a real model file. Out-of-vocab tokens collapse to position 0 (robustness).

Pipeline integration — pipeline.rs:

• Pipeline gains a teacher: Box<dyn TeacherLogitsProvider + Send> field
• Pipeline::new() defaults to FixtureTeacher::new(32) so legacy tests behave identically
• Pipeline::with_teacher() lets Phase 4 callers swap in RealizarTeacher
• train() pulls teacher_logits from self.teacher.logits_for_batch() and reshapes to Array2<f32> shape [batch, vocab]
• Legacy build_synthetic_logits() retained only for the student path (Phase 2 replaces it with CudaTransformerTrainer.forward_backward_kd)

Why "online" instead of cached (spec revision)

The original SPEC-DISTILL-001 v1.0.0 plan cached top-K=64 logits to disk. The storage math (1.24B tokens × 64 × 6 bytes ≈ 476 GB) exceeds the lambda-vector NVMe budget. Online teacher inference (the DistilBERT / Distil-Qwen actual practice) trades ~2× student-step time for zero cache footprint. Spec was revised to v1.1.0 in #1785 with this design.

Tests

6 new teacher_provider::tests::* unit tests pass:

• vocab_size reporting
• one logits vec per batch element
• argmax recovers last input token (F-DISTILL-FIXTURE-001)
• out-of-vocab token collapses to position 0 (robustness)
• empty batch returns empty vec
• deterministic across calls and instances

All 41 aprender-train-distill lib tests pass.

What's next

• Phase 1b — PMAT-692: implement RealizarTeacher (the real backend that delegates to aprender-serve's Model::forward_logits path). The trait is now in place; Phase 1b is just adding a second impl.
• Phase 2 — PMAT-693: wire CudaTransformerTrainer.forward_backward_kd_batch into the student path (replaces the remaining build_synthetic_logits call site for the student).
• Phase 3-6: per SPEC-DISTILL-001.

Test plan

• 6 new teacher_provider tests pass
• All 41 aprender-train-distill lib tests pass
• cargo check -p aprender-train-distill clean
• Phase 1b — wire RealizarTeacher to the real teacher inference path

🤖 Generated with Claude Code