diff --git a/README.md b/README.md index 3b735eb18..a775f4dac 100755 --- a/README.md +++ b/README.md @@ -115,7 +115,7 @@ Minimum requirements: Ubuntu Version 22.04 and Python 3.12. ## Getting Started -For your first time running MaxDiffusion, we provide specific [instructions](docs/getting_started/first_run.md). +For your first time running MaxDiffusion, we provide specific [instructions](docs/getting_started/first_run.md). To run multi-host jobs on GKE, see [Running MaxDiffusion with Cluster Toolkit](docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md). ## NVIDIA DGX Spark @@ -127,7 +127,7 @@ After installation completes, run the training script. ## Wan 2.1 Training - In the first part, we'll run on a single host VM to get familiar with the workflow, then run on xpk for large scale training. + In the first part, we'll run on a single host VM to get familiar with the workflow, then run on a GKE cluster with Cluster Toolkit for large scale training. Although not required, attaching an external disk is recommended as weights take up a lot of disk space. [Follow these instructions if you would like to attach an external disk](https://cloud.google.com/tpu/docs/attach-durable-block-storage). @@ -211,7 +211,7 @@ After installation completes, run the training script. ``` ```bash - export LIBTPU_INIT_ARGS='--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ + export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ --xla_tpu_megacore_fusion_allow_ags=false \ --xla_enable_async_collective_permute=true \ --xla_tpu_enable_ag_backward_pipelining=true \ @@ -235,7 +235,7 @@ After installation completes, run the training script. --xla_latency_hiding_scheduler_rerun=2 \ --xla_tpu_use_minor_sharding_for_major_trivial_input=true \ --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 \ - --xla_tpu_assign_all_reduce_scatter_layout=true' + --xla_tpu_assign_all_reduce_scatter_layout=true" ``` ```bash @@ -273,7 +273,7 @@ After installation completes, run the training script. It is important to note a couple of things: - per_device_batch_size can be fractional, but must be a whole number when multiplied by number of devices. In this example, 0.25 * 4 (devices) = effective global batch size = 1. - - The step time in v5p-8 with global batch size = 1 is large due to the aggressive remat policy (`HIDDEN_STATE_WITH_OFFLOAD`). On a larger number of chips we can run larger batch sizes, greatly increasing MFU, as we will see in the next section on deploying with xpk. + - The step time in v5p-8 with global batch size = 1 is large due to the aggressive remat policy (`HIDDEN_STATE_WITH_OFFLOAD`). On a larger number of chips we can run larger batch sizes, greatly increasing MFU, as we will see in the next section on deploying with Cluster Toolkit. - To enable eval during training set `eval_every` to a value > 0. - In Wan2.1, the ici_fsdp_parallelism axis is used for sequence parallelism, the ici_tensor_parallelism axis is used for head parallelism. - You can enable both, keeping in mind that Wan2.1 has 40 heads and 40 must be evenly divisible by ici_tensor_parallelism. @@ -315,11 +315,14 @@ After installation completes, run the training script. completed step: 6, seconds: 36.006, TFLOP/s/device: 135.900, loss: 0.169 ``` - ### Deploying with XPK + ### Deploying with Cluster Toolkit - This assumes the user has already created an xpk cluster, installed all dependencies and also created the dataset from the step above. For getting started with MaxDiffusion and xpk see [this guide](docs/getting_started/run_maxdiffusion_via_xpk.md). + This assumes the user has already created a GKE cluster with TPUs, installed `gcluster` (the [Cluster Toolkit](https://github.com/GoogleCloudPlatform/cluster-toolkit) CLI), pushed a MaxDiffusion image built from this checkout to Artifact Registry and also created the dataset from the step above. For getting started with MaxDiffusion and Cluster Toolkit see [this guide](docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md). - Using v5p-256, the command to run on xpk is as follows: + > [!NOTE] + > [XPK is deprecated](https://github.com/AI-Hypercomputer/xpk) in favor of Cluster Toolkit. The `gcluster` commands below were translated from the previous XPK commands using the [official migration guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/migration/xpk_to_clustertoolkit.md) and have not yet been validated end-to-end on a cluster. The legacy XPK commands are kept in the collapsed sections for users with existing XPK clusters. + + Using v5p-256, the command to run with `gcluster` is as follows: ```bash RUN_NAME=jfacevedo-wan-v5p-8-${RANDOM} @@ -330,7 +333,7 @@ After installation completes, run the training script. ``` ```bash - LIBTPU_INIT_ARGS='--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ + LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ --xla_tpu_megacore_fusion_allow_ags=false \ --xla_enable_async_collective_permute=true \ --xla_tpu_enable_ag_backward_pipelining=true \ @@ -354,9 +357,76 @@ After installation completes, run the training script. --xla_latency_hiding_scheduler_rerun=2 \ --xla_tpu_use_minor_sharding_for_major_trivial_input=true \ --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 \ - --xla_tpu_assign_all_reduce_scatter_layout=true' + --xla_tpu_assign_all_reduce_scatter_layout=true" ``` + ```bash + # Machine type and topology of the slice. Examples with 128 JAX devices, matching + # ici_data_parallelism=32 x ici_fsdp_parallelism=4 below (TPU7x exposes two devices per chip): + # v5p-256 -> COMPUTE_TYPE=ct5p-hightpu-4t TOPOLOGY=4x4x8 (128 chips) + # v6e-128 -> COMPUTE_TYPE=ct6e-standard-4t TOPOLOGY=8x16 (128 chips) + # tpu7x-4x4x4 -> COMPUTE_TYPE=tpu7x-standard-4t TOPOLOGY=4x4x4 (64 chips) + COMPUTE_TYPE=ct5p-hightpu-4t + TOPOLOGY=4x4x8 + # Image built from this checkout and pushed to Artifact Registry in step 3 of + # docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md. PROJECT_ID, LOCATION, CLUSTER_NAME, + # REGION and AR_REPO are the variables set in step 2 of that guide. + IMAGE=${REGION}-docker.pkg.dev/${PROJECT_ID}/${AR_REPO}/maxdiffusion_runner:latest + + gcluster job submit \ + --name=${RUN_NAME} \ + --cluster=$CLUSTER_NAME \ + --project=$PROJECT_ID \ + --location=$LOCATION \ + --compute-type=${COMPUTE_TYPE} \ + --topology=${TOPOLOGY} \ + --num-slices=1 \ + --image=${IMAGE} \ + --env="LIBTPU_INIT_ARGS=${LIBTPU_INIT_ARGS}" \ + --verbose \ + --command=" \ + pip install --no-deps . && HF_HUB_CACHE=/mnt/disks/external_disk/maxdiffusion_hf_cache/ python src/maxdiffusion/train_wan.py \ + src/maxdiffusion/configs/base_wan_14b.yml \ + attention='flash' \ + weights_dtype=bfloat16 \ + activations_dtype=bfloat16 \ + guidance_scale=5.0 \ + flow_shift=5.0 \ + fps=16 \ + skip_jax_distributed_system=False \ + run_name=${RUN_NAME} \ + output_dir=${OUTPUT_DIR} \ + train_data_dir=${DATASET_DIR} \ + load_tfrecord_cached=True \ + height=1280 \ + width=720 \ + num_frames=81 \ + num_inference_steps=50 \ + jax_cache_dir=${OUTPUT_DIR}/jax_cache/ \ + enable_profiler=True \ + dataset_save_location=${SAVE_DATASET_DIR} \ + remat_policy='HIDDEN_STATE_WITH_OFFLOAD' \ + flash_min_seq_length=0 \ + seed=$RANDOM \ + skip_first_n_steps_for_profiler=3 \ + profiler_steps=3 \ + per_device_batch_size=0.25 \ + ici_data_parallelism=32 \ + ici_fsdp_parallelism=4 \ + ici_tensor_parallelism=1 \ + max_train_steps=5000 \ + eval_every=100 \ + eval_data_dir=${EVAL_DATA_DIR} \ + enable_generate_video_for_eval=True" \ + --priority=medium \ + --restarts=0 + ``` + + Environment variables are not forwarded from your shell automatically, which is why `LIBTPU_INIT_ARGS` is passed with `--env` above (alternatively set them inline in `--command` as done for `HF_HUB_CACHE`). `--verbose` enables the same TPU debug logging as XPK's `--enable-debug-logs`. `pip install --no-deps .` reinstalls `maxdiffusion` from the checkout shipped in the image so that your local changes are the ones imported; see the [guide](docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md#4-submit-your-first-workload) for details. + +
+ Legacy XPK command (deprecated) + ```bash python3 ~/xpk/xpk.py workload create \ --cluster=$CLUSTER_NAME \ @@ -405,6 +475,8 @@ After installation completes, run the training script. --max-restarts=0 ``` +
+ ## Wan 2.2 Training Wan 2.2 introduces a **dual-expert DiT architecture** (High-Noise Expert and Low-Noise Expert, ~27B total parameters). MaxDiffusion supports joint dual-expert training where samples are dynamically routed to the appropriate expert based on `boundary_ratio` (default `0.875`) using the Flow Match time shift schedule. @@ -434,9 +506,43 @@ After installation completes, run the training script. save_final_checkpoint=True ``` - ### Multi-Host Training with XPK + ### Multi-Host Training with Cluster Toolkit - For large-scale multi-host training across TPU pods or clusters (e.g. v5p, v6e, v7x). The following example is configured for a 128-chip slice (such as `v5p-256`, `v6e-128`, or `tpu7x-4x4x4`): + This section covers large-scale multi-host training across TPU pods (e.g. v5p, v6e, v7x). The following example is configured for a slice with 128 JAX devices (such as `v5p-256`, `v6e-128`, or `tpu7x-4x4x4`; TPU7x exposes two devices per chip) and reuses the `COMPUTE_TYPE`, `TOPOLOGY` and `IMAGE` variables from the [Wan 2.1 Cluster Toolkit example](#deploying-with-cluster-toolkit). Pick a new `RUN_NAME` for each submission: + + ```bash + gcluster job submit \ + --name=${RUN_NAME} \ + --cluster=$CLUSTER_NAME \ + --project=$PROJECT_ID \ + --location=$LOCATION \ + --compute-type=${COMPUTE_TYPE} \ + --topology=${TOPOLOGY} \ + --num-slices=1 \ + --image=${IMAGE} \ + --command=" \ + pip install --no-deps . && python3 src/maxdiffusion/train_wan.py \ + src/maxdiffusion/configs/base_wan_27b.yml \ + run_name=${RUN_NAME} \ + output_dir=${OUTPUT_DIR} \ + train_data_dir=${DATASET_DIR} \ + dataset_save_location=${SAVE_DATASET_DIR} \ + boundary_ratio=0.875 \ + weights_dtype=bfloat16 \ + activations_dtype=bfloat16 \ + per_device_batch_size=1 \ + ici_fsdp_parallelism=4 \ + ici_data_parallelism=32 \ + remat_policy='HIDDEN_STATE_WITH_OFFLOAD' \ + max_train_steps=5000 \ + checkpoint_every=1000 \ + save_final_checkpoint=True" \ + --priority=medium \ + --restarts=0 + ``` + +
+ Legacy XPK command (deprecated) ```bash python3 ~/xpk/xpk.py workload create \ @@ -468,6 +574,8 @@ After installation completes, run the training script. --max-restarts=0 ``` +
+ ### Checkpointing & Export Wan 2.2 training checkpoints are managed with Orbax via `WanCheckpointer2_2`. Checkpoints store optimizer states and weights for both experts (`low_noise_transformer_state` and `high_noise_transformer_state`) alongside model configurations. MaxDiffusion inference pipelines (`WanPipeline2_2`) can load directly from these Orbax checkpoints for downstream sampling. @@ -971,7 +1079,7 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence ``` ## Getting Started: Multihost development -Multihost training for Stable Diffusion 2 base can be run using the following command: +For multi-host training on GKE, see [Running MaxDiffusion with Cluster Toolkit](docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md). On a multi-host TPU VM slice, multihost training for Stable Diffusion 2 base can be run using the following command: ```bash TPU_NAME= ZONE= diff --git a/docs/README.md b/docs/README.md index 5c41a595c..cb35b58b9 100644 --- a/docs/README.md +++ b/docs/README.md @@ -5,7 +5,8 @@ This folder contains documentation for getting started with and using MaxDiffusi ## Getting Started * **[First Run](getting_started/first_run.md)** - Provides instructions for setting up and running MaxDiffusion for the first time. -* **[Running MaxDiffusion via XPK](getting_started/run_maxdiffusion_via_xpk.md)** - Explains how to run MaxDiffusion on GKE using XPK. +* **[Running MaxDiffusion via Cluster Toolkit](getting_started/run_maxdiffusion_via_cluster_toolkit.md)** - Recommended way to run MaxDiffusion at scale on GKE using Cluster Toolkit's `gcluster` CLI. +* **[Running MaxDiffusion via XPK](getting_started/run_maxdiffusion_via_xpk.md)** - Explains how to run MaxDiffusion on GKE using XPK (deprecated, kept for existing XPK clusters). * **[NVIDIA DGX Spark](dgx_spark.md)** - Explains how to run MaxDiffusion on an NVIDIA DGX Spark. ## Contributing & Community diff --git a/docs/getting_started/first_run.md b/docs/getting_started/first_run.md index ccc452557..3c68f4750 100644 --- a/docs/getting_started/first_run.md +++ b/docs/getting_started/first_run.md @@ -33,4 +33,6 @@ source ~/$venv_name/bin/activate ## Getting Started: Multihost development -[GKE, recommended] [Running MaxDiffusion with xpk](run_maxdiffusion_via_xpk.md) - Quick Experimentation and Production support +[GKE, recommended] [Running MaxDiffusion with Cluster Toolkit](run_maxdiffusion_via_cluster_toolkit.md) - Quick Experimentation and Production support + +[GKE, deprecated] [Running MaxDiffusion with XPK](run_maxdiffusion_via_xpk.md) - Kept for existing XPK clusters only; XPK is deprecated in favor of Cluster Toolkit diff --git a/docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md b/docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md new file mode 100644 index 000000000..d6c2edd2b --- /dev/null +++ b/docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md @@ -0,0 +1,262 @@ + + +# How to run MaxDiffusion at scale with Cluster Toolkit (`gcluster`) + +This guide describes the recommended workflow for running MaxDiffusion on Google Kubernetes Engine (GKE) using +**Cluster Toolkit's `gcluster` CLI**. It replaces the previous [XPK guide](run_maxdiffusion_via_xpk.md): +[XPK is deprecated](https://github.com/AI-Hypercomputer/xpk) (see the notice at the top of its README) and new TPU/GPU +generations are only supported through Cluster Toolkit. + +For a complete reference see the [Cluster Toolkit repository](https://github.com/GoogleCloudPlatform/cluster-toolkit), +the [Google Cloud documentation](https://docs.cloud.google.com/cluster-toolkit/docs/overview), the +[`gcluster` job submission guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/gcluster_job_guide.md) +and the official [XPK → Cluster Toolkit migration guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/migration/xpk_to_clustertoolkit.md). + +> [!IMPORTANT] +> The `gcluster` commands in this guide (and the matching ones in the main [README](../../README.md)) were translated +> from the previous XPK commands using the official migration guide's 1:1 flag mapping. They have **not yet been +> validated end-to-end on a cluster** by the MaxDiffusion team. If something does not work as written, please +> [open an issue](https://github.com/AI-Hypercomputer/maxdiffusion/issues) so we can correct it. + +## Overview of the workflow + +1. **Package MaxDiffusion into a container image.** Build the dependency image once with + `docker_build_dependency_image.sh`, layer your current checkout on top of it with `maxdiffusion_runner.Dockerfile` + (takes seconds) and push the result to Artifact Registry. +2. **Submit the workload.** `gcluster job submit --image ...` inspects your GKE cluster, generates the Kubernetes + resources (a `JobSet` admitted through Kueue) and launches the multi-host job. + +```none ++--------------------------+ +--------------------+ +-------------------+ +| Your Development Machine +------> Artifact Registry +------> GKE Cluster | +| (docker + gcluster CLI) | | (stores images) | | (with TPUs/GPUs) | +| | | | | | +| 1. Build & push image | | 2. Nodes pull the | | 3. gcluster runs | +| from local checkout | | image | | multi-host job | ++--------------------------+ +--------------------+ +-------------------+ +``` + +## 1. Prerequisites + +### Required tools + +* **Google Cloud CLI (`gcloud`)** – install from [here](https://docs.cloud.google.com/sdk/docs/install-sdk) and run `gcloud init`. +* **kubectl and the GKE auth plugin**: + ```bash + gcloud components install kubectl gke-gcloud-auth-plugin + ``` + If `gcloud` was installed through a package manager (apt or snap, which is the case on TPU VMs), `gcloud components` + is disabled; install `kubectl` and `google-cloud-cli-gke-gcloud-auth-plugin` with that package manager instead (the + [XPK guide](run_maxdiffusion_via_xpk.md#steps-to-setup-xpk-on-tpu-vm) lists the apt commands). +* **`gcluster` CLI** (v1.103.0 or later, see the Cluster Toolkit + [security bulletins](https://docs.cloud.google.com/cluster-toolkit/docs/security-bulletins)) – follow the + [Cluster Toolkit setup guide](https://docs.cloud.google.com/cluster-toolkit/docs/setup/configure-environment) + or download a release bundle from the [releases page](https://github.com/GoogleCloudPlatform/cluster-toolkit/releases) + and make sure `gcluster` is on your `$PATH`: + ```bash + TAG=vX.Y.Z # see https://github.com/GoogleCloudPlatform/cluster-toolkit/releases + # Bundles are also published for linux_arm64, mac_amd64 and mac_arm64. + mkdir -p cluster-toolkit && \ + curl -fL https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${TAG?}/gcluster_bundle_linux_amd64.tgz \ + | tar -xz -C cluster-toolkit + export PATH="$PWD/cluster-toolkit:$PATH" + ``` +* **Docker** – used to build the MaxDiffusion images. Configure credentials for your Artifact Registry region: + ```bash + gcloud auth configure-docker -docker.pkg.dev --quiet + ``` + +### Google Cloud APIs and permissions + +```bash +gcloud services enable \ + container.googleapis.com \ + artifactregistry.googleapis.com \ + storage.googleapis.com +``` + +Your account needs at least the following IAM roles in the target project: + +* Artifact Registry Writer +* Kubernetes Engine Admin (or Developer on an existing cluster) +* Storage Admin (GCS buckets for datasets, checkpoints and outputs) +* Logging / Monitoring Viewer + +### A GKE cluster with accelerators + +This guide assumes you already have a GKE cluster with TPU (or GPU) node pools and the JobSet and Kueue controllers +installed. Clusters are provisioned with declarative blueprints via `gcluster deploy`; see the +[Cluster Toolkit GKE examples](https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/examples) and the +*Cluster Infrastructure Migration* section of the +[migration guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/migration/xpk_to_clustertoolkit.md) +for a worked example. Clusters created earlier with `xpk cluster create` keep running, but in-place migration is not +supported, so plan to recreate them with `gcluster deploy`. + +## 2. Environment configuration + +```bash +export PROJECT_ID= +export LOCATION= # e.g. us-east5-a (zonal cluster) or europe-west4 (regional cluster) +export CLUSTER_NAME= +export REGION= # Artifact Registry region, e.g. us-east5 +export AR_REPO=maxdiffusion-images # Artifact Registry repository name + +# Workload name: max 28 characters, lowercase alphanumerics and hyphens only (adjust if ${USER} does not comply). +export RUN_NAME=${USER}-first-job + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${CLUSTER_NAME?} --location ${LOCATION?} --project ${PROJECT_ID?} + +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${CLUSTER_NAME?} +gcluster job config set location ${LOCATION?} +``` + +Before submitting, `gcluster` runs prerequisite checks (gcloud authentication and Application Default Credentials, +`kubectl`, the GKE auth plugin, the Docker credential helper and the Artifact Registry API) and prints remediation +commands for anything that is missing. It then verifies (and installs if needed) the JobSet CRD on the cluster and +auto-discovers the Kueue `LocalQueue` to submit to (override it with `--queue`). + +## 3. Build and push the MaxDiffusion image + +MaxDiffusion does not publish a public base image, so build one from this repository. Run everything in this section +from the **root of your MaxDiffusion checkout**. + +```bash +# One-time: create the Artifact Registry repository. +gcloud artifacts repositories create ${AR_REPO?} \ + --repository-format=docker \ + --location=${REGION?} \ + --description="MaxDiffusion container images" + +# 1. Dependency image (slow, minutes). Produces the local tag maxdiffusion_base_image and only needs to be +# rebuilt when dependencies change. MODE=stable (default) or MODE=nightly; see the script header for all options. +bash docker_build_dependency_image.sh MODE=stable + +# 2. Runner image (fast, seconds): copies your current checkout into /deps on top of the dependency image. +docker build --build-arg BASEIMAGE=maxdiffusion_base_image -f maxdiffusion_runner.Dockerfile -t maxdiffusion_runner . + +# 3. Push it to Artifact Registry. +export IMAGE=${REGION?}-docker.pkg.dev/${PROJECT_ID?}/${AR_REPO?}/maxdiffusion_runner:latest +docker tag maxdiffusion_runner ${IMAGE?} +docker push ${IMAGE?} +``` + +Repeat steps 2–3 whenever your local code changes (`docker_upload_runner.sh` automates them, but pushes to +`gcr.io/` rather than Artifact Registry). + +> [!NOTE] +> The MaxText guide and the migration guide use `gcluster`'s on-the-fly image build (`--base-image` + `--build-context .`) +> instead of a prebuilt image. That mode appends the build context at the **root** of the image filesystem and keeps the +> base image's working directory, whereas MaxDiffusion's images set `WORKDIR /deps` and already contain a copy of the +> source tree there. A command such as `python src/maxdiffusion/train.py` would therefore silently run the copy baked +> into the base image rather than your local changes, which is why this guide uses `--image` with a runner image. + +## 4. Submit your first workload + +`${IMAGE}` contains your checkout under `/deps`, the container's working directory, so paths in `--command` are +relative to the repository root. The dependency image also installs `maxdiffusion` into site-packages at build time, +while the runner image only refreshes the source tree; prefix the command with `pip install --no-deps . &&` (the XPK +guide used `pip install . &&`) so that your latest changes are the ones imported. You can omit it if you rebuilt the +dependency image from the same checkout. + +```bash +export COMPUTE_TYPE=v6e-8 # TPU shorthand; see "Choosing --compute-type" below +export OUTPUT_DIR=gs:/// + +gcluster job submit \ + --name=${RUN_NAME?} \ + --compute-type=${COMPUTE_TYPE?} \ + --num-slices=1 \ + --image=${IMAGE?} \ + --command="pip install --no-deps . && python src/maxdiffusion/train.py src/maxdiffusion/configs/base_2_base.yml run_name=${RUN_NAME?} output_dir=${OUTPUT_DIR?}" +``` + +`--cluster`, `--project` and `--location` can be omitted because they were stored with `gcluster job config set` +above; pass them explicitly to target a different cluster. + +For full-scale examples (Wan 2.1 / Wan 2.2 training on 128-device slices, including the recommended `LIBTPU_INIT_ARGS` +and sharding flags) see the [Deploying with Cluster Toolkit](../../README.md#deploying-with-cluster-toolkit) and +[Multi-Host Training with Cluster Toolkit](../../README.md#multi-host-training-with-cluster-toolkit) sections of the +main README. + +### Choosing `--compute-type` + +* Common shorthands such as `v4-8`, `v6e-8`, `v6e-16`, `l4-8` or `h100-80gb-8` can be passed directly; `gcluster` + resolves the machine type and topology for you. +* For shapes that are not in the shorthand map (for example `v5e-*`), for multi-host slices where you want to pin the + shape, and **always for TPU7x**, pass the GCE machine type plus an explicit `--topology`, e.g. + `--compute-type=ct5p-hightpu-4t --topology=4x4x8` (v5p-256, 128 chips), + `--compute-type=ct6e-standard-4t --topology=8x16` (v6e-128, 128 chips) or + `--compute-type=tpu7x-standard-4t --topology=4x4x4` (64 chips; TPU7x exposes two JAX devices per chip, so this is + also 128 devices). +* `--num-nodes` is for GPU/CPU jobs only; omit it for TPU jobs. + +### Environment variables + +Environment variables from your shell are **not** forwarded to the job. Pass them explicitly with `--env`, e.g. +`--env HF_HUB_ENABLE_HF_TRANSFER=1` or `--env "LIBTPU_INIT_ARGS=${LIBTPU_INIT_ARGS}"`. Alternatively prefix them inside +`--command` as the README examples do (`HF_HUB_CACHE=... python ...`). + +### Mounting storage + +`xpk storage attach` is replaced by the inline `--mount` flag, which accepts GCS buckets, Filestore instances or +PVC claim names, e.g. `--mount 'gs:///datasets;/mnt/data;ro'` (add `;options=implicit-dirs` for GCS Fuse +options). MaxDiffusion reads `gs://` paths directly, so mounting is optional. + +## 5. Monitoring and managing workloads + +```bash +# List active or queued workloads. +gcluster job list + +# Stream logs (or use the Cloud Logging link printed by gcluster job submit). +gcluster job logs ${RUN_NAME?} + +# Pod status and logs of an individual pod. +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} +kubectl logs -f + +# Cancel a running workload / clean up a finished JobSet. +gcluster job cancel ${RUN_NAME?} +``` + +## 6. Mapping from the previous XPK commands + +The flags used by MaxDiffusion's old XPK commands map to `gcluster job submit` as follows (full table in the +[migration guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/migration/xpk_to_clustertoolkit.md#7-xpk--ct-command-mappings-table)): + +| `xpk workload create` | `gcluster job submit` | Notes | +| :--- | :--- | :--- | +| `--workload NAME` | `--name NAME` | Max 28 characters | +| `--cluster` / `--project` | same | Optional once set via `gcluster job config set` | +| `--zone ZONE` | `--location ZONE_OR_REGION` | | +| `--tpu-type` / `--device-type TYPE` | `--compute-type TYPE [--topology T]` | Shorthand or machine type + topology | +| `--num-slices N` | `--num-slices N` | | +| `--base-docker-image IMG` | `--image IMG` | `IMG` must be a registry image, not a local Docker tag such as `maxdiffusion_base_image`: push a runner image built from your checkout (section 3). `--base-image IMG --build-context .` is not suitable for MaxDiffusion images, see the note in section 3 | +| `--docker-image IMG` | `--image IMG` | Pre-built image containing the code | +| `--command "..."` | `--command "..."` | | +| `--env KEY=VAL` | `--env KEY=VAL` | | +| `--priority P` | `--priority P` | `low`, `medium`, `high` | +| `--max-restarts N` | `--restarts N` | | +| `--enable-debug-logs` | `--verbose` | Both set `TPU_STDERR_LOG_LEVEL=0`, `TPU_MIN_LOG_LEVEL=0`, `TF_CPP_MIN_LOG_LEVEL=0` and `TPU_VMODULE=real_program_continuator=1` in the containers | +| `xpk workload list` / `delete` | `gcluster job list` / `gcluster job cancel NAME` | | +| `xpk inspector` | `gcluster job logs NAME` | | + +When `run_name` is empty, MaxDiffusion ([pyconfig.py](../../src/maxdiffusion/pyconfig.py)) falls back to the +`JOBSET_NAME` environment variable. XPK injected that variable into every pod; `gcluster` does not, so always pass +`run_name=...` explicitly (as all examples here do) or add `--env JOBSET_NAME=${RUN_NAME?}`. diff --git a/docs/getting_started/run_maxdiffusion_via_xpk.md b/docs/getting_started/run_maxdiffusion_via_xpk.md index 35452a5ab..3e3e4f028 100644 --- a/docs/getting_started/run_maxdiffusion_via_xpk.md +++ b/docs/getting_started/run_maxdiffusion_via_xpk.md @@ -1,6 +1,9 @@ # How to run MaxDiffusion with XPK? -This document focuses on steps required to setup XPK on TPU VM and assumes you have gone through the [README](https://github.com/google/xpk/blob/main/README.md) to understand XPK basics. +> [!WARNING] +> **XPK is deprecated.** Per the [XPK deprecation notice](https://github.com/AI-Hypercomputer/xpk), the project is in maintenance mode through Q3 2026 and will then be archived; new TPU and GPU generations are only supported through Cluster Toolkit. Existing XPK clusters keep running, but in-place migration is not supported. For new work please follow [Running MaxDiffusion with Cluster Toolkit](run_maxdiffusion_via_cluster_toolkit.md) and the official [XPK → Cluster Toolkit migration guide](https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/docs/migration/xpk_to_clustertoolkit.md). The instructions below are kept for users with existing XPK clusters. + +This document focuses on steps required to setup XPK on TPU VM and assumes you have gone through the [README](https://github.com/AI-Hypercomputer/xpk/blob/main/README.md) to understand XPK basics. ## Steps to setup XPK on TPU VM @@ -64,7 +67,7 @@ after which log out and log back in to the machine. 3. After building the dependency image `maxdiffusion_base_image`, xpk can handle updates to the working directory when running `xpk workload create` and using `--base-docker-image`. - See details on docker images in xpk here: https://github.com/google/xpk/blob/main/README.md#how-to-add-docker-images-to-a-xpk-workload + See details on docker images in xpk here: https://github.com/AI-Hypercomputer/xpk/blob/main/docs/usage/docker.md **Note:** When using the XPK command, ensure you include `pip install .` to install the package from the current directory. This is necessary because the container is created from a copy of your local directory, and `pip install .` ensures any local changes you've made are applied within the container. @@ -91,10 +94,10 @@ after which log out and log back in to the machine. --command "pip install . && python src/maxdiffusion/train.py src/maxdiffusion/configs/base_2_base.yml run_name="my_run" output_dir="gs://your-bucket/"" ``` - __Using [xpk github repo](https://github.com/google/xpk.git)__ + __Using [xpk github repo](https://github.com/AI-Hypercomputer/xpk.git)__ ```shell - git clone https://github.com/google/xpk.git + git clone https://github.com/AI-Hypercomputer/xpk.git # Make sure you are still in the MaxDiffusion github root directory when running this command python3 xpk/xpk.py workload create \ diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index d3f4ed5dc..20bb66de0 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -214,7 +214,7 @@ def user_init(raw_keys): raw_keys["weights_dtype"] = jax.numpy.dtype(raw_keys["weights_dtype"]) raw_keys["activations_dtype"] = jax.numpy.dtype(raw_keys["activations_dtype"]) if raw_keys["run_name"] == "": - raw_keys["run_name"] = os.environ.get("JOBSET_NAME") # using XPK default + raw_keys["run_name"] = os.environ.get("JOBSET_NAME") # injected by XPK; not set by Cluster Toolkit run_name = raw_keys["run_name"] base_output_directory = raw_keys["output_dir"] if run_name: