I wonder it is possible to pre-extract the teacher logics (say, on a TPU) and just load them in the distiller code.

I think this can be achieved with datasets library, we can try to cache the logits along with the examles and while loading example , load its’ corresponding logits as well, so the dataset  could return dict which could look something like this
{"input_ids": [], "decoder_input_ids": [], "precomputed_logits": []}