[Distillation] Fix eod masking + strategy refactoring#3478

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

vladk/distill-eod-refactorAI-Hypercomputer/maxtext:vladk/distill-eod-refactorCopy head branch name to clipboard