Docs migration from XPK to CTK(2/5) #5176
+212
−128
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -41,8 +41,8 @@ The following recipe demonstrates the process of offline distillation using **Qw | |
| #### a. Setup environment variables | ||
|
|
||
| ```bash | ||
| 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=<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> \ | ||
| --type=hyperdisk-balanced \ | ||
| --zone=<ZONE> | ||
| ``` | ||
|
|
||
| Then, attach the disk to your TPU VM: | ||
|
|
||
| ```bash | ||
| 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=<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/) 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 <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=<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=<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=<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 <GCS_BUCKET>/qwen3-32b \ | ||
| --use-chat-template \ | ||
| --num-prompts 5120 \ | ||
| --num-generations 2 \ | ||
| --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=<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=<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=<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=<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> \ | ||
| 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=<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=<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. For local runs or custom images, install Tunix from GitHub: | ||
|
|
||
| ```bash | ||
| pip install "git+https://github.com/google/tunix@348959d18a4a09c75e58a7d49aec9d8b0eb4a8b6" | ||
|
Collaborator
There was a problem hiding this comment. if you check the We need to update this, and also add some background info how/when to change this commit number. @SurbhiJainUSC and @YixuanWang-99 |
||
| ``` | ||
|
|
||
| > **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=<GCS_BUCKET>/distillation/online \ | ||
| tokenizer_path=meta-llama/Llama-3.1-8B 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=<GCS_BUCKET>/pruned-llama3.1-8b-24L/0/items \ | ||
| teacher_overrides.model_name=llama3.1-8b \ | ||
| 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`). | ||
|
|
||
| #### Cluster Toolkit multi-host submission | ||
|
|
||
| Submit the distillation trainer directly as a Cluster Toolkit JobSet: | ||
|
|
||
| ```bash | ||
| 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> | ||
|
Collaborator
There was a problem hiding this comment. When is this env variable used? |
||
| 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 | ||
|
Collaborator
There was a problem hiding this comment. If the above 3 variables use the same <GCS_BUCKET>, we should export GCS_BUCKET as a env variable and re-use it in the 3 variables using ${GCS_BUCKET} |
||
| 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]" | ||
| ``` | ||
|
|
||
| #### Monitor and clean up | ||
|
|
||
| Monitor the workload and stream logs with Cluster Toolkit: | ||
|
|
||
| ```bash | ||
| # Check job status | ||
| gcluster job list | ||
|
|
||
| # Stream logs | ||
| gcluster job logs <RUN_NAME> | ||
|
|
||
| # Inspect JobSet and pods | ||
| kubectl get jobset -l gcluster.google.com/workload=<RUN_NAME> | ||
| kubectl get pods -l gcluster.google.com/workload=<RUN_NAME> | ||
|
|
||
| # Cancel workload | ||
| gcluster job cancel <RUN_NAME> | ||
| ``` | ||
|
|
||
| ### Offline top-k logits variant | ||
|
|
||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should keep the macros using env variables (${DISK_NAME?}, ${DISK_SIZE?}, ${ZONE?}). No need to change here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a common issue across this PR. We should use env variable
${...}, instead of the<...>when the variable has been defined.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed it back to env variable