Create a workflow to run benchmarks #122

Workflow file for this run

.github/workflows/benchmarks.yml at a7721cd

	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: Wait For Connection
	# uses: google-ml-infra/actions/ci_connection@main
	# with:
	# halt-dispatch-input: ${{ inputs.halt-for-connection }}

	- name: Print machine specs
	run: \|
	lscpu
	free -h # Memory information
	df -h # Disk space information
	uname -a # Kernel information

	- name: Set up Python 3.10 # Choose your desired Python version
	uses: actions/setup-python@v4
	with:
	python-version: '3.10'

	- name: Install pip, setuptools, and wheel
	run: python -m pip install --upgrade pip setuptools wheel

	- name: Create and activate virtual environment
	working-directory: xla
	run: \|
	python -m venv xla/venv # Create venv inside the xla directory
	source xla/venv/bin/activate # Activate the environment

	- name: Run setup.sh for E2E benchmarks flax_2b (within venv)
	working-directory: xla/xla/backends/cpu/benchmarks/e2e/gemma2/flax_2b
	run: \|
	source ../../../../../venv/bin/activate # Activate venv from this directory
	bash setup.sh

	- name: Run run.sh for E2E benchmarks flax_2b (within venv)
	working-directory: xla/xla/backends/cpu/benchmarks/e2e/gemma2/flax_2b
	run: \|
	source ../../../../../venv/bin/activate # Activate venv from this directory
	bash run.sh


	# - 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: 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

	# - name: Wait For Connection
	# uses: google-ml-infra/actions/ci_connection@main
	# with:
	# halt-dispatch-input: ${{ inputs.halt-for-connection }}

	- name: Create gpu_hlo_backend.hlo
	working-directory: xla
	run: \|
	cat << EOF > gpu_hlo_backend.hlo
	HloModule module
	// CHECK: is_scheduled=true

	ENTRY computation {
	p = f32[5000,6000]{1,0} parameter(0)
	e = f32[5000,6000]{1,0} sqrt(p)
	c = f32[6000,5000] transpose(p), dimensions={1,0}
	r = f32[300,20,5000] reshape(c)
	ROOT out = (f32[5000,6000], f32[300,20,5000]) tuple(e,r)
	}
	EOF

	- name: Wait For Connection
	uses: google-ml-infra/actions/ci_connection@main
	with:
	halt-dispatch-input: ${{ inputs.halt-for-connection }}

	# - name: Run an HLO file
	# working-directory: xla
	# run: \|
	# ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning gpu_hlo_backend.hlo &> results/gpu_hlo_backend.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/gpu_hlo_backend.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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Create a workflow to run benchmarks #122

Workflow file

Create a workflow to run benchmarks #122

Uh oh!

Jobs

Run details

Workflow file for this run