test(nightly): distributed SFT functionality suite

.github/workflows/regression.yml

@@ -2,7 +2,7 @@ name: Regression
 
 # Scheduled regression suites, one job group per suite. Dispatch with `suite`
 # to run one of them.
-#   parallel: test_parallel_invariance.py on the 2xH200 runners
+#   parallel: distributed/test_parallel_invariance.py on the 2xH200 runners
 #   eval:     every configs/eval/*.toml through `uv run eval` against Prime
 #             Inference, compared against baselines/eval/. reports/eval/ holds
 #             the error bars of each baseline. Also runs on PRs that touch the
@@ -86,7 +86,7 @@ jobs:
                   WANDB_MODE: offline
                   HF_TOKEN: ${{ secrets.HF_TOKEN }}
                   PYTEST_OUTPUT_DIR: /tmp/outputs
-              run: uv run pytest -vvs tests/regression/test_parallel_invariance.py
+              run: uv run pytest -vvs tests/regression/distributed/test_parallel_invariance.py
             - name: Upload run logs on failure
               if: failure()
               uses: actions/upload-artifact@v7

.gitignore

@@ -4,6 +4,7 @@
 .codex/
 *.ipynb
 *.jsonl
+!tests/regression/configs/distributed/sft/data/train.jsonl
 configs/endpoints.py
 
 # Run artifacts

docs/development.md

@@ -24,7 +24,7 @@ The test suite is split into three tiers.
 
 - **`tests/unit/`** — fast-running, hermetic tests for isolated logic: config parsing and validation, advantage / loss / scheduler / packer math, individual dataset paths, model-conversion roundtrips, etc. Tests that need a GPU are tagged with the `gpu` marker.
 - **`tests/integration/`** — full-stack RL/SFT runs on a tiny model end-to-end through inference + orchestrator + trainer.
-- **`tests/regression/`** — nightly regression suites: `test_parallel_invariance.py`, which checks that DP/CP/EP layouts train the same first step as a single GPU, and evals from `configs/eval/` compared against `baselines/eval/`.
+- **`tests/regression/`** — nightly regression suites: `distributed/test_parallel_invariance.py`, which checks that DP/CP/EP layouts train the same first step as a single GPU, and evals from `configs/eval/` compared against `baselines/eval/`.
 
 ### Running Tests Locally
 
@@ -43,7 +43,7 @@ uv run pytest tests/integration/test_reverse_text.py -vvs  # one specific scenar
 |---|---|---|---|
 | [`tests.yml`](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/.github/workflows/tests.yml) | CPU jobs: every PR + push to `main`. `gsm8k_eval`: relevant changes in non-draft PRs, every push to `main`, manual dispatch | CPU: `pytest tests/unit -m "not gpu"`, plus a slim-wheel install check that `prime-rl-configs` imports cleanly without heavy deps (no torch / vllm / transformers / wandb / verifiers / datasets / liger / loguru in `sys.modules`). Separate hosted `gsm8k_eval` job | `ubuntu-latest` |
 | [`gpu.yml`](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/.github/workflows/gpu.yml) | On a PR: applying the `gpu-ci` label (Triage or higher). Also every merge to `main` and manual dispatch | Full GPU suite: `pytest tests/unit -m gpu`, plus RL/SFT integration scenarios (`reverse_text`, `reverse_text_sft`, `reverse_text_sft_lora`, `reverse_text_lora`, `reverse_text_moe`, `reverse_text_rl_opd`, `reverse_text_rl_sft`, `alphabet_sort`) with `max-parallel: 5` | Self-hosted `prime-rl-2xh200` k8s scale set (2xH200); `ubuntu-latest` for the trigger check |
-| [`regression.yml`](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/.github/workflows/regression.yml) | 03:00 PST daily + manual `workflow_dispatch` (`suite` picks one suite); `eval` also on PRs that touch it | `parallel`: `tests/regression/test_parallel_invariance.py` (layouts that need more than 2 GPUs skip). `eval`: `tests/regression/configs/eval/*.toml` against Prime Inference, compared against `tests/regression/baselines/eval/` | `eval`: `ubuntu-latest`. `parallel`: `prime-rl-2xh200` |
+| [`regression.yml`](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/.github/workflows/regression.yml) | 03:00 PST daily + manual `workflow_dispatch` (`suite` picks one suite); `eval` also on PRs that touch it | `parallel`: `tests/regression/distributed/test_parallel_invariance.py` (layouts that need more than 2 GPUs skip). `eval`: `tests/regression/configs/eval/*.toml` against Prime Inference, compared against `tests/regression/baselines/eval/` | `eval`: `ubuntu-latest`. `parallel`: `prime-rl-2xh200` |
 
 Pushes to a PR do not run GPU jobs. To run the GPU suite on a PR's current commit, apply the `gpu-ci` label (drafts included). The workflow removes the label when it starts, so re-apply it to run again after new pushes or to retry. A new trigger cancels the PR's queued or in-flight GPU run. Fork PRs cannot run GPU CI: a maintainer can push the branch to this repo and label that PR. A PR with merge conflicts cannot run it until the conflicts are resolved.
 

scripts/ci/run_distributed_nightly.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+# Run the nightly distributed suite on one node with at least 4 GPUs. Extra args go to pytest, e.g. -k qwen3_moe or
+# --shard 2/4.
+set -euo pipefail
+
+cd "$(dirname "$0")/../.."
+uv run pytest tests/regression/distributed --durations=0 -rfEs "$@"

tests/conftest.py

@@ -94,12 +94,13 @@ def _get_wandb_project(wandb_project: str) -> str:
 class ProcessResult:
     """Result object containing process information and captured output."""
 
-    def __init__(self, process: subprocess.Popen):
+    def __init__(self, process: subprocess.Popen, timed_out: bool = False):
         self.returncode = process.returncode
         self.pid = process.pid
+        self.timed_out = timed_out
 
     def __repr__(self):
-        return f"ProcessResult(returncode={self.returncode}, pid={self.pid})"
+        return f"ProcessResult(returncode={self.returncode}, pid={self.pid}, timed_out={self.timed_out})"
 
 
 @pytest.fixture(scope="module")
@@ -109,16 +110,18 @@ def run_process() -> Callable[[Command, Environment, int], ProcessResult]:
     def _run_process(command: Command, env: Environment = {}, timeout: int | None = None) -> ProcessResult:
         """Run a subprocess with given command and environment with a timeout"""
         process = subprocess.Popen(command, env={**os.environ, **env})
+        timed_out = False
         try:
             process.wait(timeout=timeout)
         except subprocess.TimeoutExpired:
+            timed_out = True
             cleanup_process(process.pid, signal.SIGTERM)
             try:
                 process.wait(timeout=30)
             except subprocess.TimeoutExpired:
                 cleanup_process(process.pid, signal.SIGKILL)
                 process.wait()
 
-        return ProcessResult(process)
+        return ProcessResult(process, timed_out)
 
     return _run_process

tests/regression/configs/distributed/sft/base.toml

@@ -0,0 +1,5 @@
+[data]
+name = "tests/regression/configs/distributed/sft/data"
+batch_size = 16
+micro_batch_size = 1
+seq_len = 4096