Blockwise parallel distillation in the scan mode

[vlad-karp]

Description

Blockwise parallel distillation implementation in the scan mode.

The rest of the description includes relevant details and context, examples:

• This is experimental disillaion mode which might be useful in general and for layer-wise distillaions later
• Quality of distillation might be improved if this mode is used in some combination for the standard distillation.
• It uses the same approach as intermediate per-layer output cosine/l2 loss and is implemented for the efficient scan mode.
• Is enabled by blockwise_distill configuration flag defined in the teacher model config. It then saves the teacher model attention output states and passes them to the student. The student model detect them and propagates to its own attention modules and replaces own attention inputs. Then finally if blockwise_distill is enabled, the loss is coputed only as the feature loss thus counting only l2/cosine distances between teacher vs student attention outputs
• works only in the scan mode now

Tests

Added a new test case in tests/post_training/unit/train_distill_test.py as test_train_step_passes_injected_attention_inputs_under_blockwise_distill()

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

Files with missing lines	Patch %	Lines
.../trainers/post_train/distillation/train_distill.py	0.00%	7 Missing ⚠️
src/maxtext/utils/maxtext_utils.py	0.00%	5 Missing ⚠️
src/maxtext/layers/decoders.py	40.00%	2 Missing and 1 partial ⚠️
src/maxtext/layers/attentions.py	0.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!