GPU Test Sharding: Weighted Distribution Across H100 Nodes

Each GPU test shard in our CI uses 8 H100 GPUs. We run 12 shards in parallel. That’s 96 H100s burning money every time someone pushes a PR with GPU test changes. If one shard takes 45 minutes while the others finish in 10, you’re paying for 96 GPUs but only using 8 of them for the last 35 minutes. Distribution matters.

The problem with naive sharding

The simplest approach: take all GPU tests, divide by 12, assign each chunk to a shard. Round-robin or alphabetical. Done.

This fails immediately for ML tests because test duration varies by orders of magnitude. A test that loads a model checkpoint, runs forward+backward passes, and checks gradient norms takes 8 minutes. A test that checks a configuration parsing function takes 3 seconds. Put the first in shard 1 and the second in shard 2 and you’ve created a 100x duration imbalance.

With 96 H100s at ~$30/GPU-hour, a 35-minute tail shard costs about $1,400 in wasted compute. Per CI run. On a busy day we might trigger 10 GPU test runs. The naive approach wastes more compute in idle shards than most companies spend on their entire CI budget.

Weighted sharding

The fix: weight each test by its expected duration, then bin-pack tests into shards to minimize the maximum shard duration.

Every test has an entry in gpu_test_durations.json:

{
  "tests/gpu/test_distributed_training.py::test_ddp_gradient_sync": 487.2,
  "tests/gpu/test_model_loading.py::test_load_checkpoint": 312.5,
  "tests/gpu/test_attention.py::test_flash_attention_backward": 45.8,
  "tests/gpu/test_config.py::test_parse_training_config": 2.1
}

The calculate_shards step reads this file, sorts tests by weight (heaviest first), and assigns each test to the shard with the lowest current total — a greedy bin-packing algorithm. This isn’t optimal (that’s NP-hard), but it’s close enough. The heaviest tests get distributed first, preventing any single shard from accumulating all the big ones.

Result: shard durations cluster within 10-15% of each other instead of having a 10x spread. The tail shard finishes minutes after the others, not half an hour later.

Four execution modes

Not all GPU tests are equal. Some need GPUs, some don’t. Some can run in parallel, some can’t. Some corrupt global state just by existing. Four modes handle this:

The isolated mode is the interesting one. Some tests configure CUDA at the process level — setting the default device, initializing CUDA contexts, configuring NCCL communicators. These are process-global operations. Once you call torch.cuda.set_device(0), every subsequent CUDA operation in that process uses device 0. If another test expected device 1, it’s too late.

The per_file isolation means each test file runs in its own pytest process. No shared state. The cost is process startup overhead — importing PyTorch and initializing CUDA takes 5-10 seconds per process. For a test that takes 3 seconds to run, the initialization overhead doubles the wall time. But correctness beats speed.

Why GPU state is awful

Three facts about CUDA that make test isolation hard:

1. CUDA contexts are process-global. When you first touch a GPU in a process, CUDA initializes a context for that device. This context persists until the process exits. You can’t uninitialize it. You can’t start fresh. The C runtime doesn’t support it.

2. Model loading is slow. Loading a large model onto GPUs takes 30-60 seconds — reading weights from disk, transferring to GPU memory, initializing optimizer state. You want to share this across tests in the same file. Loading once and running 20 tests against it is 60 seconds + 20×3 seconds = 2 minutes. Loading per-test is 20×63 seconds = 21 minutes.

3. GPU memory is finite. An H100 has 80GB of HBM3 memory. A single large model might use 60GB. You can’t load two models simultaneously. Tests that use different models must run sequentially or in separate processes.

These three constraints interact to create the four modes. Tests that share a model and don’t corrupt global state: concurrent mode, same process, parallel. Tests that modify CUDA settings: isolated mode, separate processes. Tests that don’t need GPUs at all: cpu mode, skip the GPU queue entirely.

The CI chain

GPU testing is the most expensive part of CI, so it’s gated and staged:

find_changed
  → wait-for-label ("run-gpu-tests")
  → build_image
  → calculate_shards
  → gpu_tests (12 parallel shards)
  → gpu_tests_passed
  → tag image as 'main'

Label gating

