feat(distill): CudaStudentProvider + forward_backward_with_grad (SPEC-DISTILL-001 Phase 2d)

[noahgift]

Summary

Closes the last engineering gap before Phase 3 (500-step E2E smoke). Real GPU student backend wraps CudaTransformerTrainer and implements the Phase 2b StudentLogitsProvider trait.

Stacks on top of #1792 (Phase 2c pipeline integration).

Two complementary pieces

1. New pub method on CudaTransformerTrainer (cuda_trainer.rs:1247-1311):

pub fn forward_backward_with_grad(
    &mut self,
    input_ids: &[u32],
    logit_gradient: &[f32],
) -> Option<()>

Runs forward, uploads the caller-supplied logit gradient into the last-position slice of logits_buf (replacing what gpu_forward wrote), runs gpu_backward (which back-props from the in-place gradient through the transformer stack), and runs embed_backward. Matches the KAIZEN-052 in-place gradient convention that fused_cross_entropy_cuda uses for the CE path — except the gradient now comes from kd_step::kd_logit_gradient (Phase 2a).

2. New CudaStudentProvider in aprender-train-distill (cuda-gated):

pub struct CudaStudentProvider {
    trainer: CudaTransformerTrainer,
    vocab_size: usize,
    last_input_ids: Option<Vec<u32>>,  // batch_size=1 cache
}

impl StudentLogitsProvider for CudaStudentProvider {
    // logits_for_batch → trainer.forward_logits per batch element
    // apply_kd_gradient → trainer.forward_backward_with_grad on cached
    //                      last input_ids with the last gradient
}

Phase 2d limitation

batch_size=1 only — the trait's apply_kd_gradient doesn't take input_ids, so the provider caches from the most-recent logits_for_batch call. With batches >1, only the last element gets a real gradient update.

Phase 2e (PMAT-698, follow-up) generalizes via a fused-step trait method that takes input_ids + gradient together so all batch elements process correctly.

Tests

All 58 aprender-train-distill lib tests pass under both --features (none) and --features cuda. cargo check clean under both gates.

The CudaStudentProvider itself doesn't have a unit test — exercising it needs CUDA at test time, and the F-DISTILL-CUDA-STUDENT-001 falsifier (logits parity within 1e-6 vs a standalone forward_logits call) is integration-tested in Phase 4 production runs.

What's next

With Phase 2d landed, Phase 3 (500-step E2E smoke run with real CudaTrainerTeacher + CudaStudentProvider) is unblocked. Phase 4 is the actual 50K-step distillation training run for albor-370m-v2.

Test plan

• cargo check -p aprender-train-distill clean (no cuda)
• cargo check -p aprender-train-distill --features cuda clean
• cargo check -p aprender-train --features cuda clean (new method)
• All 58 distill lib tests pass under both feature configs
• Phase 3: 500-step E2E smoke on lambda-vector (real CUDA)
• Phase 2e: fused-step trait method for batch_size>1 (PMAT-698)

🤖 Generated with Claude Code

[noahgift]

Subsumed by #1797 squash-merge (chain-PR leapfrog pattern per memory rule). All content landed on main at aee8716.