Implementation for soft offline distillation using saved top-k teacher logits

[ajkv-google]

Description

This PR introduces an end-to-end offline distillation training pipeline. Previously, the distillation loop executed in an "online" mode, which required both the frozen Teacher model and the learning Student model to be loaded and executed simultaneously during training. This change allows the trainer to load pre-computed, top-K Teacher logits from .array_record files, which allows us to bybass the forward pass for the teacher model during the training loop.

Tests

Tested this code change by running the following command:

python3 src/maxtext/trainers/post_train/distillation/train_distill.py src/maxtext/configs/post_train/distillation.yml steps=100 tokenizer_path="/mnt/ajkv/disks/codebase/maxtext/src/maxtext/assets/tokenizers/tokenizer_llama3.tiktoken" --offline_distillation --offline_data_dir="/mnt/ajkv/disks/teacher_logits_output/teacher_top_k_global.array_record"

Truncated output showing the successful run: https://paste.googleplex.com/6342987127848960#l=8.

Verified that the training happened sucessfully and finished the distillation run.

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

Files with missing lines	Patch %	Lines
.../trainers/post_train/distillation/train_distill.py	13.88%	31 Missing ⚠️
...ners/post_train/distillation/distillation_utils.py	28.57%	25 Missing ⚠️

📢 Thoughts on this report? Let us know!