GPU tests #29
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: | |
| workflow_dispatch: | |
| inputs: | |
| instance_type: | |
| description: 'EC2 instance type' | |
| required: false | |
| type: choice | |
| default: 'g6.2xlarge' | |
| options: | |
| - g4ad.xlarge # 4 vCPUs, 16GB RAM, AMD V520 GPU, ≈$0.38/hr | |
| - g4ad.2xlarge # 8 vCPUs, 32GB RAM, AMD V520 GPU, ≈$0.55/hr | |
| - g5.xlarge # 4 vCPUs, 16GB RAM, A10G GPU, ≈$1.11/hr | |
| - g5.2xlarge # 8 vCPUs, 32GB RAM, A10G GPU, ≈$1.33/hr | |
| - g5.4xlarge # 16 vCPUs, 64GB RAM, A10G GPU, ≈$1.79/hr | |
| - g6.xlarge # 4 vCPUs, 16GB RAM, L4 GPU, ≈$0.89/hr | |
| - g6.2xlarge # 8 vCPUs, 32GB RAM, L4 GPU, ≈$1.08/hr | |
| - g6.4xlarge # 16 vCPUs, 64GB RAM, L4 GPU, ≈$1.46/hr | |
| ec2_image_id: | |
| description: 'AMI ID (leave empty for auto-selection based on instance type)' | |
| required: false | |
| type: string | |
| workflow_call: | |
| inputs: | |
| instance_type: | |
| description: 'EC2 instance type' | |
| required: true | |
| type: string | |
| ec2_image_id: | |
| description: 'AMI ID (leave empty for auto-selection based on instance type)' | |
| required: false | |
| type: string | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| ec2: | |
| name: Start EC2 runner | |
| uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@rw/hooks | |
| with: | |
| action_ref: rw/hooks | |
| ec2_instance_type: ${{ inputs.instance_type || 'g6.2xlarge' }} | |
| # AMI selection: Use provided AMI or default based on instance type | |
| # For g4ad (AMD GPU): Ubuntu 22.04 base (ROCm must be installed) | |
| # For g5/g6 (NVIDIA GPU): Deep Learning AMI with CUDA/PyTorch pre-installed | |
| ec2_image_id: ${{ inputs.ec2_image_id || (startsWith(inputs.instance_type || 'g6.2xlarge', 'g4ad') && 'ami-021589336d307b577' || 'ami-0aee7b90d684e107d') }} | |
| # ami-021589336d307b577: Ubuntu 22.04 LTS (20250801) - for AMD GPU instances | |
| # 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: | | |
| INSTANCE_TYPE="${{ inputs.instance_type || 'g6.2xlarge' }}" | |
| if [[ "$INSTANCE_TYPE" == g4ad.* ]]; then | |
| # For AMD GPU instances, we need to set up Python manually | |
| echo "Setting up Python for AMD GPU instance..." | |
| sudo apt-get update | |
| sudo apt-get install -y python3-pip python3-venv | |
| python3 -m venv /opt/venv | |
| echo "/opt/venv/bin" >> $GITHUB_PATH | |
| else | |
| # Use the DLAMI's pre-installed PyTorch conda environment | |
| echo "/opt/conda/envs/pytorch/bin" >> $GITHUB_PATH | |
| echo "CONDA_DEFAULT_ENV=pytorch" >> $GITHUB_ENV | |
| fi | |
| - name: Check GPU | |
| run: | | |
| INSTANCE_TYPE="${{ inputs.instance_type || 'g6.2xlarge' }}" | |
| if [[ "$INSTANCE_TYPE" == g4ad.* ]]; then | |
| echo "AMD GPU instance detected - skipping nvidia-smi" | |
| # Could add rocm-smi or other AMD GPU checks here if ROCm is installed | |
| else | |
| nvidia-smi | |
| fi | |
| - 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 || 'g6.2xlarge' }}" | |
| # Set GPU architecture based on instance type | |
| if [[ "$INSTANCE_TYPE" == g4ad.* ]]; then | |
| # AMD V520 GPU - ROCm support required | |
| echo "Setting up for AMD GPU (V520)..." | |
| echo "Installing ROCm and PyTorch with ROCm support..." | |
| # Install ROCm (version 6.0+) | |
| wget https://repo.radeon.com/amdgpu-install/24.04.1/ubuntu/jammy/amdgpu-install_24.04.1-1.61_all.deb | |
| sudo apt-get install -y ./amdgpu-install_24.04.1-1.61_all.deb | |
| sudo amdgpu-install --usecase=rocm --no-dkms -y | |
| # Add user to render and video groups for GPU access | |
| sudo usermod -a -G render,video $USER | |
| # Install PyTorch with ROCm support | |
| pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.0 | |
| # Set ROCm environment variables | |
| export ROCM_PATH=/opt/rocm | |
| export PATH=$ROCM_PATH/bin:$PATH | |
| export LD_LIBRARY_PATH=$ROCM_PATH/lib:$LD_LIBRARY_PATH | |
| export HIP_ARCHITECTURES=gfx906 # V520 GPU architecture | |
| echo "WARNING: ROCm support for Mamba-SSM is experimental." | |
| echo "The build may fail or tests may not pass on AMD GPUs." | |
| elif [[ "$INSTANCE_TYPE" == g5.* ]]; then | |
| # TORCH_CUDA_ARCH_LIST tells PyTorch which specific architecture to compile for | |
| export TORCH_CUDA_ARCH_LIST="8.6" # A10G GPU | |
| export CUDA_VISIBLE_DEVICES=0 | |
| export NVCC_GENCODE="-gencode arch=compute_86,code=sm_86" | |
| elif [[ "$INSTANCE_TYPE" == g6.* ]]; then | |
| export TORCH_CUDA_ARCH_LIST="8.9" # L4 GPU (Ada Lovelace) | |
| export CUDA_VISIBLE_DEVICES=0 | |
| export NVCC_GENCODE="-gencode arch=compute_89,code=sm_89" | |
| 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 -vs --maxfail=10 tests/ |