hpc-slurm6-tpu-maxtext.yaml

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

---

blueprint_name: slurm6-tpu-v4

vars:
  project_id: ## Set GCP Project ID Here ##
  deployment_name: slurm6-tpu-v4
  region: us-central2
  zone: us-central2-b

deployment_groups:
- group: primary
  modules:
  - id: network
    source: modules/network/vpc

  - id: script
    source: modules/scripts/startup-script
    settings:
      runners:
      - type: shell
        destination: tpu_setup.sh
        content: |
          #!/bin/bash
          mkdir -p /opt/apps/scripts/tpu-test
          chmod a+rwx /opt/apps/scripts/tpu-test

      - type: data
        destination: /opt/apps/scripts/tpu-test/run_maxtext.sh
        content: |
          #!/bin/bash
          # Update the following parameters in the python3 command in this script.
          # RUN_NAME: (required, see https://github.com/google/maxtext?tab=readme-ov-file#overview)
          # STORAGE_BUCKET: GCS bucket where dataset is stored. (i.e.: gs://dataset-tpu/dataset)
          # ATTENTION: (i.e. dot_product, flash)
          # STEPS: steps to run (i.e. 100, 1000)

          #SBATCH --nodes=1
          #SBATCH --ntasks-per-node=1
          #SBATCH --mem=50000
          #SBATCH --partition=tpu

          set -e -o pipefail

          # Start virtual environment.
          virtualenv venv
          source venv/local/bin/activate

          # Clone maxtext repository, lock version of maxtext and
          # install dependencies.
          git clone https://github.com/google/maxtext
          cd maxtext/
          git reset --hard 39a3f19e832016741be803ef5253333e2f434cb8
          bash setup.sh

          # Run maxtext benchmark test.
          # (i.e.) python3 MaxText/train.py MaxText/configs/base.yml run_name=1xv4-128 base_output_directory=${PWD}/output/ dataset_path=gs://dataset_tpu/dataset async_checkpointing=False attention=dot_product steps=100
          python3 MaxText/train.py MaxText/configs/base.yml run_name=<RUN_NAME> base_output_directory=${PWD}/output/ dataset_path=<STORAGE_BUCKET> async_checkpointing=False attention=<ATTENTION> steps=<STEPS>

  - id: tpu_nodeset
    source: community/modules/compute/schedmd-slurm-gcp-v6-nodeset-tpu
    use: [network]
    settings:
      node_type: v4-8
      tf_version: 2.14.0
      # Preemptible TPUs cost much less than non-preemptible TPUs.
      # The Cloud TPU service might preempt (shut down) these TPUs at any time.
      # https://cloud.google.com/tpu/docs/preemptible
      preemptible: false
      # Specify whether to preserve TPU on suspend.
      # If set to true, suspended VM will be stopped.
      # If set to false, suspended VM will be deleted.
      preserve_tpu: false
      node_count_dynamic_max: 1

  - id: tpu_partition
    source: community/modules/compute/schedmd-slurm-gcp-v6-partition
    use: [tpu_nodeset]
    settings:
      partition_name: tpu

  - id: compute_nodeset
    source: community/modules/compute/schedmd-slurm-gcp-v6-nodeset
    use: [network]
    settings:
      name: ns2
      node_count_dynamic_max: 20
      bandwidth_tier: gvnic_enabled
      allow_automatic_updates: false

  - id: compute_partition
    source: community/modules/compute/schedmd-slurm-gcp-v6-partition
    use: [compute_nodeset]
    settings:
      partition_name: compute
      is_default: true

  - id: slurm_login
    source: community/modules/scheduler/schedmd-slurm-gcp-v6-login
    use: [network]
    settings:
      enable_login_public_ips: true
      machine_type: n2-standard-16

  - id: slurm_controller
    source: community/modules/scheduler/schedmd-slurm-gcp-v6-controller
    use:
    - tpu_partition
    - compute_partition
    - slurm_login
    - network
    settings:
      enable_controller_public_ips: true
      machine_type: n2-standard-16
      login_startup_script: $(script.startup_script)