ci(github): add GPU unittest workflow

.github/workflows/unittest.yml

@@ -0,0 +1,99 @@
+name: GPU Unit Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: gpu-unittest-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unittest:
+    name: Unit Tests (H20, 4 GPUs)
+    runs-on: [self-hosted, linux, x64, h20, gpu]
+    timeout-minutes: 30
+    env:
+      CI_IMAGE: ghcr.io/redai-studio/relaxrl:latest
+      CI_CONTAINER_NAME: relax-unittest-${{ github.run_id }}-${{ github.run_attempt }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v7
+        with:
+          persist-credentials: false
+
+      - name: Create test container
+        env:
+          CUDA_VISIBLE_DEVICES: 0,1,2,3
+          HOST_UV_CACHE_DIR: /home/relax-ci/.cache/uv
+          UV_CACHE_DIR: /uv-cache
+          UV_LINK_MODE: copy
+          ARTIFACT_DIR: ${{ runner.temp }}/unittest-${{ github.run_id }}-${{ github.run_attempt }}
+        run: |
+          set -euo pipefail
+          : "${CI_GPU_DEVICES:?Set CI_GPU_DEVICES to the four GPUs assigned to this runner}"
+          IFS=',' read -r -a gpu_devices <<< "$CI_GPU_DEVICES"
+          if [[ ${#gpu_devices[@]} -ne 4 ]]; then
+            echo "CI_GPU_DEVICES must list exactly four GPU indices or UUIDs" >&2
+            exit 1
+          fi
+          mkdir -p "$HOST_UV_CACHE_DIR" "$ARTIFACT_DIR"
+
+          # Keep network, IPC and /tmp private for parallel Ray/NCCL tests.
+          # Create before starting so cancellation during a pull cannot start tests.
+          docker create --pull always --rm --init \
+            --name "$CI_CONTAINER_NAME" \
+            --gpus "\"device=$CI_GPU_DEVICES\"" \
+            --shm-size=32g \
+            -v "$GITHUB_WORKSPACE:/source:ro" \
+            -v "$ARTIFACT_DIR:/artifacts" \
+            -v "$HOST_UV_CACHE_DIR:$UV_CACHE_DIR" \
+            --workdir /workspace/Relax \
+            --env CUDA_VISIBLE_DEVICES \
+            --env UV_CACHE_DIR \
+            --env UV_LINK_MODE \
+            --env http_proxy \
+            --env https_proxy \
+            --env no_proxy \
+            --entrypoint /bin/bash "$CI_IMAGE" -c 'sleep infinity'
+          docker inspect --format 'Test image: {{.Config.Image}} ({{.Image}})' "$CI_CONTAINER_NAME"
+          docker start "$CI_CONTAINER_NAME"
+
+      - name: Prepare workspace and show GPUs
+        run: |
+          docker exec "$CI_CONTAINER_NAME" bash -euc '
+            cp -a --no-preserve=ownership /source/. /workspace/Relax/
+            nvidia-smi
+          '
+
+      - name: Install dependencies
+        run: |
+          docker exec "$CI_CONTAINER_NAME" uv pip install \
+            --system --break-system-packages -r requirements.txt \
+            --extra-index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
+
+      - name: Run unittest
+        run: |
+          docker exec "$CI_CONTAINER_NAME" bash -euo pipefail -c '
+            python -m pytest tests/ -v -ra --tb=short --junitxml=/artifacts/junit.xml \
+              2>&1 | tee /artifacts/pytest.log
+          '
+
+      - name: Remove test container
+        if: always()
+        run: docker rm --force "$CI_CONTAINER_NAME" >/dev/null 2>&1 || true
+
+      - name: Upload test results
+        if: always()
+        uses: actions/upload-artifact@v7
+        with:
+          name: unittest-${{ github.run_id }}-${{ github.run_attempt }}
+          path: ${{ runner.temp }}/unittest-${{ github.run_id }}-${{ github.run_attempt }}/
+          if-no-files-found: warn
+          retention-days: 14

tests/backends/megatron/test_chunked_mtp_loss.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import types
+from collections.abc import Iterator
 
 import pytest
 import torch
@@ -35,6 +36,18 @@
 S, B, H, V, LAYERS = 16, 2, 8, 32, 2
 
 
+@pytest.fixture(scope="module", autouse=True)
+def _single_threaded_cpu() -> Iterator[None]:
+    """Run bitwise CPU comparisons with one thread, then restore the previous
+    setting."""
+    previous_num_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    try:
+        yield
+    finally:
+        torch.set_num_threads(previous_num_threads)
+
+
 class _FakeHead(nn.Module):
     """Stand-in for ColumnParallelLinear output_layer.