Docs migration from XPK to CTK(2/5) - #5176
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request migrates the post-training documentation (covering knowledge distillation, LoRA, and SFT) from using the XPK launcher to using the Cluster Toolkit (gcluster) for submitting and managing multi-host TPU workloads on GKE. Feedback was provided regarding the updated gcluster job submit command for multi-host distillation, which is missing critical parameters such as tokenizer paths and student model overrides required for successful execution.
KevinYenky
changed the title
| # Note: Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores or periods). | ||
| export RUN_NAME=<RUN_NAME> | ||
| # 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="sft-$(date +%m%d%H%M%S)" |
There was a problem hiding this comment.
I used this date postfix in my test, but we can remove it here and stay consistent to the previous.
There was a problem hiding this comment.
reverted it to the previous name.
|
|
||
| # 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?} --zone ${ZONE?} --project ${PROJECT_ID?} |
There was a problem hiding this comment.
May I know is it relevant to CTK migration?
I checked that both location and zone works but location is preferred ref: doc
There was a problem hiding this comment.
Sure. Changed it to the recommended --location
| --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?} \ | ||
| --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?}" \ |
There was a problem hiding this comment.
inconsistent on using --command= or
…evinYenky/maxtext into docs/ctk-migration-core-guides
2 participants
Description
This PR is the second part of decomposing XPK to Cluster Toolkit document migration. It focuses on updating our post-training tutorial documentation to use Cluster Toolkit (CTK) instead of the deprecated XPK launcher.
Specifically, this PR updates the following post-training guides:
xpkcommands withgcluster job submit.xpk workload create-pathwayswith the multi-hostgclusterequivalent, ensuring all student/teacher model overrides and tokenizer paths are properly documented.Why this is being made:
CTK is replacing XPK as the primary submission mechanism for Google Kubernetes Engine (GKE) workloads. Breaking down the massive PR #5175 into smaller, logical chunks (starting with these core guides) makes the migration significantly easier to review and test.
Tests
tested by scripts that run the gcluster command provided in the documents, results stored in gcs bucket.
Checklist
gemini-reviewlabel.