GPU tests #23
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: GPU tests | |
| on: | |
| push: | |
| branches: [main] | |
| pull_request: | |
| branches: [main] | |
| workflow_dispatch: | |
| inputs: | |
| instance_type: | |
| description: 'EC2 instance type' | |
| required: false | |
| type: choice | |
| default: 'g5.2xlarge' | |
| options: | |
| - g4dn.xlarge # 4 vCPUs, 16GB RAM, T4 GPU | |
| - g4dn.2xlarge # 8 vCPUs, 32GB RAM, T4 GPU | |
| - g4dn.4xlarge # 16 vCPUs, 64GB RAM, T4 GPU | |
| - g4dn.8xlarge # 32 vCPUs, 128GB RAM, T4 GPU | |
| - g5.xlarge # 4 vCPUs, 16GB RAM, A10G GPU | |
| - g5.2xlarge # 8 vCPUs, 32GB RAM, A10G GPU | |
| - g5.4xlarge # 16 vCPUs, 64GB RAM, A10G GPU | |
| - g5.8xlarge # 32 vCPUs, 192GB RAM, A10G GPU | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| ec2: | |
| name: Start EC2 runner | |
| uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2 | |
| with: | |
| ec2_instance_type: ${{ inputs.instance_type || 'g4dn.xlarge' }} | |
| ec2_image_id: ami-0aee7b90d684e107d # Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.4.1 (Ubuntu 22.04) 20250623 | |
| secrets: | |
| GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }} | |
| test: | |
| name: GPU tests | |
| needs: ec2 | |
| runs-on: ${{ needs.ec2.outputs.id }} | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Setup Python environment | |
| run: | | |
| # Use the DLAMI's pre-installed PyTorch conda environment | |
| echo "/opt/conda/envs/pytorch/bin" >> $GITHUB_PATH | |
| echo "CONDA_DEFAULT_ENV=pytorch" >> $GITHUB_ENV | |
| - name: Check GPU and PyTorch | |
| run: | | |
| nvidia-smi | |
| echo "Timing PyTorch import..." | |
| time python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}, version: {torch.version.cuda}')" | |
| echo "Second import (should be cached):" | |
| time python -c "import torch; print('Imported torch')" | |
| - name: Install mamba-ssm and test dependencies | |
| run: | | |
| # Use all available CPUs for compilation (we're only building for 1 GPU arch) | |
| export MAX_JOBS=$(nproc) | |
| INSTANCE_TYPE="${{ inputs.instance_type || 'g4dn.xlarge' }}" | |
| # Set CUDA architecture based on GPU type | |
| # CRITICAL: Also set CUDA_VISIBLE_DEVICES to force PyTorch to detect the right arch | |
| if [[ "$INSTANCE_TYPE" == g4dn.* ]]; then | |
| export TORCH_CUDA_ARCH_LIST="7.5" # T4 GPU | |
| export CUDA_VISIBLE_DEVICES=0 | |
| # Force nvcc to only compile for our GPU | |
| export NVCC_GENCODE="-gencode arch=compute_75,code=sm_75" | |
| elif [[ "$INSTANCE_TYPE" == g5.* ]]; then | |
| export TORCH_CUDA_ARCH_LIST="8.6" # A10G GPU | |
| export CUDA_VISIBLE_DEVICES=0 | |
| export NVCC_GENCODE="-gencode arch=compute_86,code=sm_86" | |
| fi | |
| echo "Building with MAX_JOBS=$MAX_JOBS for $INSTANCE_TYPE" | |
| # Install mamba-ssm with causal-conv1d and dev dependencies | |
| # Note: causal-conv1d will download pre-built wheels when available | |
| pip install -v --no-build-isolation -e .[causal-conv1d,dev] | |
| - name: Run tests | |
| run: pytest -xvs tests/ |