Merge branch 'master' into peng-auto-ep-optimizer-fix-v1

Author: tohtana

.github/workflows/TEST_SELECTION.md

@@ -25,8 +25,9 @@ one config — see [Adding a workflow](#adding-a-new-workflow).
   the import graph from your changed files to the impacted tests, and writes the
   list to `ci/.test_selection/test_list.txt`.
 - The trusted controller validates that list, then fetches, installs, and tests
-  the exact candidate SHA inside a no-secret Modal Sandbox. `push` to `master`
-  and manual runs always run everything.
+  the exact candidate SHA inside a no-secret Modal Sandbox. The Sandbox runs
+  only for merge queue entries, `push` to `master` (always the full suite), and
+  manual runs — a plain PR event never spends Modal quota.
 - It is **fail-safe**: anything it can't reason about safely → run the *full* suite.
   It never silently runs *fewer* tests than reality.
 - Preview locally:
@@ -63,12 +64,12 @@ The design is a small, self-contained take on HuggingFace `transformers`'
 ### Job flow
 
 ```
-                pull_request_target / push / workflow_dispatch
+           pull_request_target / merge_group / push / workflow_dispatch
                                   │
                     ┌─────────────┴─────────────┐
                     │        collect-tests       │   (no secrets)
                     │  checkout exact base SHA   │
-                    │  public-fetch exact PR SHA │
+                    │ public-fetch exact head SHA│
                     │  self-test the fetcher     │
                     │  parse candidate as data   │
                     │   → mode (all|subset|none) │
@@ -83,7 +84,7 @@ The design is a small, self-contained take on HuggingFace `transformers`'
                     │  validate mode/list + SHA   │
                     │  create one L40S:2 Sandbox  │
                     │    no secrets or mounts     │
-                    │    fetch exact PR SHA       │
+                    │    fetch exact head SHA      │
                     │    install + run pytest     │
                     │  terminate/observe Sandbox  │
                     └────────────────────────────┘
@@ -96,6 +97,14 @@ The design is a small, self-contained take on HuggingFace `transformers`'
 - **`subset`** → `deploy` runs pytest on exactly the impacted files.
 - **`all`** → `deploy` runs the whole scope (`tests/unit/v1`).
 
+Independent of `mode`, `deploy` is also skipped on `pull_request_target` runs,
+so pushing to a PR never spends Modal quota — the `collect-tests` summary still
+previews what the queue will run. The Sandbox actually executes on
+`merge_group` (the merged tree, gating the merge), `push` to `master`, and
+`workflow_dispatch`. On `merge_group` the candidate is the merge-group commit
+in the base repository, diffed against the queue's base SHA, so the selection
+covers exactly what the entry would introduce.
+
 
 ## How a decision is made
 
@@ -280,6 +289,17 @@ context and the `deploy` controller authenticates to Modal. The trust boundary i
   mitigation, **not** the trust boundary. Exact trusted base code plus Sandbox
   isolation is the primary protection.
 
+**`merge_group` runs are a separate trust context.** GitHub runs a queued
+entry's workflows from the merge-group commit — the PR's merged tree — so the
+workflow YAML itself is candidate-controlled. The jobs restore the repo's
+invariant by resolving every GitHub-side checkout to the trusted base revision
+(`merge_group.base_sha`): the selector, the controller, and the accelerate
+launcher all come from master there, and the merged candidate enters only as
+validated git data (fetched by exact SHA into a separate root) or inside the
+no-secret Modal Sandbox. The residual exposure — a queued PR rewriting the
+workflow YAML itself to echo secrets — is inherent to GitHub's merge queue;
+review PRs that touch `.github/workflows/*` with that in mind.
+
 > **Consequence:** changes to `ci/*` (including `tests_fetcher.py` itself) take
 > effect under `pull_request_target` only after they're **merged**. A PR that
 > changes this launcher cannot prove its new PR-triggered end-to-end path by
@@ -339,4 +359,6 @@ are still publicly reachable. Do not replace this with a moving branch fallback.
 **Why didn't this PR exercise its new `pull_request_target` controller?**
 GitHub intentionally runs that event's workflow from the trusted base. The new
 controller becomes the trusted code only after merge; before then, rely on the
-focused static/unit evidence described in the security model.
+focused static/unit evidence described in the security model. (The merge-queue
+run *does* use the PR's merged controller, so `ci/*` changes are first
+exercised live there.)

.github/workflows/modal-accelerate.yml

@@ -12,6 +12,11 @@ name: modal-accelerate
 # collect-tests job to do the filtering for us so that the job can be skipped and satisfy the
 # Required status for PRs to pass.
 #
+# Modal GPUs are a limited quota, so the deploy job only runs for merge queue entries, pushes to
+# master, and manual runs. Plain PR events keep the job skipped (which still satisfies the
+# Required status); the merge queue re-runs it on the merged tree. On merge queue runs the
+# launcher (this checkout) stays on trusted master and the merged candidate is fetched as data,
+# so queued PR code never executes on this runner next to the Modal/HF tokens.
 
 
 on:
@@ -35,6 +40,14 @@ on:
     branches:
       - master
 
+  # Required checks must also trigger on merge queue entries, otherwise the
+  # merge fails waiting for a status that is never reported. This is the only
+  # PR-related event that spends Modal quota.
+  merge_group:
+    types: [checks_requested]
+    branches:
+      - master
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
@@ -55,6 +68,8 @@ jobs:
         with:
           lfs: true
 
+      # paths-filter v4.0.1+ covers merge_group natively (it diffs the event's
+      # base/head commits), so no extra inputs are needed here.
       - name: Filter changed files
         uses: dorny/paths-filter@v4
         id: filter
@@ -81,13 +96,35 @@ jobs:
       # and it too is then updated at https://github.com/deepspeedai/deepspeed/settings/secrets/actions
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
-    if: needs.collect-tests.outputs.deepspeed == 'true'
+    # Plain PR events keep this job skipped so they never spend Modal quota (a
+    # skipped dependent job still satisfies the Required status); the merge
+    # queue re-runs it on the merged tree.
+    if: needs.collect-tests.outputs.deepspeed == 'true' && github.event_name != 'pull_request_target'
     steps:
       - name: Checkout Repository
         uses: actions/checkout@v7
         with:
+          # Trusted base revision on merge queue runs (master otherwise), so
+          # ci/accelerate.py holding the Modal/HF tokens is never queued-PR code.
+          ref: ${{ github.event.merge_group.base_sha || github.sha }}
           lfs: true
 
+      - name: Fetch merge queue candidate tree
+        if: github.event_name == 'merge_group'
+        env:
+          HEAD_REPOSITORY: ${{ github.repository }}
+          HEAD_SHA: ${{ github.sha }}
+          BASE_REPOSITORY: ${{ github.repository }}
+          BASE_SHA: ${{ github.event.merge_group.base_sha }}
+          CANDIDATE_ROOT: ${{ runner.temp }}/deepspeed-candidate
+        run: |
+          python3 ci/torch_latest.py checkout-candidate \
+            --head-repository "$HEAD_REPOSITORY" \
+            --head-sha "$HEAD_SHA" \
+            --base-repository "$BASE_REPOSITORY" \
+            --base-sha "$BASE_SHA" \
+            --destination "$CANDIDATE_ROOT"
+
       - name: Install Python
         uses: actions/setup-python@v7
         with:
@@ -100,5 +137,9 @@ jobs:
           uv pip install --system modal
 
       - name: Run tests
+        env:
+          # Tree under test: the fetched merge-group candidate on queue runs,
+          # empty (this checkout) everywhere else.
+          DS_CI_CANDIDATE_ROOT: ${{ github.event_name == 'merge_group' && format('{0}/deepspeed-candidate', runner.temp) || '' }}
         run: |
           modal run -m ci.accelerate

.github/workflows/modal-torch-latest.yml

@@ -3,7 +3,10 @@ name: modal-torch-latest
 # This Required workflow selects tests on a no-secret GitHub runner, then tests
 # the exact candidate SHA inside a no-secret Modal Sandbox. Under
 # pull_request_target, every GitHub-side checkout and Python entrypoint comes
-# from the trusted base commit. See .github/workflows/TEST_SELECTION.md.
+# from the trusted base commit; merge queue runs resolve that trusted base from
+# merge_group.base_sha. The Sandbox only runs for merge queue entries, pushes to
+# master, and manual runs, so plain PR events never spend Modal quota. See
+# .github/workflows/TEST_SELECTION.md.
 
 on:
   workflow_dispatch:
@@ -42,6 +45,14 @@ on:
     branches:
       - master
 
+  # Required checks must also trigger on merge queue entries, otherwise the
+  # merge fails waiting for a status that is never reported. This is the only
+  # PR-related event that spends Modal quota.
+  merge_group:
+    types: [checks_requested]
+    branches:
+      - master
+
 permissions:
   contents: read
 
@@ -63,19 +74,21 @@ jobs:
       - name: Checkout trusted workflow code
         uses: actions/checkout@v7
         with:
-          ref: ${{ github.event.pull_request.base.sha || github.sha }}
+          # Trusted base revision for every event: the merge-group commit must
+          # never control the selector that runs on this GitHub runner.
+          ref: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha || github.sha }}
           fetch-depth: 1
           persist-credentials: false
           lfs: false
           submodules: false
 
       - name: Fetch candidate tree as data
-        if: github.event_name == 'pull_request_target'
+        if: github.event_name == 'pull_request_target' || github.event_name == 'merge_group'
         env:
-          HEAD_REPOSITORY: ${{ github.event.pull_request.head.repo.full_name }}
-          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
-          BASE_REPOSITORY: ${{ github.event.pull_request.base.repo.full_name }}
-          BASE_SHA: ${{ github.event.pull_request.base.sha }}
+          HEAD_REPOSITORY: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
+          BASE_REPOSITORY: ${{ github.event.pull_request.base.repo.full_name || github.repository }}
+          BASE_SHA: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha }}
           CANDIDATE_ROOT: ${{ runner.temp }}/deepspeed-candidate
         run: |
           python3 ci/torch_latest.py checkout-candidate \
@@ -94,7 +107,7 @@ jobs:
           EVENT_NAME: ${{ github.event_name }}
           CANDIDATE_ROOT: ${{ runner.temp }}/deepspeed-candidate
         run: |
-          if [ "$EVENT_NAME" = "pull_request_target" ]; then
+          if [ "$EVENT_NAME" = "pull_request_target" ] || [ "$EVENT_NAME" = "merge_group" ]; then
             python3 ci/tests_fetcher.py \
               --workflow modal-torch-latest \
               --repo-root "$CANDIDATE_ROOT" \
@@ -125,14 +138,17 @@ jobs:
   deploy:
     name: modal-torch-latest / DeepSpeedAI CI
     runs-on: ubuntu-latest
-    timeout-minutes: 75
+    timeout-minutes: 105
     permissions:
       contents: read
     needs: collect-tests
 
     # Preserve the Required check on failures and no-test selections. A selector
     # failure enters this job and fails explicitly; mode=none skips the Sandbox.
-    if: ${{ !cancelled() && (needs.collect-tests.result != 'success' || needs.collect-tests.outputs.mode != 'none') }}
+    # pull_request_target runs skip the Sandbox entirely (a skipped dependent
+    # job still satisfies the Required check), so the tests gate the merge queue
+    # entry instead of every PR update.
+    if: ${{ !cancelled() && (needs.collect-tests.result != 'success' || (needs.collect-tests.outputs.mode != 'none' && github.event_name != 'pull_request_target')) }}
     steps:
       - name: Fail if test selection failed
         if: needs.collect-tests.result != 'success'
@@ -141,7 +157,9 @@ jobs:
       - name: Checkout trusted controller code
         uses: actions/checkout@v7
         with:
-          ref: ${{ github.event.pull_request.base.sha || github.sha }}
+          # Trusted base revision for every event: the merge-group commit is
+          # only ever fetched by exact SHA inside the no-secret Sandbox.
+          ref: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha || github.sha }}
           fetch-depth: 1
           persist-credentials: false
           lfs: false

.github/workflows/nv-flash-attn.yml

@@ -5,7 +5,7 @@ on:
   pull_request:
     paths:
       - 'deepspeed/sequence/**'
-      - 'tests/unit/sequence_parallelism/**'
+      - 'tests/unit/v1/sequence_parallelism/**'
       - '.github/workflows/nv-flash-attn.yml'
   schedule:
     - cron: "0 0 * * *"
@@ -57,7 +57,7 @@ jobs:
         run: |
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           cd tests
-          python -m pytest --color=yes --durations=0 --verbose -rF unit/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
+          python -m pytest --color=yes --durations=0 --verbose -rF unit/v1/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
       - name: Open GitHub issue if nightly CI fails
         if: ${{ failure() && (github.event_name == 'schedule') }}
         uses: JasonEtco/create-an-issue@v2

AGENTS.md

@@ -25,6 +25,16 @@
 - Comments should explain **why**, not **what**. Describe the purpose and reasoning, not the mechanics that the code already shows.
 - New features must include corresponding tests and documentation updates.
 
+### Test discipline
+
+- Tests verify contracts, not implementations: a test is well-formed only if a different correct implementation of the same contract passes it.
+- Before writing a test, name the concrete incorrect behavior it would catch; if you cannot name one, do not write it.
+- Assert on observable outcomes through public/stable interfaces; do not assert private method return values or exact internal strings unless pinning a specific fixed bug (justify in a comment).
+- Anchor to an external oracle or an independently derived reference instead of re-implementing the logic under test.
+- Mocks must stand in for a collaborator's documented contract (schema, protocol), never for internals of the module under test.
+- Changes that affect the external contract at the training-loop or inference level require integration tests, not just unit tests (e.g. a minimal training loop with `SimpleModel`).
+- Integration tests must be executed on actual devices, not merely written: report the execution hardware spec and results in the PR.
+
 ## Tool Caveats
 
 ### Edit tool auto-formatter

CLAUDE.md

@@ -25,6 +25,16 @@
 - Comments should explain **why**, not **what**. Describe the purpose and reasoning, not the mechanics that the code already shows.
 - New features must include corresponding tests and documentation updates.
 
+### Test discipline
+
+- Tests verify contracts, not implementations: a test is well-formed only if a different correct implementation of the same contract passes it.
+- Before writing a test, name the concrete incorrect behavior it would catch; if you cannot name one, do not write it.
+- Assert on observable outcomes through public/stable interfaces; do not assert private method return values or exact internal strings unless pinning a specific fixed bug (justify in a comment).
+- Anchor to an external oracle or an independently derived reference instead of re-implementing the logic under test.
+- Mocks must stand in for a collaborator's documented contract (schema, protocol), never for internals of the module under test.
+- Changes that affect the external contract at the training-loop or inference level require integration tests, not just unit tests (e.g. a minimal training loop with `SimpleModel`).
+- Integration tests must be executed on actual devices, not merely written: report the execution hardware spec and results in the PR.
+
 ## Tool Caveats
 
 ### Edit tool auto-formatter

CONTRIBUTING.md

@@ -53,10 +53,12 @@ make test
 
 ### Diff-based CI test selection
 Some GPU CI workflows (currently `modal-torch-latest`, which runs `tests/unit/v1/`)
-don't run the whole suite on every PR. Instead, `ci/tests_fetcher.py` looks at the
-files your PR changes, builds an import graph over `deepspeed/` and the `unit` test
-helpers, and runs only the tests that could be affected. This keeps CI fast without
-losing coverage (`push` to `master` always runs everything). The full design — and
+run their modal tests on the merge queue entry instead of on every PR push, to
+conserve Modal GPU quota. `ci/tests_fetcher.py` looks at the files your PR
+changes, builds an import graph over `deepspeed/` and the `unit` test helpers,
+and selects only the tests that could be affected. The selection is previewed on
+the PR by a cheap no-secret job and executed once the PR enters the merge queue
+(`push` to `master` always runs everything). The full design — and
 how to drive and extend it — is in
 [`.github/workflows/TEST_SELECTION.md`](.github/workflows/TEST_SELECTION.md).
 
@@ -91,9 +93,11 @@ cat ci/.test_selection/test_list.txt
 python ci/tests_fetcher.py --base origin/master --explain
 ```
 
-> Note: under `pull_request_target` the `deploy` job runs the PR's `deepspeed/` and
-> `tests/` but restores `ci/` from the base branch (the CI scripts hold the modal/HF
-> secrets). So changes to `ci/*` only take effect once merged — validate them via a
+> Note: plain PR events never start the modal `deploy` job (it stays skipped to
+> conserve GPU quota), and under `pull_request_target` GitHub runs the base branch's
+> CI scripts anyway. Merge queue runs also use trusted master for the CI scripts —
+> your merged tree is only fetched as the candidate under test — so changes to
+> `ci/*` take effect once **merged**; validate them before that via a
 > `pull_request`-triggered run or the `modal` CLI.
 
 ### Model Tests

ci/accelerate.py

@@ -3,26 +3,33 @@
 
 # DeepSpeed Team
 
+import os
 from pathlib import Path
 
 import modal
 
 ROOT_PATH = Path(__file__).parents[1]
 
+# Tree under test. On merge queue runs this launcher is checked out from
+# trusted master, so the merged candidate is fetched separately and pointed at
+# here. Remotely the env var is unset and this module is materialized inside
+# the image at /root, where it must keep resolving to ROOT_PATH.
+CANDIDATE_PATH = Path(os.environ.get("DS_CI_CANDIDATE_ROOT") or ROOT_PATH)
+
 # yapf: disable
 image = (modal.Image
          .from_registry("pytorch/pytorch:2.9.1-cuda12.8-cudnn9-devel", add_python="3.10")
          .apt_install("git")
          .pip_install("uv")
          # uv_pip_install already includes --compile-bytecode
          .uv_pip_install("datasets==3.6.0", extra_options="--system")
-         .pip_install_from_requirements(ROOT_PATH / "requirements/requirements.txt", gpu="any")
-         .pip_install_from_requirements(ROOT_PATH / "requirements/requirements-dev.txt", gpu="any")
-         .add_local_dir(ROOT_PATH , remote_path="/root/", copy=True)
+         .pip_install_from_requirements(CANDIDATE_PATH / "requirements/requirements.txt", gpu="any")
+         .pip_install_from_requirements(CANDIDATE_PATH / "requirements/requirements-dev.txt", gpu="any")
+         .add_local_dir(CANDIDATE_PATH, remote_path="/root/", copy=True)
          .run_commands("pip install /root")
-         .add_local_dir(ROOT_PATH / "accelerator", remote_path="/root/deepspeed/accelerator")
-         .add_local_dir(ROOT_PATH / "csrc", remote_path="/root/deepspeed/ops/csrc")
-         .add_local_dir(ROOT_PATH / "op_builder", remote_path="/root/deepspeed/ops/op_builder")
+         .add_local_dir(CANDIDATE_PATH / "accelerator", remote_path="/root/deepspeed/accelerator")
+         .add_local_dir(CANDIDATE_PATH / "csrc", remote_path="/root/deepspeed/ops/csrc")
+         .add_local_dir(CANDIDATE_PATH / "op_builder", remote_path="/root/deepspeed/ops/op_builder")
         )
 
 app = modal.App("deepspeedai-accelerate-ci", image=image)