Skip to content
CI / GPU

Fix ONNX export of channels_last_3d models; fail fast on requested ex… #328

Sign in to view logs

Fix ONNX export of channels_last_3d models; fail fast on requested ex…

Fix ONNX export of channels_last_3d models; fail fast on requested ex… #328

Workflow file for this run

.github/workflows/ci-gpu.yml at a7ba11c
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# GPU tests live in a separate workflow because NVIDIA self-hosted runners
# block pull_request events entirely. Keeping them here avoids a confusing
# "Skipped" entry with unresolved matrix names on every PR.
name: CI / GPU
on:
workflow_dispatch:
push:
branches:
- main
- "pull-request/[0-9]+"
merge_group:
types:
- checks_requested
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
env:
PIP_NO_CACHE_DIR: "1"
PIP_DISABLE_PIP_VERSION_CHECK: "1"
PIP_PREFER_BINARY: "1"
jobs:
gpu-tests:
runs-on: linux-amd64-gpu-rtxpro6000-latest-1
# contents: write gives the token push access, which is required to see the
# permanent-draft release that fetch-models downloads the weights from.
permissions:
contents: write
container:
image: ubuntu:24.04
options: -u root --security-opt seccomp=unconfined --shm-size 16g
env:
NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
timeout-minutes: 45
strategy:
fail-fast: false
matrix:
python-version: ["3.11", "3.12", "3.13"]
# cu128 = representative CUDA 12.x wheel; cu130 = CUDA 13.0 wheel.
torch-cuda: ["cu128", "cu130"]
name: "gpu / py${{ matrix.python-version }} / ${{ matrix.torch-cuda }}"
steps:
- name: Install system dependencies
run: |
export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get install -y git gcc
- uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- uses: actions/checkout@v4
- name: Fetch pre-trained models
uses: ./.github/actions/fetch-models
- name: Verify GPU
run: nvidia-smi
- name: Install dependencies and run tests
run: bash code/scripts/check_python_compat.sh
env:
PYTHON_BIN: python
MODE: train
SKIP_TESTS: "0"
REQUIRE_GPU: "1"
TORCH_CUDA: ${{ matrix.torch-cuda }}
VENV_DIR: .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}
REQ_FILE: code/requirements_public_gpu_${{ matrix.torch-cuda == 'cu130' && 'cu13' || 'cu12' }}.txt
- name: Training + inference with LER check
shell: bash
run: |
source .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}/bin/activate
bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_train.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
# 0.35: short run (16k samples, 4 epochs for stable LER across py versions)
python code/scripts/check_ler_from_log.py /tmp/ci_train.log --max-ler 0.35
env:
EXPERIMENT_NAME: ci_short
PREDECODER_TRAIN_SAMPLES: "16384"
PREDECODER_VAL_SAMPLES: "2048"
PREDECODER_TEST_SAMPLES: "2048"
PREDECODER_TRAIN_EPOCHS: "4"
- name: Training + inference with multi-worker DataLoader (num_workers=2)
shell: bash
run: |
source .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}/bin/activate
bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_multiworker.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
python code/scripts/check_ler_from_log.py /tmp/ci_multiworker.log --max-ler 0.35
env:
EXPERIMENT_NAME: ci_multiworker
PREDECODER_TRAIN_SAMPLES: "16384"
PREDECODER_VAL_SAMPLES: "2048"
PREDECODER_TEST_SAMPLES: "2048"
PREDECODER_TRAIN_EPOCHS: "4"
PREDECODER_INFERENCE_NUM_WORKERS: "2"
# ---------------------------------------------------------------------------
# Mid-tier (~5-10 min): extended training + inference with LER check.
# Runs only after merge to main (not on PR branches) to save GPU time.
# Single Python version — multi-version coverage is handled by gpu-tests.
# ---------------------------------------------------------------------------
mid-gpu-tests:
if: github.ref == 'refs/heads/main'
needs: gpu-tests
runs-on: linux-amd64-gpu-rtxpro6000-latest-1
container:
image: ubuntu:24.04
options: -u root --security-opt seccomp=unconfined --shm-size 16g
env:
NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
timeout-minutes: 40
steps:
- name: Install system dependencies
run: |
export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get install -y git gcc
- uses: actions/setup-python@v5
with:
python-version: "3.13"
- uses: actions/checkout@v4
- name: Verify GPU
run: nvidia-smi
- name: Install Python dependencies
run: |
python -m venv .venv_mid
. .venv_mid/bin/activate
python -m pip install --upgrade pip setuptools wheel
# TODO: matrix by CUDA major version [cu12, cu13]
pip install -r code/requirements_public_train-cu12.txt
- name: Mid-tier training + inference with LER check (32k train, 4 epochs)
shell: bash
run: |
. .venv_mid/bin/activate
bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_mid.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
# 0.2: mid-tier (32k/4 epochs); loosen if flaky
python code/scripts/check_ler_from_log.py /tmp/ci_mid.log --max-ler 0.2
env:
EXPERIMENT_NAME: ci_mid
PREDECODER_TRAIN_SAMPLES: "32768"
PREDECODER_VAL_SAMPLES: "4096"
PREDECODER_TEST_SAMPLES: "4096"
PREDECODER_TRAIN_EPOCHS: "4"
- name: HE compile tests (torch.compile + autotune on GPU)
run: |
. .venv_mid/bin/activate
PYTHONPATH=code python -m unittest discover -s code/tests/mid -p "test_*.py" -v
# ---------------------------------------------------------------------------
# Multi-GPU tests: validates NCCL, DDP gradient sync, and per-rank data
# generation across 2 GPUs.
#
# Runner requirement: a self-hosted runner with >=2 GPUs.
# NVIDIA GHA runners follow the naming pattern
# linux-amd64-gpu-<model>-latest-<gpu-count>
# so the 2-GPU variant of the existing rtxpro6000 runner would be:
# linux-amd64-gpu-rtxpro6000-latest-2
# Confirm this label with your runner pool before enabling; if no 2-GPU
# runner exists the job will queue indefinitely.
#
# Runs only after merge to main (not on PR branches) to conserve GPU quota.
# ---------------------------------------------------------------------------
multi-gpu-tests:
if: github.ref == 'refs/heads/main'
needs: gpu-tests
runs-on: linux-amd64-gpu-rtxpro6000-latest-2
container:
image: ubuntu:24.04
options: -u root --security-opt seccomp=unconfined --shm-size 16g
env:
NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
timeout-minutes: 35
steps:
- name: Install system dependencies
run: |
export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get install -y git python3 python3-pip python3-venv
- uses: actions/checkout@v4
- name: Verify 2 GPUs are visible
run: |
nvidia-smi
count=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
echo "GPU count: ${count}"
[ "${count}" -ge 2 ] || { echo "ERROR: expected >=2 GPUs, found ${count}"; exit 1; }
- name: Install Python dependencies
run: |
python3 -m venv .venv_multigpu
. .venv_multigpu/bin/activate
python -m pip install --upgrade pip setuptools wheel
# TODO: matrix by CUDA major version [cu12, cu13]
pip install -r code/requirements_public_train-cu12.txt
- name: Run multi-GPU unit tests
run: |
. .venv_multigpu/bin/activate
PYTHONPATH=code python -m unittest discover \
-s code/tests -p "test_multi_gpu.py" -v
- name: Multi-GPU smoke training (2 GPUs, DDP)
# smoke_run.sh hardcodes GPUS=1; call local_run.sh directly so we
# can pass GPUS=2 and exercise the torch.distributed.run path.
shell: bash
run: |
. .venv_multigpu/bin/activate
export PREDECODER_TIMING_RUN=1
export PREDECODER_DISABLE_SDR=1
export PREDECODER_LER_FINAL_ONLY=1
export PREDECODER_INFERENCE_NUM_SAMPLES=32
export PREDECODER_INFERENCE_LATENCY_SAMPLES=0
export PREDECODER_INFERENCE_MEAS_BASIS=both
export PREDECODER_INFERENCE_NUM_WORKERS=0
EXPERIMENT_NAME=ci_multi_gpu WORKFLOW=train GPUS=2 \
bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_train.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
EXPERIMENT_NAME=ci_multi_gpu WORKFLOW=inference GPUS=2 \
bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_infer.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
# [LER Validation] lines are emitted during training, not inference
python code/scripts/check_ler_from_log.py /tmp/ci_multigpu_train.log --max-ler 0.35
env:
PREDECODER_TRAIN_SAMPLES: "16384"
PREDECODER_VAL_SAMPLES: "2048"
PREDECODER_TEST_SAMPLES: "2048"
PREDECODER_TRAIN_EPOCHS: "4"
- name: Multi-GPU smoke training with parallel spacelike HE (2 GPUs, DDP)
# Additive coverage on top of the default-config multi-GPU smoke above.
# Forces data.use_compile=True + data.use_parallel_spacelike=True so the
# parallel + compiled spacelike HE path runs end-to-end under DDP on
# 2 GPUs. Failure modes specific to this combination (per-rank device
# pinning of the partition, torch.compile cache contention across
# ranks, deadlocks during the compiled inner loop) surface as a
# training crash here. The existing default-config step above is
# intentionally left untouched so we do not regress on coverage of
# the default path.
shell: bash
run: |
. .venv_multigpu/bin/activate
export PREDECODER_TIMING_RUN=1
export PREDECODER_DISABLE_SDR=1
export PREDECODER_LER_FINAL_ONLY=1
export PREDECODER_INFERENCE_NUM_SAMPLES=32
export PREDECODER_INFERENCE_LATENCY_SAMPLES=0
export PREDECODER_INFERENCE_MEAS_BASIS=both
export PREDECODER_INFERENCE_NUM_WORKERS=0
EXPERIMENT_NAME=ci_multi_gpu_he WORKFLOW=train GPUS=2 \
EXTRA_PARAMS="data.use_compile=True data.use_parallel_spacelike=True" \
bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_he_train.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
EXPERIMENT_NAME=ci_multi_gpu_he WORKFLOW=inference GPUS=2 \
EXTRA_PARAMS="data.use_compile=True data.use_parallel_spacelike=True" \
bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_he_infer.log
r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
python code/scripts/check_ler_from_log.py /tmp/ci_multigpu_he_train.log --max-ler 0.35
env:
PREDECODER_TRAIN_SAMPLES: "16384"
PREDECODER_VAL_SAMPLES: "2048"
PREDECODER_TEST_SAMPLES: "2048"
PREDECODER_TRAIN_EPOCHS: "4"
# ---------------------------------------------------------------------------
# GPU coverage: captures GPU-specific code paths missed by the CPU coverage job
# ---------------------------------------------------------------------------
gpu-coverage:
runs-on: linux-amd64-gpu-rtxpro6000-latest-1
# contents: write — see gpu-tests: needed to read the draft release.
permissions:
contents: write
container:
image: ubuntu:24.04
options: -u root --security-opt seccomp=unconfined --shm-size 16g
env:
NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
timeout-minutes: 20
steps:
- name: Install system dependencies
run: |
export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get install -y git python3 python3-pip python3-venv
- uses: actions/checkout@v4
- name: Fetch pre-trained models
uses: ./.github/actions/fetch-models
- name: Verify GPU
run: nvidia-smi
- name: Install Python dependencies
run: |
python3 -m venv .venv_gpu_cov
. .venv_gpu_cov/bin/activate
python -m pip install --upgrade pip setuptools wheel
# TODO: matrix by CUDA major version [cu12, cu13]
pip install -r code/requirements_public_train-cu12.txt
pip install -r code/requirements_ci.txt
- name: Run tests with GPU coverage
run: |
. .venv_gpu_cov/bin/activate
PYTHONPATH=code coverage run -m unittest discover -s code/tests -p "test_*.py"
coverage report
coverage html -d htmlcov-gpu
coverage xml -o coverage-gpu.xml
- name: Upload GPU coverage artifacts
if: always()
uses: actions/upload-artifact@v4
with:
name: gpu-coverage-report
path: |
htmlcov-gpu/
coverage-gpu.xml