Skip to content
Benchmarks

Create a workflow to run benchmarks #90

Sign in to view logs

Create a workflow to run benchmarks

Create a workflow to run benchmarks #90

Workflow file for this run

.github/workflows/benchmarks.yml at 808d88e
name: Benchmarks
on:
pull_request:
branches:
- main
workflow_dispatch:
inputs:
halt-for-connection:
description: 'Should this workflow run wait for a remote connection?'
type: choice
required: true
default: 'no'
options:
- 'yes'
- 'no'
jobs:
build-xla-gpu-and-test:
runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
container:
image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
options: --gpus all --privileged # Might need privileged mode, use with caution
steps:
# - name: Checkout XLA
# uses: actions/checkout@v3
# with:
# repository: openxla/xla # Replace with your fork if needed
# path: xla
- name: Checkout repository
uses: actions/checkout@v3
with:
repository: juliagmt-google/xla
path: xla
- name: Create results directory
working-directory: xla
run: mkdir -p results
- name: Get GPU spec
working-directory: xla
continue-on-error: true
run: nvidia-smi
- name: Configure XLA
working-directory: xla
run: ./configure.py --backend CUDA --nccl
- name: Set TF_CPP_MAX_VLOG_LEVEL
working-directory: xla
run: echo "TF_CPP_MAX_VLOG_LEVEL=1" >> $GITHUB_ENV # Use GITHUB_ENV to persist across steps
- name: Check TF_CPP_MAX_VLOG_LEVEL
working-directory: xla
run: echo "$TF_CPP_MAX_VLOG_LEVEL"
- name: Wait For Connection
uses: google-ml-infra/actions/ci_connection@main
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Build hlo_runner_main
working-directory: xla
run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu
- name: Build test_gpu_profiler
working-directory: xla
run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:test_gpu_profiler
# - name: Wait For Connection
# uses: google-ml-infra/actions/ci_connection@main
# with:
# halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Create isolated_convolution.hlo
working-directory: xla
run: |
cat << EOF > isolated_convolution.hlo
HloModule convolution.167:
ENTRY %convolution.167 (parameter.0: f32[64,28,28,128]{2,1,0,3}, parameter.1: f32[3,3,128,128]{2,1,0,3}) -> f32[64,28,28,128]{2,1,0,3} {
%parameter.0 = f32[64,28,28,128]{2,1,0,3} parameter(0)
%parameter.1 = f32[3,3,128,128]{2,1,0,3} parameter(1)
ROOT %convolution.167 = f32[64,28,28,128]{2,1,0,3} convolution(f32[64,28,28,128]{2,1,0,3} %parameter.0, f32[3,3,128,128]{2,1,0,3} %parameter.1), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_01oi->bf01
}
EOF
# - name: Wait For Connection
# uses: google-ml-infra/actions/ci_connection@main
# with:
# halt-dispatch-input: ${{ inputs.halt-for-connection }}
# - name: Run specific HLO file
# working-directory: xla
# run: |
# nvidia-smi --query-gpu=utilization.gpu --format=csv -l 1 > gpu_utilization.log & ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu --device_type=gpu --log_output=True --use_spmd_partitioning gemm_00881937d6d49056045c3325a12b108b.hlo &> results/gemm_00881937d6d49056045c3325a12b108b.log
# - name: Wait For Connection
# uses: google-ml-infra/actions/ci_connection@main
# with:
# halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Run specific HLO file
working-directory: xla
run: |
nvidia-smi --query-gpu=utilization.gpu --format=csv -l 1 > results/gpu_utilization.log & ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu --device_type=gpu --log_output=True --use_spmd_partitioning isolated_convolution.hlo &> results/isolated_convolution.hlo.log
- name: Run test_gpu_profiler
working-directory: xla
run: |
nvidia-smi --query-gpu=utilization.gpu --format=csv -l 1 > results/gpu_utilization_profiler.log & ./bazel-bin/xla/tools/multihost_hlo_runner/test_gpu_profiler &> results/test_gpu_profiler.log
- name: Wait For Connection
uses: google-ml-infra/actions/ci_connection@main
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Download parse_xla_logs.py
working-directory: xla
run: wget https://raw.githubusercontent.com/juliagmt-google/xla/main/.github/workflows/parse_xla_logs.py
- name: Parse XLA logs
working-directory: xla
run: python parse_xla_logs.py results/isolated_convolution.hlo.log
- name: Upload Results
uses: actions/upload-artifact@v4
with:
name: gpu-xla-benchmarks
path: xla/results
# # jax-build-and-test:
# # runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
# # container:
# # image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
# # env:
# # JAXCI_HERMETIC_PYTHON_VERSION: 3.11
# # steps:
# # - name: Checkout JAX Fork
# # uses: actions/checkout@v3
# # with:
# # repository: 'google-ml-infra/jax-fork'
# # path: jax-fork
# # - name: Install JAX Dependencies
# # working-directory: jax-fork
# # run: |
# # python -m pip install --upgrade pip
# # pip install pytest
# # pip install absl-py
# # pip install "jax[cuda12_pip]" # Adjust CUDA version if needed
# # pip install google-benchmark
# # - name: Run JAX Multiprocess GPU Test
# # working-directory: jax-fork
# # continue-on-error: true
# # run: python -m pytest tests/multiprocess_gpu_test.py
# # - name: Run HLO Module Benchmarks withg GPU in xla/tests/fuzz
# # working-directory: xla
# # continue-on-error: true
# # run: |
# # for file in xla/tests/fuzz/*.hlo; do
# # filename=$(basename "$file")
# # # Skip expected failed hlo files.
# # if [[ "$filename" == "rand_000060.hlo" || "$filename" == "rand_000067.hlo" || "$filename" == "rand_000072.hlo" ]]; then
# # echo "Skipping benchmark on $file"
# # continue
# # fi
# # echo "Running benchmark on $file" &> results/"$filename".log
# # ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --use_spmd_partitioning "$file" &> results/"$filename".log
# # done
# # - name: Upload Results
# # uses: actions/upload-artifact@v4
# # with:
# # name: gpu-xla-benchmarks
# # path: xla/results