integration_test_8gpu_exp_models.yaml

name: 8 GPU Experimental Model Tests

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

concurrency:
  group: unit-test${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
  cancel-in-progress: true

defaults:
  run:
    shell: bash -l -eo pipefail {0}

jobs:
  build-test:
    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
    with:
      runner: linux.g5.48xlarge.nvidia.gpu
      gpu-arch-type: cuda
      gpu-arch-version: "12.6"
      # This image is faster to clone than the default, but it lacks CC needed by triton
      # (1m25s vs 2m37s).
      docker-image: torchtitan-ubuntu-20.04-clang12
      repository: pytorch/torchtitan
      upload-artifact: outputs
      script: |
        set -eux

        # The generic Linux job chooses to use base env, not the one setup by the image
        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
        conda activate "${CONDA_ENV}"

        # Log CUDA driver version for debugging.
        DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 || true)
        echo "CUDA driver version: ${DRIVER_VERSION}"

        pip config --user set global.progress_bar off

        python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126

        USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126

        mkdir artifacts-to-be-uploaded
        python -m tests.integration_tests.run_tests --test_suite experimental_models artifacts-to-be-uploaded --ngpu 8