test_fx_automatic_parallelism.yml

name: FX Automatic Parallelism on GPU / Python - Test

on:
  workflow_dispatch:
  # TODO: Uncomment when fixed
  # push:
  #   branches:
  #     - main
  #   paths:
  #     - 'optimum/fx/parallelization/**.py'
  # pull_request:
  #   branches:
  #     - main
  #   paths:
  #     - 'optimum/fx/parallelization/**.py'

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

env:
  TRANSFORMERS_IS_CI: true

jobs:
  run_gpu_tests:
    runs-on:
      group: aws-g5-12xlarge-plus

    container:
      image: nvidia/cuda:12.4.1-devel-ubuntu22.04
      options: --mount type=tmpfs,destination=/tmp --shm-size 64gb --gpus all --ipc host -v /mnt/hf_cache:/mnt/cache/
      env:
        NCCL_DEBUG: INFO
    defaults:
      run:
        shell: bash

    steps:
      - name: Checkout optimum
        uses: actions/checkout@v6
        with:
          fetch-depth: 1

      - uses: actions/setup-python@v6
        with:
          python-version: "3.10"

      - name: Run nvidia-smi
        run: |
          nvidia-smi

      - name: Install dependencies
        run: |
          pip install -U pip
          pip install .[tests]

      - name: Run automatic model parallelism tests
        run: |
          pytest tests/fx/parallelization -s -v -o log_cli=true