|
| 1 | +<!-- |
| 2 | + # Copyright 2023–2025 Google LLC |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + --> |
| 16 | + |
| 17 | +# Supervised Fine-Tuning (SFT) with Deepseek-V3 Model |
| 18 | +This guide provides step by step instructions to run SFT with Deepseek-V3 model on TPU v6e-256. Deepseek-V3 is a Mixture-of-Experts (MoE) language model with 671B parameters. |
| 19 | + |
| 20 | +## 1. Build and Upload MaxText Docker Image |
| 21 | +This section guides you through cloning the MaxText repository, building MaxText Docker image with dependencies, and uploading the docker image to your project's Artifact Registry. |
| 22 | + |
| 23 | +### 1.1. Clone the MaxText Repository |
| 24 | +```bash |
| 25 | +git clone https://github.com/google/maxtext.git |
| 26 | +cd maxtext |
| 27 | +``` |
| 28 | + |
| 29 | +### 1.2. Build MaxText Docker Image |
| 30 | +```bash |
| 31 | +bash docker_build_dependency_image.sh MODE=jax_ai_image BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest |
| 32 | +``` |
| 33 | +This creates a local Docker image named `maxtext_base_image`. |
| 34 | + |
| 35 | +### 1.3. Upload the Docker Image to Artifact Registry |
| 36 | +```bash |
| 37 | +# Replace `$USER_runner` with your desired image name |
| 38 | +export DOCKER_IMAGE_NAME=${USER}_runner |
| 39 | +bash docker_upload_runner.sh CLOUD_IMAGE_NAME=$DOCKER_IMAGE_NAME |
| 40 | +``` |
| 41 | +The `docker_upload_runner.sh` script uploads your Docker image to Artifact Registry. |
| 42 | + |
| 43 | +## 2. Install XPK |
| 44 | +Install XPK by following the instructions in the [official documentation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation-via-pip). |
| 45 | + |
| 46 | +## 3. Create GKE Cluster |
| 47 | +If you don't already have a GKE cluster with a `v6e-256` TPU slice available, create one by following the [XPK cluster creation guide](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#cluster-create). |
| 48 | + |
| 49 | +## 4. Environment Configuration |
| 50 | +```bash |
| 51 | +# -- Google Cloud Configuration -- |
| 52 | +export PROJECT=<Google Cloud Project ID> |
| 53 | +export CLUSTER_NAME=<Name of GKE Cluster> |
| 54 | +export ZONE=<GKE Cluster Zone> |
| 55 | + |
| 56 | +# -- Workload Configuration -- |
| 57 | +export WORKLOAD_NAME="sft-$(date +%Y-%m-%d-%H-%M-%S)" # Or your desired workload name |
| 58 | +export TPU_TYPE=v6e-256 |
| 59 | +export TPU_SLICE=1 |
| 60 | +export DOCKER_IMAGE="gcr.io/${PROJECT}/${DOCKER_IMAGE_NAME}" |
| 61 | + |
| 62 | +# -- MaxText Configuration -- |
| 63 | +export OUTPUT_PATH=<GCS Bucket Path for output/logs> |
| 64 | +export STEPS=100 # Number of fine-tuning steps to run |
| 65 | +export HF_TOKEN=<Hugging Face access token> |
| 66 | +export MODEL_CHECKPOINT_PATH=<GCS path to model checkpoint> |
| 67 | +``` |
| 68 | + |
| 69 | +## 5. Submit Workload on GKE Cluster |
| 70 | +This section provides the command to run SFT with Deepseek-v3 model on a v6e-256 GKE cluster. |
| 71 | +```bash |
| 72 | +xpk workload create \ |
| 73 | +--cluster=${CLUSTER_NAME} \ |
| 74 | +--project=${PROJECT} \ |
| 75 | +--zone=${ZONE} \ |
| 76 | +--docker-image=${DOCKER_IMAGE} \ |
| 77 | +--workload=${WORKLOAD_NAME} \ |
| 78 | +--tpu-type=${TPU_TYPE} \ |
| 79 | +--num-slices=${TPU_SLICE} \ |
| 80 | +--command "python3 -m MaxText.sft.sft_trainer MaxText/configs/sft.yml run_name=$WORKLOAD_NAME base_output_directory=$OUTPUT_PATH model_name=deepseek3-671b load_parameters_path=$MODEL_CHECKPOINT_PATH hf_access_token=$HF_TOKEN tokenizer_path=deepseek-ai/DeepSeek-V3 per_device_batch_size=1 steps=$STEPS profiler=xplane megablox=False sparse_matmul=False ici_expert_parallelism=16 ici_fsdp_parallelism=16 weight_dtype=bfloat16 dtype=bfloat16 remat_policy=full decoder_layer_input=offload sa_block_q=2048 sa_block_q_dkv=2048 sa_block_q_dq=2048 opt_type=sgd attention=flash capacity_factor=1.0 max_target_length=2048" |
| 81 | +``` |
| 82 | +Once the fine-tuning is completed, you can access your model checkpoint at `${OUTPUT_PATH}/${WORKLOAD_NAME}/checkpoints/${STEPS}/model_params`. |
| 83 | + |
0 commit comments