Docs migration from XPK to CTK(2/5)

docs/guides/distillation.md

@@ -176,7 +176,7 @@ Because it's additive, β's *absolute* magnitude matters relative to the logit l
 
 ### Layer indices for feature loss
 
-`distill_layer_indices` selects which scanned-layer slices contribute to `feature_loss`. The XPK launcher's default is `[0,1,2,...,7]` — the first 8 layers, irrespective of model depth. Better defaults usually exist:
+`distill_layer_indices` selects which scanned-layer slices contribute to `feature_loss`. The default configuration is `[0,1,2,...,7]` — the first 8 layers, irrespective of model depth. Better defaults usually exist:
 
 | Goal                                        | Llama-8B (32 layers)                        | Llama-70B (80 layers)      |
 | ------------------------------------------- | ------------------------------------------- | -------------------------- |

docs/tutorials/posttraining/knowledge_distillation.md

@@ -41,8 +41,8 @@ The following recipe demonstrates the process of offline distillation using **Qw
 #### a. Setup environment variables
 
 ```bash
-export HF_TOKEN=<your-hf-token> # e.g., hf_BA6...
-export RUN_NAME=<your-run-name> # e.g., distill-20260115
+export HF_TOKEN=<HF_TOKEN> # e.g., hf_BA6...
+export RUN_NAME=<RUN_NAME> # e.g., distill-20260115
 ```
 
 #### b. Install dependencies
@@ -56,23 +56,23 @@ To store large models and datasets, attach a Hyperdisk to your TPU VM. Refer to
 First, create a Hyperdisk:
 
 ```bash
-export ZONE=<your-tpu-zone>  # e.g., us-central1-a
-export TPU_VM_NAME=<your-tpu-vm-name>
-export DISK_NAME=<your-disk-name>  # e.g., my-hyperdisk
-export DISK_SIZE=<disk-size>  # e.g., 500GB
+export ZONE=<ZONE>  # e.g., us-central1-a
+export TPU_VM_NAME=<TPU_VM_NAME>
+export DISK_NAME=<DISK_NAME>  # e.g., my-hyperdisk
+export DISK_SIZE=<DISK_SIZE>  # e.g., 500GB
 
-gcloud compute disks create ${DISK_NAME?} \
-  --size=${DISK_SIZE?} \
+gcloud compute disks create <DISK_NAME> \
+  --size=<DISK_SIZE> \
   --type=hyperdisk-balanced \
-  --zone=${ZONE?}
+  --zone=<ZONE>
 ```
 
 Then, attach the disk to your TPU VM:
 
 ```bash
-gcloud compute instances attach-disk ${TPU_VM_NAME?} \
-  --disk=${DISK_NAME?} \
-  --zone=${ZONE?}
+gcloud compute instances attach-disk <TPU_VM_NAME> \
+  --disk=<DISK_NAME> \
+  --zone=<ZONE>
 ```
 
 Inside the TPU VM, format and mount the disk (if not already mounted):
@@ -87,22 +87,22 @@ sudo mount /dev/sdb /mnt/hyperdisk
 Update the BASE_OUTPUT_DIRECTORY to point to the mounted disk and create the directory:
 
 ```bash
-export BASE_NAME=<your-base-directory>  # e.g., knowledge-distillation
-export BASE_OUTPUT_DIRECTORY=/mnt/hyperdisk/${BASE_NAME?}
-mkdir -p ${BASE_OUTPUT_DIRECTORY?}
+export BASE_NAME=<BASE_DIRECTORY>  # e.g., knowledge-distillation
+export BASE_OUTPUT_DIRECTORY=/mnt/hyperdisk/<BASE_DIRECTORY>
+mkdir -p <GCS_BUCKET>
 ```
 
 > **Note:** This tutorial uses a mounted Hyperdisk for performance and reproducibility, because writing large model files and many small I/O operations directly to `gs://` can be significantly slower.
 
 ### Obtain and prepare the teacher model
 
-For the teacher model, we will use **vLLM** to run inference. vLLM can load Hugging Face checkpoints directly, so **no conversion to MaxText format is needed** for the teacher. Ensure the teacher model is supported on TPU vLLM (refer to the [vLLM TPU recommended models](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features) for the latest list).
+For the teacher model, we will use **vLLM** to run inference. vLLM can load Hugging Face checkpoints directly, so **no conversion to MaxText format is needed** for the teacher. Ensure the teacher model is supported on TPU vLLM (refer to the [vLLM TPU recommended models](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models/) for the latest list).
 
 You can simply download the model from Hugging Face to your local directory:
 
 ```bash
-huggingface-cli login --token ${HF_TOKEN?}
-huggingface-cli download Qwen/Qwen3-32B --repo-type model --local-dir ${BASE_OUTPUT_DIRECTORY?}/qwen3-32b
+huggingface-cli login --token <HF_TOKEN>
+huggingface-cli download Qwen/Qwen3-32B --repo-type model --local-dir <GCS_BUCKET>/qwen3-32b
 ```
 
 ### Obtain and prepare the student model
