diff --git a/docs/guides/distillation.md b/docs/guides/distillation.md index d2c7eebfc2..82552c9f54 100644 --- a/docs/guides/distillation.md +++ b/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) | | ------------------------------------------- | ------------------------------------------- | -------------------------- | diff --git a/docs/tutorials/posttraining/knowledge_distillation.md b/docs/tutorials/posttraining/knowledge_distillation.md index 7db716cc0e..3bb645cd56 100644 --- a/docs/tutorials/posttraining/knowledge_distillation.md +++ b/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= # e.g., hf_BA6... -export RUN_NAME= # e.g., distill-20260115 +export HF_TOKEN= # e.g., hf_BA6... +export RUN_NAME= # e.g., distill-20260115 ``` #### b. Install dependencies @@ -56,10 +56,10 @@ To store large models and datasets, attach a Hyperdisk to your TPU VM. Refer to First, create a Hyperdisk: ```bash -export ZONE= # e.g., us-central1-a -export TPU_VM_NAME= -export DISK_NAME= # e.g., my-hyperdisk -export DISK_SIZE= # e.g., 500GB +export ZONE= # e.g., us-central1-a +export TPU_VM_NAME= +export DISK_NAME= # e.g., my-hyperdisk +export DISK_SIZE= # e.g., 500GB gcloud compute disks create ${DISK_NAME?} \ --size=${DISK_SIZE?} \ @@ -87,7 +87,7 @@ 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= # e.g., knowledge-distillation +export BASE_NAME= # e.g., knowledge-distillation export BASE_OUTPUT_DIRECTORY=/mnt/hyperdisk/${BASE_NAME?} mkdir -p ${BASE_OUTPUT_DIRECTORY?} ``` @@ -96,7 +96,7 @@ mkdir -p ${BASE_OUTPUT_DIRECTORY?} ### 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: @@ -231,13 +231,13 @@ python3 -m maxtext.checkpoint_conversion.to_maxtext \ #### 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==` install. +> **Note:** Once a Tunix PyPI release ships, this will become a versioned `google-tunix==` install. ### Configuration @@ -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= -export XPK_PROJECT= -export XPK_ZONE= # e.g. us-central1-a -export XPK_DEVICE_TYPE= # e.g. tpu7x-4x4x4, v5p-128 -export XPK_BASE_OUTPUT_DIR=gs:///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= +export GKE_CLUSTER= +export LOCATION= # e.g., 'europe-west4' (region) or 'us-central1-a' (zone) +export RUN_NAME= +export IMAGE_URI= +export COMPUTE_TYPE= +export TOPOLOGY= +export BASE_OUTPUT_DIRECTORY=gs:///distillation +export STUDENT_CKPT_PATH=gs:////checkpoints/0/items +export TEACHER_CKPT_PATH=gs:////checkpoints/0/items +export TOKENIZER_PATH=meta-llama/Llama-3.1-8B +export HF_TOKEN= + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${GKE_CLUSTER?} \ + --location ${LOCATION?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${GKE_CLUSTER?} +gcluster job config set location ${LOCATION?} + +gcluster job submit \ + --image=${IMAGE_URI?} \ + --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=${BASE_OUTPUT_DIRECTORY?}/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_CKPT_PATH?} \ + teacher_overrides.model_name=llama3.1-8b \ + teacher_overrides.load_parameters_path=${TEACHER_CKPT_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 diff --git a/docs/tutorials/posttraining/lora_on_multi_host.md b/docs/tutorials/posttraining/lora_on_multi_host.md index cf2908b29f..fbe222d61e 100644 --- a/docs/tutorials/posttraining/lora_on_multi_host.md +++ b/docs/tutorials/posttraining/lora_on_multi_host.md @@ -36,8 +36,8 @@ Before starting, ensure you have: - Access to a Google Cloud Project with TPU quotas. - A Hugging Face account with an access token for downloading models. - Permissions for Google Artifact Registry (Artifact Registry Writer role). -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. +- A GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components (see [Cluster Toolkit documentation](https://cloud.google.com/cluster-toolkit/docs/overview)). - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Build and upload MaxText Docker image @@ -67,8 +67,7 @@ export HF_TOKEN= export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs # An arbitrary string to identify this specific run. -# We recommend to include the model, user, and timestamp. -# Note: Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores or periods). +# Note: Workload names cannot exceed 28 characters and must be valid DNS labels (lowercase alphanumeric and hyphens). export RUN_NAME= # -- Workload configuration -- @@ -77,22 +76,24 @@ export RUN_NAME= # gcloud config get-value project export PROJECT_ID= -# The GCP location (listed as "Location" in the UI) and name of your +# The GCP location (region or zone) and name of your # TPU-enabled GKE cluster. Both can be found on the # [Cloud Console](https://console.cloud.google.com/kubernetes/list). -export ZONE= # e.g., 'us-central1' +export LOCATION= # e.g., 'europe-west4' (region) or 'us-central1-a' (zone) export GKE_CLUSTER= # For a full list of MaxText-supported TPU types, see: `src/maxtext/utils/accelerator_to_spec_map.py`. To see the TPU type # of your cluster: # 1. Connect to the cluster (required for kubectl commands later): -# gcloud container clusters get-credentials ${GKE_CLUSTER?} --location ${ZONE?} --project ${PROJECT_ID?} +# gcloud container clusters get-credentials ${GKE_CLUSTER?} --location ${LOCATION?} --project ${PROJECT_ID?} # 2. Find your TPU type (e.g., 'v6e-256') by checking the accelerator labels on your nodes: # kubectl get nodes -l cloud.google.com/gke-tpu-accelerator -o jsonpath='{.items[*].metadata.labels.cloud\.google\.com/gke-tpu-accelerator}' | tr ' ' '\n' | sort -u export TPU_TYPE= export NUM_SLICES= +export COMPUTE_TYPE= +export TOPOLOGY= # The Docker image you pushed in the prerequisite step export CLOUD_IMAGE_NAME= @@ -111,6 +112,7 @@ export DATASET_NAME= # e.g., openai/gsm8k export TRAIN_SPLIT= # e.g., train export HF_DATA_DIR= # e.g., main export TRAIN_DATA_COLUMNS= # e.g., ['question','answer'] +export CHAT_TEMPLATE_PATH= # e.g., src/maxtext/examples/chat_templates/gemma_chat.json # -- LoRA Conversion configuration (Optional) -- export HF_LORA_ADAPTER_PATH= # e.g., 'username/adapter-name' @@ -149,7 +151,7 @@ checkpoint_storage_use_ocdbt=False ### Option 2: Converting a Hugging Face checkpoint -Refer to the steps in [Hugging Face to MaxText](https://maxtext.readthedocs.io/en/maxtext-v0.2.1/guides/checkpointing_solutions/convert_checkpoint.html#hugging-face-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on. +Refer to the steps in [Hugging Face to MaxText](hf-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on. ```bash export MAXTEXT_CKPT_PATH= # gs://my-bucket/my-checkpoint-directory/0/items @@ -158,19 +160,26 @@ export MAXTEXT_CKPT_PATH= # gs://my-bucket/my-checkpoint-directory/0/ ## Submit workload on GKE cluster This section provides the command to run LoRA Fine-Tuning on a GKE cluster. +Before submitting a job, configure access to the cluster with `gcloud` and `gcluster`: + +```bash +gcloud container clusters get-credentials ${GKE_CLUSTER?} \ + --location ${LOCATION?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${GKE_CLUSTER?} +gcluster job config set location ${LOCATION?} +``` ### Run a Fresh LoRA Fine-Tuning on Hugging Face Dataset ```bash -xpk workload create-pathways \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload=${RUN_NAME?} \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=${NUM_SLICES?} \ ---command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --command="\ python3 -m maxtext.trainers.post_train.sft.train_sft \ run_name=${RUN_NAME?} \ base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ @@ -190,8 +199,7 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \ lora.lora_rank=${LORA_RANK?} \ lora.lora_alpha=${LORA_ALPHA?} \ checkpoint_storage_use_zarr3=False \ - checkpoint_storage_use_ocdbt=False \ - enable_single_controller=True" + checkpoint_storage_use_ocdbt=False" ``` Once the fine-tuning is completed, you can access your model checkpoints at `${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints`. @@ -202,18 +210,17 @@ If you want to resume training from a previous run or further fine-tune an exist #### Step 1: Convert HF LoRA adapter to MaxText format +> For new deployments, run this conversion with Cluster Toolkit after configuring the cluster with `gcloud container clusters get-credentials` and `gcluster job config set` above. + If your LoRA adapter is currently in Hugging Face format, you must convert it to MaxText format before it can be loaded. Use the integrated conversion utility: ```sh -xpk workload create \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload=${RUN_NAME?} \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=${NUM_SLICES?} \ ---command="python3 -m maxtext.checkpoint_conversion.to_maxtext \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?}-convert \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --command="python3 -m maxtext.checkpoint_conversion.to_maxtext \ model_name=${MODEL?} \ hf_lora_adapter_path=${HF_LORA_ADAPTER_PATH?} \ base_output_directory=${BASE_OUTPUT_DIRECTORY?}/converted_adapter \ @@ -240,15 +247,15 @@ Once your environment variables and checkpoints are ready, you can start the LoR Execute the following command to begin training: ```bash -xpk workload create-pathways \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload=${RUN_NAME?} \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=${NUM_SLICES?} \ ---command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=${NUM_SLICES:-1} \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="\ python3 -m maxtext.trainers.post_train.sft.train_sft \ run_name=${RUN_NAME?} \ base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ @@ -277,23 +284,21 @@ Your fine-tuned model checkpoints will be saved here: `$BASE_OUTPUT_DIRECTORY/$R ## (Optional) Convert Fine-tuned LoRA to Hugging Face Format +> For new deployments, run this conversion with Cluster Toolkit after configuring the cluster with `gcloud container clusters get-credentials` and `gcluster job config set` above. + After completing the fine-tuning process, your LoRA weights are stored in MaxText/Orbax format. To use these weights with the Hugging Face ecosystem (e.g., for inference or sharing), convert them back using the `to_huggingface.py` script. ```sh -xpk workload create \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload="${RUN_NAME?}-to-hf" \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=1 \ ---command="python3 -m maxtext.checkpoint_conversion.to_huggingface \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name="${RUN_NAME?}-to-hf" \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --command="python3 -m maxtext.checkpoint_conversion.to_huggingface \ model_name=${MODEL?} \ lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints//model_params \ base_output_directory=${BASE_OUTPUT_DIRECTORY?}/hf_lora_adapter \ hf_access_token=${HF_TOKEN?}" - ``` - `lora.lora_restore_path`: Point this to the specific checkpoint directory (e.g., `.../checkpoints/1000/items`) that you want to export. diff --git a/docs/tutorials/posttraining/sft_on_multi_host.md b/docs/tutorials/posttraining/sft_on_multi_host.md index e9d0b0423c..5eb785c9a1 100644 --- a/docs/tutorials/posttraining/sft_on_multi_host.md +++ b/docs/tutorials/posttraining/sft_on_multi_host.md @@ -34,18 +34,27 @@ Before starting, ensure you have: - **Artifact Registry Writer** (`roles/artifactregistry.writer`) to upload Docker images. - **Storage Admin** (`roles/storage.admin`) or **Storage Object Admin** (`roles/storage.objectAdmin`) combined with **Storage Legacy Bucket Reader** (`roles/storage.legacyBucketReader`) on your GCS bucket to read/write checkpoints and logs. (Note: A bucket-level read permission like `storage.buckets.get` is required by JAX/TensorStore to verify bucket existence and metadata; using `roles/storage.objectAdmin` alone will cause a misleading "bucket not found" error). - A Hugging Face account with an access token for downloading models. -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. - **Important:** Modern GKE clusters require the GKE auth plugin. If you encounter `gke-gcloud-auth-plugin not found` when running `kubectl` commands, you must install it locally (e.g., `sudo apt-get install google-cloud-cli-gke-gcloud-auth-plugin` for `apt` installations, or `gcloud components install gke-gcloud-auth-plugin` for standalone archive installations). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- A GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components (see [Cluster Toolkit documentation](https://cloud.google.com/cluster-toolkit/docs/overview)). - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Build and upload MaxText Docker image For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](build-docker). -## Create GKE cluster +## Configure GKE cluster -Use a pathways ready GKE cluster as described [here](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster). +Configure access to the target cluster with `gcloud`, then configure the project, cluster, and location with `gcluster` as described in [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md): + +```bash +gcloud container clusters get-credentials ${GKE_CLUSTER?} \ + --location ${LOCATION?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${GKE_CLUSTER?} +gcluster job config set location ${LOCATION?} +``` ## Environment configuration @@ -71,8 +80,7 @@ export HF_TOKEN= export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs # An arbitrary string to identify this specific run. -# We recommend to include the model, user, and timestamp. -# Note: Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores or periods). +# Note: Workload names cannot exceed 28 characters (or 22 characters when using Pathways due to Kubernetes 63-byte coordinator label limits) and must be valid DNS labels (lowercase alphanumeric and hyphens). export RUN_NAME= # -- Workload configuration -- @@ -81,23 +89,28 @@ export RUN_NAME= # gcloud config get-value project export PROJECT_ID= -# The GCP location (listed as "Location" in the UI) and name of your +# The GCP location (region or zone) and name of your # TPU-enabled GKE cluster. Both can be found on the # [Cloud Console](https://console.cloud.google.com/kubernetes/list). -export ZONE= # e.g., 'us-central1' +export LOCATION= # e.g., 'europe-west4' (region) or 'us-central1-a' (zone) export GKE_CLUSTER= # For a full list of MaxText-supported TPU types, see: `src/maxtext/utils/accelerator_to_spec_map.py`. To see the TPU type # of your cluster: # 1. Connect to the cluster (required for kubectl commands later): -# gcloud container clusters get-credentials ${GKE_CLUSTER?} --location ${ZONE?} --project ${PROJECT_ID?} +# gcloud container clusters get-credentials ${GKE_CLUSTER?} --location ${LOCATION?} --project ${PROJECT_ID?} # 2. Find your TPU type (e.g., 'v5p-128') by checking the accelerator labels on your nodes: # kubectl get nodes -l cloud.google.com/gke-tpu-accelerator -o jsonpath='{.items[*].metadata.labels.cloud\.google\.com/gke-tpu-accelerator}' | tr ' ' '\n' | sort -u export TPU_TYPE= export NUM_SLICES= +# Cluster Toolkit workload placement. See the Cluster Toolkit guide for the +# compute type and topology matching your TPU slice. +export COMPUTE_TYPE= +export TOPOLOGY= + # The Docker image you pushed in the prerequisite step export CLOUD_IMAGE_NAME= export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" @@ -154,33 +167,71 @@ This section provides the command to run SFT on a GKE cluster. ### SFT with Multi-Controller JAX (McJAX) ```bash -xpk workload create \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload=${RUN_NAME?} \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=${NUM_SLICES?} \ ---command "python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${RUN_NAME?} base_output_directory=${BASE_OUTPUT_DIRECTORY?} model_name=${MODEL?} load_parameters_path=${MAXTEXT_CKPT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane hf_path=${DATASET_NAME?} train_split=${TRAIN_SPLIT?} train_data_columns=${TRAIN_DATA_COLUMNS?}" +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --command="python3 -m maxtext.trainers.post_train.sft.train_sft \ + run_name=${RUN_NAME?} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + model_name=${MODEL?} \ + load_parameters_path=${MAXTEXT_CKPT_PATH?} \ + hf_access_token=${HF_TOKEN?} \ + per_device_batch_size=1 \ + steps=${STEPS?} \ + profiler=xplane \ + hf_path=${DATASET_NAME?} \ + train_split=${TRAIN_SPLIT?} \ + train_data_columns=${TRAIN_DATA_COLUMNS?}" ``` Once the fine-tuning is completed, you can access your model checkpoints at `${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints`. ### SFT with Pathways +To submit an SFT workload with Pathways using Cluster Toolkit, use `gcluster job submit` with the `--pathways` flag: + ```bash export USE_PATHWAYS=1 -xpk workload create-pathways \ ---cluster=${GKE_CLUSTER?} \ ---project=${PROJECT_ID?} \ ---zone=${ZONE?} \ ---docker-image=${DOCKER_IMAGE?} \ ---workload=${RUN_NAME?} \ ---tpu-type=${TPU_TYPE?} \ ---num-slices=${NUM_SLICES?} \ ---command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 python3 -m maxtext.trainers.post_train.sft.train_sft run_name=${RUN_NAME?} base_output_directory=${BASE_OUTPUT_DIRECTORY?} model_name=${MODEL?} load_parameters_path=${MAXTEXT_CKPT_PATH?} hf_access_token=${HF_TOKEN?} per_device_batch_size=1 steps=${STEPS?} profiler=xplane checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) enable_single_controller=True" +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=${NUM_SLICES:-1} \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.sft.train_sft \ + run_name=${RUN_NAME?} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + model_name=${MODEL?} \ + load_parameters_path=${MAXTEXT_CKPT_PATH?} \ + hf_access_token=${HF_TOKEN?} \ + per_device_batch_size=1 \ + steps=${STEPS?} \ + profiler=xplane \ + checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \ + checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \ + enable_single_controller=True" ``` Once the fine-tuning is completed, you can access your model checkpoints at `${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints`. + +## Monitor and clean up + +```bash +gcluster job list +# Note: For Pathways workloads (> 5 pods), specify --main-only=false to retrieve logs from all pods: +gcluster job logs ${RUN_NAME?} --main-only=false +gcluster job cancel ${RUN_NAME?} +``` + +You can also inspect the Kubernetes resources directly: + +```bash +kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?} +# In Pathways workloads, use the jobset-name label to select all pods (both pathways-head and worker pods): +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} +``` diff --git a/src/maxtext/trainers/post_train/distillation/README.md b/src/maxtext/trainers/post_train/distillation/README.md index 768098099d..c18c1d9308 100644 --- a/src/maxtext/trainers/post_train/distillation/README.md +++ b/src/maxtext/trainers/post_train/distillation/README.md @@ -9,6 +9,8 @@ and [`post_training_index.md`](../../../../../docs/tutorials/post_training_index Canonical launcher: [`scripts/run_distill_xpk.sh`](scripts/run_distill_xpk.sh) (see its header for all env vars). +> **Note:** For new deployments on GKE, Cluster Toolkit (`gcluster job submit`) is recommended. See [`knowledge_distillation.md`](../../../../../docs/tutorials/posttraining/knowledge_distillation.md#cluster-toolkit-multi-host-submission) for submitting distillation jobs via Cluster Toolkit. + ## 1. Pick a config @@ -34,11 +36,11 @@ HBM — ≥ v5p-16 in practice; a v5p-8 only fits with bf16 weights or a shrunken student (see below): ```bash -source /bin/activate +source /bin/activate PYTHONPATH=$PWD/src python -m maxtext.trainers.post_train.distillation.train_distill \ src/maxtext/configs/post_train/distillation.yml \ run_name=local_smoke \ - base_output_directory=gs:///distill_smoke \ + base_output_directory=gs:///distill_smoke \ steps=5 ``` @@ -55,8 +57,8 @@ pip install git+https://github.com/AI-Hypercomputer/xpk.git # Kubeconfig (use --dns-endpoint; IP endpoints are often stale). # Use --zone for zonal clusters, --region for regional ones. -gcloud container clusters get-credentials \ - --zone= --project= --dns-endpoint +gcloud container clusters get-credentials \ + --zone= --project= --dns-endpoint # Verify RBAC in the default namespace: kubectl auth can-i create roles --namespace=default # must print: yes @@ -72,7 +74,7 @@ The flow is: build the MaxText base → `prep_image` rebuilds `$XPK_BASE_IMAGE` bakes the workspace `./src` into a runner image and pushes it to GCR. ```bash -export XPK_PROJECT= +export XPK_PROJECT= export XPK_BASE_IMAGE=maxtext_base_image # local tag prep_image rebuilds export XPK_RUNNER_IMAGE_NAME=maxtext_base_image # short name pushed under gcr.io/$XPK_PROJECT/ # XPK_RUNNER_IMAGE_TAG defaults to ${USER}-distill; override (or set USER) @@ -98,13 +100,13 @@ layout and xpk's crane-overlay fallback. ## 5. Submit ```bash -export XPK_CLUSTER= -export XPK_PROJECT= -export XPK_ZONE= +export XPK_CLUSTER= +export XPK_PROJECT= +export XPK_ZONE= export XPK_DEVICE_TYPE=tpu7x-4x4x4 export XPK_BASE_IMAGE=gcr.io/${XPK_PROJECT}/${XPK_RUNNER_IMAGE_NAME}:${XPK_RUNNER_IMAGE_TAG:-${USER}-distill} # slash → --docker-image auto-selected -export XPK_BASE_OUTPUT_DIR=gs:///distillation -export XPK_RUN_NAME= # default: distill_run; set per experiment +export XPK_BASE_OUTPUT_DIR=gs:///distillation +export XPK_RUN_NAME= # default: distill_run; set per experiment # to scope checkpoints + TB under # ${XPK_BASE_OUTPUT_DIR}/${XPK_WORKLOAD}/${XPK_RUN_NAME}/