Update offline distillation and saving top-k teacher logits to be efficient

[ajkv-google]

Description

This PR implements an optimized Offline Distillation pipeline in MaxText, allowing student models to train using pre-saved teacher logits to significantly reduce compute costs.

Key changes:

• The save_top_k_teacher_logits.py script now supports parallel execution across hosts and uses asynchronous GCS uploads to prevent TPU idling
• Offline distillation now uses grain input pipeline without having to use the custom offline arrayrecord iterator. This speeds up the offline distillation process without having to read logits on a single thread.
• Introduced sparse KL divergence in distillation_utils.py that computes KL divergence using only the teacher's top-k predictions when running offline distillation

Tests

Used the following to run offline distillation:

• Command: https://paste.googleplex.com/4852897526448128
• Yaml: https://paste.googleplex.com/4981440528908288
• Example run for offline distillation (using our team's LTI checkpoint for 30b model): Tensorboard Link

Used the following to run saving of top-k teacher logits:

• Command: https://paste.googleplex.com/6383382496935936
• Yaml: https://paste.googleplex.com/5898105013796864
• GS Bucket Path for saved logits (1B tokens of climbmix): Saved offline logits

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 40.81633% with 29 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/input_pipeline_utils.py	25.00%	7 Missing and 5 partials ⚠️
...ners/post_train/distillation/distillation_utils.py	56.25%	6 Missing and 1 partial ⚠️
.../trainers/post_train/distillation/train_distill.py	0.00%	6 Missing ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	63.63%	2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[vlad-karp]

should we also compute teacher logits softmax over entire vocabulary before saving to files?

[ajkv-google]

We save the raw logits instead of the full softmax so we can still tweak the temperature on the fly during training. Also, saving full-vocab probabilities would make the files bigger, which could reduce read speeds during offline training. I think as of now, it would be good to stick to the current approach where we store the information that is needed, and read it quickly during the offline training.

[vlad-karp]

"lso, saving full-vocab probabilities would make the files bigger" - why do you consider saving the full vocab? the entire idea of offline is to use a limited set of logits.
My concern is purely mathematical - you normalize only over top-k while the student logits will be normalized over entire vocabulary, and then you calculate kl divergence over those distributions with completely different normalization scales.