Skip to content

Add support for reducing across the middle dimension for 3D matrices using the sum Triton kernel #2300

Add support for reducing across the middle dimension for 3D matrices using the sum Triton kernel

Add support for reducing across the middle dimension for 3D matrices using the sum Triton kernel #2300

Workflow file for this run

name: TorchBench PR Test on A10G
on:
pull_request:
workflow_dispatch:
push:
branches:
- main
env:
CONDA_ENV: "torchbench"
DOCKER_IMAGE: "ghcr.io/pytorch/torchbench:latest"
SETUP_SCRIPT: "/workspace/setup_instance.sh"
HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
jobs:
pr-test:
# AWS A10G GPU instance label: linux.g5.4xlarge.nvidia.gpu
# OS version: Amazon Linux 2
runs-on: [self-hosted, linux.g5.4xlarge.nvidia.gpu]
timeout-minutes: 1440 # 24 hours
environment: docker-s3-upload
steps:
- name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
uses: pytorch/test-infra/.github/actions/setup-ssh@main
with:
github-secret: ${{ secrets.TORCHBENCH_ACCESS_TOKEN }}
- name: Checkout TorchBench
uses: actions/checkout@v3
with:
path: benchmark
- name: Pull docker image
uses: pytorch/test-infra/.github/actions/pull-docker-image@main
with:
docker-image: ${{ env.DOCKER_IMAGE }}
- name: Install NVIDIA Driver, docker runtime, set GPU_FLAG
id: install-nvidia-driver
uses: pytorch/test-infra/.github/actions/setup-nvidia@main
- name: Install and Test TorchBench
run: |
container_name=$(docker run \
-e CONDA_ENV="${CONDA_ENV}" \
-e SETUP_SCRIPT="${SETUP_SCRIPT}" \
-e HUGGING_FACE_HUB_TOKEN="${HUGGING_FACE_HUB_TOKEN}" \
--tty \
--detach \
--shm-size=32gb \
-v "${PWD}/benchmark:/benchmark" \
--gpus all \
-w / \
"${{ env.DOCKER_IMAGE }}" \
tail -f /dev/null
)
echo "Container name: ${container_name}"
docker exec -t -w "/" "${container_name}" bash -c "sudo chown -R runner /benchmark; sudo chgrp -R runner /benchmark"
docker exec -t -w "/benchmark" "${container_name}" bash /benchmark/scripts/torchbench_install.sh
docker exec -t -w "/benchmark" "${container_name}" bash /benchmark/scripts/torchbench_test.sh
- name: Teardown Linux
uses: pytorch/test-infra/.github/actions/teardown-linux@main
if: always()
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
cancel-in-progress: true