From 06bebcb65f2b2826d94f47ca0c2b36ea0ea80c67 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Tue, 2 Sep 2025 13:13:23 -0400
Subject: [PATCH] run cu128-2.8.0 e2e tests on B200 (#3126)

* run cu128-2.8.0 e2e tests on B200

* not an int :facepalm:

* fix yaml
---
 .github/workflows/tests.yml | 2 ++
 cicd/single_gpu.py          | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 59011ee77..337230d4a 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -303,6 +303,7 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             num_gpus: 1
+            gpu_type: "B200"
             axolotl_extras:
     steps:
       - name: Checkout
@@ -324,6 +325,7 @@ jobs:
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV
           echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
           echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal
diff --git a/cicd/single_gpu.py b/cicd/single_gpu.py
index 0e2922e90..5a06a34f0 100644
--- a/cicd/single_gpu.py
+++ b/cicd/single_gpu.py
@@ -57,7 +57,8 @@ VOLUME_CONFIG = {
 }
 
 N_GPUS = int(os.environ.get("N_GPUS", 1))
-GPU_CONFIG = f"L40S:{N_GPUS}"
+GPU_TYPE = os.environ.get("GPU_TYPE", "L40S")
+GPU_CONFIG = f"{GPU_TYPE}:{N_GPUS}"
 
 
 def run_cmd(cmd: str, run_folder: str):