The wait-for-label step pauses until someone adds the run-gpu-tests label to the PR. This is a manual gate — not every PR needs GPU testing, and each run burns 96 H100s.

Without the label, the entire GPU test pipeline doesn’t trigger. The PR gets CPU tests, linting, type checking — everything cheap. Only when a reviewer decides “this change affects GPU codepaths” do they add the label and unlock the expensive tests.

This saves approximately 96 GPU-hours per skipped run. On a day with 30 PRs where only 5 need GPU tests, that’s ~2,400 GPU-hours saved. At H100 spot prices, that’s real money.

Image dependency

The build_image step must complete before shards start. GPU tests run inside the training container — same image researchers use for actual training. This ensures the test environment matches production exactly. No “works in CI but fails in training” discrepancies.

If the image build fails, no GPU tests run. The chain is strict — build → shard → test → gate. Each step depends on the previous one’s success.

Weight computation

The weights in gpu_test_durations.json come from an offline analysis pipeline:

1. Download pytest logs from recent CI runs (last 2-4 weeks of GPU test results)
2. Parse PYTEST-DURATION blocks — pytest outputs per-test timing when run with --durations=0
3. Run linear regression on the timing data — handles variance across runs (some tests fluctuate based on GPU thermal state, memory fragmentation, etc.)
4. Output per-test weights — the predicted duration in seconds
5. Human review — someone eyeballs the output, checks for outliers, confirms it looks reasonable
6. Manual copy to gpu_test_durations.json

The linear regression smooths out noise. A test might take 45 seconds on a cold GPU and 38 seconds on a warm one. The regression gives you an expected value that accounts for this variance, which is better for bin-packing than raw max or mean.

The manual review step exists because automated statistical analysis of GPU test durations is exactly the kind of thing that looks correct 95% of the time and catastrophically wrong 5% of the time. A test that got refactored might show up with the old test’s duration. A new test has no history. A flaky test that times out at 300 seconds skews the regression.

Humans are still better at “does this list of numbers look reasonable” than any automated check we could write in less time than it takes to just look at the list.

Container config

Each GPU shard runs in a container with specific resource requirements:

containers:
  - name: gpu-test
    resources:
      limits:
        nvidia.com/gpu: 8
    env:
      - name: NVIDIA_VISIBLE_DEVICES
        value: "all"
    securityContext:
      capabilities:
        add: ["SYS_NICE"]
    volumeMounts:
      - name: shm
        mountPath: /dev/shm
      - name: ci-data
        mountPath: /mnt/ci
volumes:
  - name: shm
    emptyDir:
      medium: Memory
      sizeLimit: 64Gi

The noteworthy settings:

--shm-size=64gb (via emptyDir): Shared memory. PyTorch uses /dev/shm for IPC between DataLoader workers. The default is 64MB, which causes Bus error when loading large batches. 64GB is generous but it’s a tmpfs backed by RAM, and these nodes have 256GB. Better to over-provision than debug random bus errors at 2am.

SYS_NICE capability: Allows the process to set CPU scheduling priority. Some GPU tests use CPU-intensive data preprocessing in parallel with GPU computation. Without SYS_NICE, the CPU scheduler might starve the data loading threads, causing GPU idle time while waiting for the next batch.

/mnt/ci mount: Shared storage with test data, model checkpoints, and cached datasets. Tests that load pre-trained models read from here instead of downloading from blob storage on every run. This mount is the difference between a 30-second model load and a 5-minute download.

The economics

Back-of-envelope math for a single GPU test run:

At 8-10 GPU test runs per day, that’s $1,400-1,700/day just for GPU CI. The weighted sharding saves maybe 20-30% by reducing tail shard waste — roughly $400/day. The label gating saves another $170 per skipped run — easily $2,000-3,000/day.

This is why the sharding algorithm matters. This is why the label gate exists. This is why someone spent time building a linear regression pipeline to predict test durations. The cost of getting distribution wrong isn’t “CI is slow.” It’s “CI costs an extra $10,000 per week in wasted GPU time.”

When each shard is 8 H100s, minutes matter. When you run 12 shards in parallel, distribution matters more than individual test speed. Optimize the schedule, not the test.