Add soft distillation training script and configuration.

[gagika]

Description

This PR introduces train_distill.py, a new training script that implements a "Post-Pruning Recovery" distillation workflow by integrating the Tunix distillation library with MaxText.

Context & Problem Solved:
We needed a way to recover model quality after structural pruning (e.g., reducing attention heads). This script enables a "Stage 2" distillation process where a smaller, trainable Student model mimics the output distribution of a larger, frozen Teacher model on general data.

Key Implementation Details:

• Dual Model Loading: Initializes two distinct MaxText models (Student and Teacher) in the same process using pyconfig with student_overrides and teacher_overrides to support different architectures.
• Tunix Integration: Wraps MaxText models in TunixMaxTextAdapter to expose a standard interface compatible with the Tunix DistillationTrainer.

Tests

I have tested this script end-to-end on a TPU VM.

1. Verification of Model Loading:
Verified that Student and Teacher models load with distinct configurations (e.g., different head counts):

--- Student Configuration ---
  Dimensions: 32 Layers, 4096 Emb Dim, 256 Head Dim
  Attention Heads: 16 Query, 16 KV

--- Teacher Configuration ---
  Dimensions: 32 Layers, 4096 Emb Dim, 128 Head Dim
  Attention Heads: 32 Query, 8 KV

python3 -m src.MaxText.distillation.train_distill src/MaxText/configs/distillation.yml   run_name=distillation_test   base_output_directory=${BASE_OUTPUT_DIRECTORY}   checkpoint_period=2000   hf_access_token=$HF_TOKEN   log_period=100  save_checkpoint_on_completion=True 

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 168 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/distillation/train_distill.py	0.00%	168 Missing ⚠️

📢 Thoughts on this report? Let us know!