@@ -121,13 +121,13 @@ python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
 ```bash
 # Set the checkpoint directory
-export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY?}/llama3.1-8b-ckpt
+export MAXTEXT_CKPT_PATH=<GCS_BUCKET>/llama3.1-8b-ckpt
 
 # Convert to MaxText format
 python3 -m maxtext.checkpoint_conversion.to_maxtext \
     model_name=llama3.1-8b \
-    hf_access_token=${HF_TOKEN?} \
-    base_output_directory=${MAXTEXT_CKPT_PATH?} \
+    hf_access_token=<HF_TOKEN> \
+    base_output_directory=<CKPT_PATH> \
     scan_layers=True skip_jax_distributed_system=True
 ```
 
@@ -138,18 +138,18 @@ Use the provided script `generate_distillation_data_vllm.py` to generate the dat
 Run the generation script:
 
 ```bash
-export OUTPUT_DATASET=${BASE_OUTPUT_DIRECTORY?}/datasets/distillation_data.parquet
+export OUTPUT_DATASET=<GCS_BUCKET>/datasets/distillation_data.parquet
 
 python3 -m tools.data_generation.generate_distillation_data_vllm \
   --dataset-path HuggingFaceH4/ultrachat_200k \
   --data-split train_sft \
   --data-columns messages \
-  --hf-access-token ${HF_TOKEN?} \
-  --teacher-model ${BASE_OUTPUT_DIRECTORY?}/qwen3-32b \
+  --hf-access-token <HF_TOKEN> \
+  --teacher-model <GCS_BUCKET>/qwen3-32b \
   --use-chat-template \
   --num-prompts 5120 \
   --num-generations 2 \
-  --output-file ${OUTPUT_DATASET?}
+  --output-file <DATASET_PATH>
 
 ```
 
@@ -163,8 +163,8 @@ The checkpoint from the student model's fine-tuning (on the teacher-generated da
 
 ```bash
 # Get the latest checkpoint for fine-tuned student model
-CHECKPOINTS_PATH=${BASE_OUTPUT_DIRECTORY?}/distillation/qwen3-32b-distill-llama3.1-8b/${RUN_NAME?}/checkpoints
-checkpoints=$(ls ${CHECKPOINTS_PATH?})
+CHECKPOINTS_PATH=<GCS_BUCKET>/distillation/qwen3-32b-distill-llama3.1-8b/<RUN_NAME>/checkpoints
+checkpoints=$(ls <CKPT_PATH>)
 integer_dirs=()
 for dir in $checkpoints; do
   dir_name=$(basename "$dir")
@@ -174,24 +174,24 @@ for dir in $checkpoints; do
 done
 sorted_dirs=($(printf '%s\n' "${integer_dirs[@]}" | sort -n))
 largest_dir="${sorted_dirs[-1]}"
-FINE_TUNED_MODEL_CKPT_PATH=${CHECKPOINTS_PATH?}/${largest_dir}/model_params
+FINE_TUNED_MODEL_CKPT_PATH=<CKPT_PATH>/${largest_dir}/model_params
 
 # Fine-tune student model on original dataset
 python3 -m maxtext.trainers.post_train.sft.train_sft \
-  run_name=${RUN_NAME?}_stage2 \
-  base_output_directory=${BASE_OUTPUT_DIRECTORY?}/distillation/qwen3-32b-distill-llama3.1-8b \
+  run_name=<RUN_NAME>_stage2 \
+  base_output_directory=<GCS_BUCKET>/distillation/qwen3-32b-distill-llama3.1-8b \
   tokenizer_path=meta-llama/Llama-3.1-8B-Instruct tokenizer_type=huggingface \
   dataset_type=hf \
   hf_path='HuggingFaceH4/ultrachat_200k' \
   train_split='train_sft' \
   train_data_columns=['messages'] \
-  load_parameters_path=${FINE_TUNED_MODEL_CKPT_PATH?} \
+  load_parameters_path=<CKPT_PATH> \
   model_name=llama3.1-8b \
   per_device_batch_size=2 \
   steps=200 \
   ici_expert_parallelism=-1 ici_fsdp_parallelism=4 \
   max_target_length=2048 \
-  hf_access_token=${HF_TOKEN?} \
+  hf_access_token=<HF_TOKEN> \
   profiler=xplane
 ```
 
