|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -xe |
| 3 | + |
| 4 | +export RANDOM_CHARS=$(LC_CTYPE=C openssl rand -base64 12 | tr -dc 'a-z0-9' | head -c 3 ; echo) |
| 5 | +#export PROJECT_ID="cloud-tpu-best-effort-colo" |
| 6 | +export PROJECT_ID="cloud-tpu-multipod-dev" |
| 7 | +export NUM_REPLICAS=${NUM_REPLICAS:-2} |
| 8 | +export JOBSET_NAME=axlearn-image-test-$RANDOM_CHARS |
| 9 | +#export JOBSET_NAME="sujinesh-64-orbax-3" |
| 10 | +export BASTION_TIER=disabled |
| 11 | +export GKE_CLUSTER=$(axlearn gcp config | grep gke_cluster | awk '{ print $3 }' | tr -d '"') |
| 12 | +# Switch to tpu-v6e-256 if on scale cluster |
| 13 | +export INSTANCE_TYPE=${INSTANCE_TYPE:-"tpu-v5p-128"} |
| 14 | +# Switch to tpu-v6e-256-4 if on scale cluster |
| 15 | +export MESH_SELECTOR=${MESH:-"tpu-v5p-128"} |
| 16 | +export CONFIG=${CONFIG:-"fuji-7B-v2-flash"} |
| 17 | +#export OUTPUT_DIR=${OUTPUT_DIR:-gs://tess-checkpoints-flat-us-east5} |
| 18 | +export OUTPUT_DIR="gs://axlearn-ml-solutions" |
| 19 | +# export DATA_DIR="gs://tess-apple-southamerica-west1/tensorflow_datasets" |
| 20 | +#export DATA_DIR="gs://tess-dataloading-us-east5/tensorflow_datasets" |
| 21 | +export DATA_DIR="gs://axlearn-public/tensorflow_datasets" |
| 22 | + |
| 23 | + |
| 24 | +# Example for v6e-256 |
| 25 | +# MESH_SELECTOR=tpu-v6e-256-4 INSTANCE_TYPE=tpu-v6e-256 ./test-orbax.sh |
| 26 | + |
| 27 | +# The bundle step is needed if you run on cloudtop |
| 28 | +# uncomment if you use cloudtop |
| 29 | +axlearn gcp bundle --name=$JOBSET_NAME \ |
| 30 | + --bundler_spec=allow_dirty=True \ |
| 31 | + --bundler_type=artifactregistry \ |
| 32 | + --bundler_spec=dockerfile=Dockerfile \ |
| 33 | + --bundler_spec=image=tpu \ |
| 34 | + --bundler_spec=target=tpu |
| 35 | + |
| 36 | +# Only enable kueue when running on scale testing cluster |
| 37 | +# --queue=multislice-queue \ |
| 38 | +# --priority_class=very-high \ |
| 39 | +# --trainer_dir=gs://tess-checkpoints-us-west1/${JOBSET_NAME}-nr-${NUM_REPLICAS}/ \ |
| 40 | +# |
| 41 | + |
| 42 | +# Check if CONFIG ends with "orbaxem" |
| 43 | +if [[ "$CONFIG" == *"orbaxem"* ]]; then |
| 44 | + echo "Running with Orbax emergency checkpointer." |
| 45 | + axlearn gcp launch run --cluster=$GKE_CLUSTER \ |
| 46 | + --runner_name gke_tpu_single \ |
| 47 | + --name=$JOBSET_NAME \ |
| 48 | + --instance_type=${INSTANCE_TYPE} \ |
| 49 | + --host_mount_spec=name=tmp,host_path=/tmp,mount_path=/host-tmp \ |
| 50 | + --num_replicas=${NUM_REPLICAS} \ |
| 51 | + --bundler_spec=allow_dirty=True \ |
| 52 | + --bundler_type=artifactregistry --bundler_spec=image=tpu \ |
| 53 | + --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu \ |
| 54 | + -- "patch /opt/venv/lib/python3.10/site-packages/jax/experimental/shard_map.py -p0 < patches/shard_map.py.patch; ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main" \ |
| 55 | + --init_module=axlearn.common.checkpointer_orbax_emergency:local_ckpt_dir=/host-tmp/checkpoints \ |
| 56 | + --module=text.gpt.c4_trainer \ |
| 57 | + --config=${CONFIG} \ |
| 58 | + --trainer_dir=gs://${PROJECT_ID}-axlearn/${JOBSET_NAME}-nr-${NUM_REPLICAS}/ \ |
| 59 | + --data_dir=gs://axlearn-public/tensorflow_datasets \ |
| 60 | + --jax_backend=tpu \ |
| 61 | + --mesh_selector=${MESH_SELECTOR} \ |
| 62 | + --initialization_timeout=1200 \ |
| 63 | + --trace_at_steps=29,59,89,119,149,179,209,239,269,299,329,359,389,419,449,479,509,539,569,599,629,659,689,719 |
| 64 | + |
| 65 | +else |
| 66 | + echo "Running Orbax regular checkpointer or AXLearn native." |
| 67 | + axlearn gcp launch run --cluster=$GKE_CLUSTER \ |
| 68 | + --runner_name gke_tpu_single \ |
| 69 | + --name=$JOBSET_NAME \ |
| 70 | + --max_tries=100 \ |
| 71 | + --instance_type=${INSTANCE_TYPE} \ |
| 72 | + --num_replicas=${NUM_REPLICAS} \ |
| 73 | + --bundler_spec=allow_dirty=True \ |
| 74 | + --bundler_type=artifactregistry --bundler_spec=image=tpu \ |
| 75 | + --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu \ |
| 76 | + -- "ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main" \ |
| 77 | + --module=text.gpt.c4_trainer \ |
| 78 | + --config=${CONFIG} \ |
| 79 | + --trainer_dir=${OUTPUT_DIR}/${JOBSET_NAME}/ \ |
| 80 | + --data_dir=$DATA_DIR \ |
| 81 | + --jax_backend=tpu \ |
| 82 | + --mesh_selector=${MESH_SELECTOR} \ |
| 83 | + --initialization_timeout=1200 \ |
| 84 | + --trace_at_steps=50,101 |
| 85 | +fi |
| 86 | + |
| 87 | + |
| 88 | + # "ulimit -n 1048576; ulimit -c 0; sed -i 's|max_step = TOTAL_TOKENS\\[version\\]\\[model_size\\] // tokens_per_batch|max_step = 200|; /max_step = 200/a \\ save_every_n_steps=50' axlearn/experiments/text/gpt/fuji.py;python3 -c 'import jax; jax.devices()'; |
0 commit comments