use repeat_interleave for input_features in kd_trainer

[zhuxiaoxuhit]

Same issue as #119 but in KnowledgeDistillationTrainer._prepare_logprob_inputs.

input_features uses .repeat(num_generations, 1, 1) while all other tensors in the same function use .repeat_interleave(num_generations, dim=0). These two produce different element orderings when batch size > 1:

• repeat_interleave: [s0, s0, s0, s1, s1, s1, ...] which matches how generated_ids are laid out
• repeat: [s0, s1, s0, s1, s0, s1, ...] which does not

So when batch size > 1, input_features ends up paired with the wrong completions during log probability computation, making both the student and teacher logprob calculations incorrect and corrupting the KD training signal.

[yuekaizhang]

@zhuxiaoxuhit Thanks.

[yuekaizhang]

@robin1001 Would you mind helping merge this one also? Thanks.