@@ -215,29 +215,29 @@ Online distillation runs the teacher inside MaxText (not vLLM), so both checkpoi
 # Student
 python3 -m maxtext.checkpoint_conversion.to_maxtext \
     model_name=llama3.1-8b \
-    hf_access_token=${HF_TOKEN?} \
-    base_output_directory=${BASE_OUTPUT_DIRECTORY?}/llama3.1-8b-ckpt \
+    hf_access_token=<HF_TOKEN> \
+    base_output_directory=<GCS_BUCKET>/llama3.1-8b-ckpt \
     scan_layers=True skip_jax_distributed_system=True
 
 # Teacher (example: same family, larger)
 python3 -m maxtext.checkpoint_conversion.to_maxtext \
     model_name=llama3.1-70b \
-    hf_access_token=${HF_TOKEN?} \
-    base_output_directory=${BASE_OUTPUT_DIRECTORY?}/llama3.1-70b-ckpt \
+    hf_access_token=<HF_TOKEN> \
+    base_output_directory=<GCS_BUCKET>/llama3.1-70b-ckpt \
     scan_layers=True skip_jax_distributed_system=True
 ```
 
 > **Note:** Student and teacher must share the same vocabulary. The trainer asserts `student_config.vocab_size == teacher_config.vocab_size` at startup.
 
 #### b. Install Tunix
 
-The online distillation trainer depends on Tunix. The XPK launcher script ([`scripts/run_distill_xpk.sh`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh)) contains a `prep_image` step that layers Tunix on top of the MaxText base image. For local runs, install the same pin used by the launcher — the default `TUNIX_SOURCE` in `run_distill_xpk.sh` is the source of truth. As of this writing:
+The online distillation trainer depends on Tunix. For local runs or custom images, install Tunix from GitHub:
 
 ```bash
 pip install "git+https://github.com/google/tunix@348959d18a4a09c75e58a7d49aec9d8b0eb4a8b6"
 ```
 
-> **Note:** The commit pin above will drift as the launcher is updated. Before installing, check the `TUNIX_SOURCE` default in [`run_distill_xpk.sh`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh) and use that spec. Once a Tunix PyPI release ships, this will become a versioned `google-tunix==<ver>` install.
+> **Note:** Once a Tunix PyPI release ships, this will become a versioned `google-tunix==<ver>` install.
 
 ### Configuration
 
@@ -294,15 +294,15 @@ The example below demonstrates **Pattern B** (pruning recovery): the student is
 ```bash
 python3 -m maxtext.trainers.post_train.distillation.train_distill \
   src/maxtext/configs/post_train/distillation.yml \
-  run_name=${RUN_NAME?} \
-  base_output_directory=${BASE_OUTPUT_DIRECTORY?}/distillation/online \
+  run_name=<RUN_NAME> \
+  base_output_directory=<GCS_BUCKET>/distillation/online \
   tokenizer_path=meta-llama/Llama-3.1-8B tokenizer_type=huggingface \
-  hf_access_token=${HF_TOKEN?} \
+  hf_access_token=<HF_TOKEN> \
   student_overrides.model_name=llama3.1-8b \
   student_overrides.base_num_decoder_layers=24 \
-  student_overrides.load_parameters_path=${BASE_OUTPUT_DIRECTORY?}/pruned-llama3.1-8b-24L/0/items \
+  student_overrides.load_parameters_path=<GCS_BUCKET>/pruned-llama3.1-8b-24L/0/items \
   teacher_overrides.model_name=llama3.1-8b \
-  teacher_overrides.load_parameters_path=${BASE_OUTPUT_DIRECTORY?}/llama3.1-8b-ckpt/0/items \
+  teacher_overrides.load_parameters_path=<GCS_BUCKET>/llama3.1-8b-ckpt/0/items \
   per_device_batch_size=2 \
   gradient_accumulation_steps=8 \
   ici_fsdp_parallelism=4 \
@@ -320,49 +320,74 @@ The schedule values above are a strong default for same-size pruning recovery. S
 
 > **Note:** `distill_layer_indices` is applied to **both** student and teacher activations identically. When the two have different depths (Pattern A or a depth-pruned Pattern B), every index must be valid on the *smaller* side, and same-numbered layers are aligned across the two models. The trainer cannot map student layer *i* to teacher layer *f(i)* for arbitrary *f*. If the depths differ significantly, prefer logit-only distillation (`distill_beta=0`).
 
-#### Multi-host on GKE via XPK
+#### Cluster Toolkit multi-host submission
 
