Workflow to save top-k teacher logits in GCS to use in distillaiton#3193
Merged
copybara-service[bot] merged 2 commits intomainfrom Mar 19, 2026
Merged
Workflow to save top-k teacher logits in GCS to use in distillaiton#3193copybara-service[bot] merged 2 commits intomainfrom
copybara-service[bot] merged 2 commits intomainfrom
Conversation
ajkv-google
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jesselu-google,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
gagika
reviewed
Feb 19, 2026
vlad-karp
reviewed
Feb 19, 2026
ajkv-google
force-pushed
the
ajkv-teacher-top-k-distillation
branch
2 times, most recently
from
044cbf7 to
7cffe43
Compare
gagika
reviewed
Feb 24, 2026
ajkv-google
requested review from
Obliviour,
SujeethJinesh,
jacoguzo,
mitalisi,
notabee and
shauryagup
as code owners
ajkv-google
requested review from
Lumosis,
gpolovets1,
jrplatin,
mailvijayasingh,
parambole,
patemotter and
shuningjin
as code owners
ajkv-google
force-pushed
the
ajkv-teacher-top-k-distillation
branch
2 times, most recently
from
4de8440 to
7cddd84
Compare
entrpn
approved these changes
Mar 3, 2026
4 tasks
vlad-karp
approved these changes
Mar 16, 2026
Moved optional keys to be cmd arguments Added script to verify writing of data to gs bucket added code to work in multihost fixed spacing and code formatting saving top-k teacher logits on one host to local with option to store to a gs bucket updated code formatting Updated to provide local filepath as cmd arg
ajkv-google
force-pushed
the
ajkv-teacher-top-k-distillation
branch
from
6a79f42 to
c5702d8
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces a standalone script (
save_top_k_teacher_logits.py) to generate, extract, and save top-K logits and their corresponding vocabulary indices from a teacher model directly to a GCS bucket for offline distillation. In the current form of distillation, both the teacher and student models are loaded together into memory, and the process of distillation is done in an "online" fashion. This can lead to slower training times due to computing the teacher's probability distribution on the fly and training the student model simultaneously. In this change, we are saving the top-k teacher outputs in a GCS bucket (instead of the full probability distribution), so that this can be used to train the student model more efficiently without having to use extra memory to compute teacher logits on the fly.Tests
Distillation YAML used to run this script: https://paste.googleplex.com/5736358690816000
Ran the following command from the head of the Maxtext directory:
python3 src/maxtext/trainers/post_train/distillation/save_top_k_teacher_logits.py src/maxtext/configs/post_train/distillation.ymlTerminal output for the above command (truncated due to size): https://paste.googleplex.com/4609623547052032
Verified that the ArrayRecord file has been written to a GS bucket in the correct path.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.