Soft distillation training support

[gagika]

Description

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

Context & Problem Solved:
This script recovers model quality after structural pruning (e.g., reducing attention heads) by enabling a 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.

Run distillation

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 

Tests

I have tested this script end-to-end on a TPU VM.
Verified that Student and Teacher models load with distinct configurations

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 180 lines in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!

[richjames0]

lgtm

[NicoGrande]

LGTM!

[NicoGrande]

Can we remove these commented lines?

[gagika]

done

[NicoGrande]

Is my understanding correct that both models will be collocated on the same Mesh? Will this need to be extended to disaggregated to work for different size student / teacher models for example?

[gagika]

yes, it's collocated setup on the same mesh, we can still run different model sizes on the same mesh with possibly different sharding.

disaggregated setup not yet planned / supported, we can extend in the future if needed.