Refactor model loading in distillation to remove Tunix adapter.

[gagika]

Description

Refactors the distillation trainer to remove the TunixMaxTextAdapter dependency, enabling direct control over model execution and configuration.

Key Changes:

• Direct Model Loading: Replaces the Tunix adapter with standard MaxText nnx.Module loading via get_maxtext_model.
• Config-Bound Forward Pass: Introduces create_forward_fn to generate distinct forward functions for Student and Teacher. This correctly binds enable_dropout from the config, allowing the Student to train with dropout while keeping the Teacher deterministic.
• Strict Configuration: Removed environment variable fallbacks; teacher_overrides.load_parameters_path is now mandatory to ensure reproducibility.

Tests

Verified locally with the following command:

python3 -m src.MaxText.distillation.train_distill src/MaxText/configs/distillation.yml \
  run_name=distill_llama3 \
  base_output_directory=${BASE_OUTPUT_DIRECTORY} \
  checkpoint_period=2000 \
  hf_access_token=$HF_TOKEN \
  log_period=10 \
  save_checkpoint_on_completion=True \
  teacher_overrides.load_parameters_path="gs://agagik-us/llama3/llama3.1-8b/scanned_chkpt/0/items"

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

Files with missing lines	Patch %	Lines
src/MaxText/distillation/train_distill.py	0.00%	12 Missing ⚠️

📢 Thoughts on this report? Let us know!

[richjames0]

lgtm

[dipannita08]

Out of curiosity, why do we remove this fallback?

[gagika]

I was testing with env TEACHER_CHECKPOINT_PATH before for teacher checkpoint path, but there is not need to have extra env, it can just be provided from distillaiton.yml file or command line argument, e.g. teacher_overrides.load_parameters_path="gs://agagik-us/...