-A reference launcher is provided at `src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh`. It handles image preparation (`prep_image` layers Tunix on top of the MaxText base image), workload submission, log streaming, and an auto-resume loop for long-running jobs.
-
-Minimum environment variables:
+Submit the distillation trainer directly as a Cluster Toolkit JobSet:
 
 ```bash
-export XPK_CLUSTER=<your-gke-cluster>
-export XPK_PROJECT=<your-gcp-project>
-export XPK_ZONE=<cluster-zone>             # e.g. us-central1-a
-export XPK_DEVICE_TYPE=<tpu-type>          # e.g. tpu7x-4x4x4, v5p-128
-export XPK_BASE_OUTPUT_DIR=gs://<bucket>/distill-runs
-
-# Distillation hyperparameters (always passed; override yml values)
-export DISTILL_ALPHA=0.9
-export DISTILL_TEMPERATURE=2.0
-export DISTILL_BETA=1.0
-# Layer indices for feature loss. Every index must be valid on the smaller side
-# (student for Pattern A, both for Pattern B). Values below assume a 32-layer
-# student; adjust for other depths — see the Distillation guide's layer-index table.
-export DISTILL_LAYER_INDICES=[3,7,11,15,19,23,27,31]   # no spaces inside brackets
+export PROJECT_ID=<PROJECT_ID>
+export GKE_CLUSTER=<CLUSTER_NAME>
+export LOCATION=<ZONE> # e.g., 'europe-west4' (region) or 'us-central1-a' (zone)
+export RUN_NAME=<RUN_NAME>
+export IMAGE_URI=<IMAGE_NAME>
+export COMPUTE_TYPE=<COMPUTE_TYPE>
+export TOPOLOGY=<TOPOLOGY>
+export BASE_OUTPUT_DIRECTORY=gs://<GCS_BUCKET>/distillation
+export STUDENT_CKPT_PATH=gs://<GCS_BUCKET>/<STUDENT_MODEL_PATH>/checkpoints/0/items
+export TEACHER_CKPT_PATH=gs://<GCS_BUCKET>/<TEACHER_MODEL_PATH>/checkpoints/0/items
+export TOKENIZER_PATH=meta-llama/Llama-3.1-8B
+export HF_TOKEN=<HF_TOKEN>
+
+gcloud config set project <PROJECT_ID>
+gcloud container clusters get-credentials <CLUSTER_NAME> \
+  --location <ZONE> \
+  --project <PROJECT_ID>
+gcluster job config set project <PROJECT_ID>
+gcluster job config set cluster <CLUSTER_NAME>
+gcluster job config set location <ZONE>
+
+gcluster job submit \
+  --image=gcr.io/<PROJECT_ID>/<IMAGE_NAME> \
+  --name=<RUN_NAME> \
+  --compute-type=<COMPUTE_TYPE> \
+  --topology=<TOPOLOGY> \
+  --command="python3 -m maxtext.trainers.post_train.distillation.train_distill \
+    src/maxtext/configs/post_train/distillation.yml \
+    run_name=<RUN_NAME> \
+    base_output_directory=<GCS_BUCKET>/online \
+    tokenizer_path=<TOKENIZER_PATH> \
+    tokenizer_type=huggingface \
+    hf_access_token=<HF_TOKEN> \
+    student_overrides.model_name=llama3.1-8b \
+    student_overrides.base_num_decoder_layers=24 \
+    student_overrides.load_parameters_path=<STUDENT_MODEL_PATH> \
+    teacher_overrides.model_name=llama3.1-8b \
+    teacher_overrides.load_parameters_path=<TEACHER_MODEL_PATH> \
+    per_device_batch_size=2 \
+    distill_alpha=0.9 \
+    distill_temperature=2.0 \
+    distill_beta=1.0 \
+    distill_layer_indices=[2,5,8,11,14,17,20,23]"
 ```
 
-Then:
-
-```bash
-# One-time: layer Tunix on top of the MaxText base image
-bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh prep_image
+#### Monitor and clean up
 
-# Bake ./src into a runner image and push to gcr.io/$XPK_PROJECT/...:${USER}-distill
-bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh upload_runner
+Monitor the workload and stream logs with Cluster Toolkit:
 
-# Submit a workload
-bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh submit
+```bash
+# Check job status
+gcluster job list
 
 # Stream logs
-bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh monitor
+gcluster job logs <RUN_NAME>
 
-# Auto-resume on failure (uses the same workload + base output dir, so checkpoint resume works)
-bash src/maxtext/trainers/post_train/distillation/scripts/run_distill_xpk.sh resume_until_done
-```
+# Inspect JobSet and pods
+kubectl get jobset -l gcluster.google.com/workload=<RUN_NAME>
+kubectl get pods -l gcluster.google.com/workload=<RUN_NAME>
 
-The script's header comment lists every supported environment variable.
+# Cancel workload
+gcluster job cancel <RUN_NAME>
+```
 
 ### Offline top-k logits variant