Ma, Guokai

[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP (#8168) AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards. - ds_to_universal.py: stage3 conversion recovers the (tp,dp) grid from checkpoint file names, extracts shards under the real tp_index, and reuses the stage<=2 TP-aware merge when tp_degree>1 (DP-only path preserved for tp_degree==1 -> no regression for plain ZeRO-3). - engine.py: _zero3_consolidated_16bit_state_dict nests GatherReplacedLayerParams inside GatheredParameters so save_16bit_model gathers both DP and TP; remove the blanket autotp+zero3 training block now that checkpoint consolidation is implemented. - stage3.py: load_hp_checkpoint_state resolves the TP shard before the ZeRO-DP partition, so universal checkpoint restore round-trips. Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com>

[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta (#8241) ## What this does Fixes #8231. `tp_shard` kept `num_kv_heads` / `num_attention_heads` / `n_embd` / `tp_grain_size` as **process-wide mutable globals**, written during AutoTP replacement. A second AutoTP model loaded into the same process overwrote them, so the first model's later sharding / gather / checkpoint conversion silently read the wrong values — making it unsafe to run more than one AutoTP model per process (teacher/student, online distillation, RL actor + reference). This moves that state onto a per-model `AutoTPMeta`, computed once from the model config and threaded through every sharding helper and TP layer, so each model carries its own kv-head / grain state. ## Stacking / merge order **Depends on #8185** (AutoTP uneven sharding). This branch is based on #8185's head and is opened as **draft** until #8185 lands — GitHub will drop #8185's commits from this diff automatically once it merges. Please **merge after #8185**. ## Changes - `AutoTPMeta` dataclass + `from_model_config` (single source for kv-head / attn-head / hidden extraction); `get_shard_size(_list)` take it as a required arg; tp_shard globals + `set_*` / `get_*` removed. - AutoTP threads `tp_meta` from `__init__` through every TP layer and fused-QKV helper. - Ulysses sequence parallelism gets its own `_ulysses_num_kv_heads`, decoupled from AutoTP (the AutoTP↔Ulysses coupling is gone; Ulysses's own multi-model case is left for a separate change). - Inference engine builds one `meta` per model (`_autotp_meta`) and threads it through the alibi head-sharding helpers; `_get_model_head_count` / `_get_model_kv_head_count` deleted. - kv-head / attn-head attribute lists unified behind `_kv_head_count_from` / `_attention_head_count_from` (covers chatglm, falcon, llama-class, dbrx, legacy `n_head_kv`). ## Tests Validated on 4×RTX 4080 (nccl) against #8185: full AutoTP / SP / checkpoint suite passes (127 passed); remaining failures are pre-existing env issues (transformers/HF network `client has been closed`, torch 2.12 `ProcessGroupGloo.perform_nocolor_split`, a cuda/cpu device-mismatch), each confirmed failing on the #8185 baseline too. `test_two_models_do_not_clobber_each_others_meta` is the direct regression test for #8231. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

fix(zero3): async grad offload + pinned offload buffers by default (#8207) ## Problem ZeRO-3's gradient GPU→CPU offload in `partition_grads()` used a blocking `copy_()` without `non_blocking`, and its destination buffer defaulted to pageable host memory. This forced the offload onto a synchronous, low-bandwidth (staged pageable) path with no overlap against backward compute, even though the copy already runs on the dedicated `reduce_and_partition_stream`. ZeRO stage 1/2 already issues this copy with `non_blocking=True` (`stage_1_and_2.py:1530`); stage 3 is the inconsistent one. ## Changes - `offload_config.py`: default `offload_optimizer/offload_param.pin_memory` to `True`. Pinned (page-locked) host memory is required for async, full-bandwidth DMA; the prior `False` silently selected the slow staged pageable copy. Disable only on hosts with tight `ulimit -l` memlock. - `stage3.py`: issue the grad offload copy with `non_blocking=True`. Stays on `reduce_and_partition_stream`. - `stage3.py` (2nd commit): remove an orphaned helper (`async_inplace_copy_grad_to_fp32_buffer_from_gpu`) that referenced an uninitialized attribute and had no callers. The live stage 1/2 version is untouched. ## Validation 4× RTX 4080-SUPER, autotp=2, `offload_optimizer`, `cpu_adam`, per-rank CPU affinity: | model | baseline BWD | fixed BWD | Δ BWD | |----------------|--------------|-----------|-------| | Qwen2.5-1.5B | 1810 ms | 1308 ms | -28% | | Qwen2.5-3B | 3097 ms | 2540 ms | -18% | Memory footprint unchanged; FWD/STEP unchanged. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>

Nightly full-suite run with classified failure triage (#8606) ## Motivation The master full suite on `modal-torch-latest` grew from ~52 to ~70 minutes within a month and keeps climbing ~7-8 min/week, repeatedly exhausting its Sandbox/job time budgets (the 3600 -> 4200 -> 5400s raises are chasing the growth). This PR adds a scheduled nightly full-suite run with automated, noise-free failure triage, so master landings no longer pay the full-suite cost on every push and nightly regressions are reported without flooding the tracker. ## Changes **1. Nightly full-suite run; no master push trigger** - A `schedule` (cron, daily `0 10 * * *`) run always executes the full suite on master; this is the only event whose failures feed automated regression triage. - The `push` trigger is dropped: master receives no direct pushes (the merge queue lands every entry, and each `merge_group` run already tested the merged tree). Cross-entry interaction regressions are what the nightly full suite exists to catch. - `concurrency` group now includes `github.event_name` so a nightly run and merge-queue runs don't cancel each other. **2. Controller classifies its own failures (`ci/torch_latest.py`)** - `infra` (exit 75): the Sandbox never got a GPU instance — no test ran, no candidate information. - `timeout` (exit 124): the Sandbox lifetime budget was exhausted mid-run (classified just inside the nominal budget, with a 120s grace for server-side kills). - `test` (exit 1): everything else, raised with a full traceback as before. - Each class prints a `DS_CI_FAILURE_CLASS=<class>` sentinel line that survives into job logs even when the job is killed before exiting. **3. Nightly triage (`nightly-triage.yml`, `ci/nightly_classify.sh`, `ci/nightly_report.sh`)** - Triggered by `workflow_run` on completed `modal-torch-latest` runs, filtered to scheduled master runs only; all script checkouts are pinned to the SHA the nightly actually ran (`github.event.workflow_run.head_sha`), so the classifier and the sentinel protocol it parses come from the same revision even if master moved on overnight. - **green** → move the `nightly-last-green` tag to the tested SHA (persistent good endpoint for future range work; operational failure days simply leave it in place). - **infra** → no-op (capacity problem, no candidate information). - **timeout** → open an issue (the suite outgrew its Sandbox budget; consider shard-splitting or raising it). - **unknown** → open an issue (no sentinel: the run died before the controller could classify itself — checkout, setup, or a job kill). Kept distinct from timeout instead of guessing. - **test failures** → open an issue listing the failing test files (parsed from the nightly log). - All reports go through `ci/nightly_report.sh`: labeled `nightly-triage`, and a recurring regression with the same stable title comments on the existing open issue instead of filing a duplicate (titles are deliberately SHA-free; the per-night SHA and run link go in the body). **4. Housekeeping**: refreshed stale self-tests that no longer matched master (sandbox timeout 4200 -> 5400 left behind by #8571, merge-group trusted-ref expression, the baked-image `FakeImage` chain, and the removed `reinstall Torch packages` runtime command). `ci/test_torch_latest.py` is 24/24 green; pre-commit passes on all changed files. ## Known limitations - The nightly run doubles as an upstream-drift sentinel (it installs transformers from git `main`); red nightlies caused by upstream drift are reported as ordinary test failures, and the issue text keeps the run link for human judgment. - Triage routes on the sentinel protocol of the nightly SHA; a controller change to that protocol requires this PR's self-tests (`ci/test_torch_latest.py`) to be updated in the same change. --------- Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

Split the modal CI budget into acquisition and test phases (#8404) Stacked on #8403 — review that one first; this PR's base retargets to `master` automatically once it merges. ## Problem A single wall-clock budget cannot tell **"we never got a GPU"** apart from **"the tests ran long"**. Both surface as the same job timeout today, and a run that spends 22 minutes waiting for capacity takes that time out of the budget the tests still need. Measured on completed master runs, the wait between `Sandbox.create()` and the first sandbox output: | run | acquisition wait | pytest | |---|---|---| | 08-22 12:53 | 0.1m | 37.5m | | 08-29 00:40 | 0.2m | 39.3m | | 08-31 08:45 | 7.9m | 48.8m | | 08-31 14:11 | 10.8m | 44.5m | | 08-31 17:36 | 0.9m | 38.4m | | 09-01 11:47 | 17.8m | 48.4m | | 09-02 04:00 | 22.0m | 38.5m | `Sandbox.create()` returns before the container exists, so the wait for the `l40s:2` reservation surfaces on the first `exec`. ## Change Bound the two phases separately: - **acquisition: 30 min** (`SANDBOX_ACQUIRE_TIMEOUT_SECONDS`). Past that the run aborts with `SandboxStartTimeout`, whose message states that no test ran, instead of holding a runner for the rest of the budget. - **tests: 70 min**, unchanged from #8403. The Sandbox lifetime clock starts when the container starts, so this budget is always fully available once a GPU is reserved, however long acquisition took. - **job timeout: 105 min**, now only a backstop covering 30 + 70 plus runner setup and cleanup. 30 min leaves headroom over the worst observed acquisition (22.0 min) while still failing fast when capacity never arrives. ## Observability The startup duration is printed with `flush=True`. Sandbox output is otherwise block-buffered by Python and lost when the job is killed, which is why the timed-out runs show a silent gap rather than any progress. This makes the acquisition wait directly visible in the log instead of something you have to reconstruct from timestamps. ## Tests `ci/test_torch_latest.py` gains coverage for the abort path: - `test_await_sandbox_start_gives_up_when_the_container_never_runs` — catches a controller that blocks forever on a reservation that is never satisfied. - `test_controller_aborts_without_running_tests_when_sandbox_never_starts` — catches a controller that spends the whole budget waiting, runs commands against a Sandbox that never started, or leaks the Sandbox when startup times out. - `test_await_sandbox_start_reports_startup_duration` Verified by mutation: with the `is_alive()` check removed, the suite hangs instead of failing, which is the exact defect these tests guard against. The acquisition budget is also pinned in `test_sandbox_kwargs_are_fixed_and_secret_free`, so changing a resource limit stays visible in review like the other fixed Sandbox parameters. ## Not addressed here Why acquisition grew from ~0.1 min to 18-22 min. Ruled out from the repo side (image name, preset table, pinned `modal==1.2.6`, sandbox kwargs all unchanged); confirming whether it is GPU queueing or an image pull needs the Modal dashboard for app `deepspeedai-torch-latest-ci`. This PR makes that case fail fast and legibly rather than fixing it. --------- Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

Skip fp16-config tests on accelerators without fp16 support (#8684) ## Description Every fp16-config test that reaches `deepspeed.initialize` crashes its sanity check (`Type fp16 is not supported on your device.`) on accelerators whose `is_fp16_supported()` is false. On CPU that maps to the AVX512-FP16 capability of the host, and GitHub's `ubuntu-24.04` runners are hardware-heterogeneous, so these tests flip between failure and skip depending on which runner they land on. #8398 added the first skipifs; the multi-rank CPU run in #8381 flushed out six more files: | File | Shape of the gap | |---|---| | `checkpoint/test_universal_checkpoint.py` | fp16 parametrizations fail **inside the baseline DistributedFixture's distributed run** — pytest reports a setup **ERROR** for every dependent test (48 on the multi-rank run) instead of a skip | | `v1/zero/test_zero_coalesce_grad_reduction.py` | `TestCoalesceFP16` forces an fp16 config | | `runtime/test_no_sync_ctxt.py` | dtype=float16 parametrizations of three methods; stages 2/3 additionally never reach their expected no_sync AssertionError on such hosts | | `checkpoint/test_moe_checkpoint.py` | whole class hardcodes fp16 | | `runtime/zero/test_zero_offloadpp.py` | `TestZeroPartialOffloadConfigSweep` hardcodes fp16 | | `checkpoint/test_pipeline.py` | fp16 enabled for zero_stage > 0; only that parametrization skips, zero_stage=0 keeps running | With this PR, **every fp16-config test in the suite guards on `is_fp16_supported()`** — the capability gap is a skip, not a failure, on any accelerator. ## Validation (executed on real hardware) - 20-core x86_64 CPU without AVX512-FP16, gloo, 2-4 ranks (`LOCAL_SIZE=2/4`) - Before: 18 failures + 48 setup ERRORs across these files on the multi-rank CPU run - After: every affected parametrization skips; adjacent non-fp16 parametrizations keep passing (e.g. pipeline zero_stage=0 runs to completion) - Full-suite evidence in #8381: the multi-rank CPU run went 8 failures -> 0 with these guards Sibling PRs from the same series: #8397, #8398, #8399, #8407, #8409, #8559, #8648. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Fix silent gradient-bucket drop on CPU in ZeRO-1/2 non-contiguous reduction (#8382) ## Problem While investigating multi-rank CPU CI (deepspeedai/DeepSpeed#8381), `TestGradientAllreduceOp` failed for two variants (`sum` + `prescale_gradients=True` + ZeRO-1/2 + `contiguous_gradients=False` + `reduce_scatter=False`) with results off by a factor of world_size, and `TestUnmanagedGradientAccumulationOffload` showed a numeric mismatch. Separately, any test comparing tensors across ranks on CPU crashed with `RuntimeError: Invalid device string` or a gloo allgather error. ## Root cause `split_half_float_double()` built legacy type strings from the accelerator device name (`torch.{device}.FloatTensor`). CPU tensors report `torch.FloatTensor` with **no device prefix**, so no tensor ever matched: the bucket list came back empty and **the gradient all-reduce was silently skipped** — every rank updated on its local gradients. - With `gradient_allreduce_op=mean` this stays invisible: identical per-rank gradients make the unreduced local value equal the mean. - With `sum` the result is wrong by a factor of world_size. The same variants pass on GPU CI (all 18 `TestGradientAllreduceOp` variants and 5 `TestGradientAllreduceOpTraining` variants are green on `modal-torch-latest`), because `torch.cuda.FloatTensor` matches on CUDA — the bug is CPU-specific, and CPU multi-rank runs were the first to reach this code. The harness crashes were two stacked bugs in `reduce_boolean_flags()` (tests/unit/common.py): `current_device()` returns a bare rank index (a `LOCAL_RANK` string on CPU) that torch cannot parse as a device, and gloo rejects 0-dim inputs to `all_gather_into_tensor`. ## Fix - Group gradient buckets by `t.dtype` instead of legacy type strings, making bucketing device-independent (behavior-equivalent on CUDA, where the legacy strings matched). - Use `current_device_name()` (valid on every accelerator) and carry the flag as a 1-element tensor. ## Verification (CPU/gloo, world sizes 2–3) | Test | Before | After | |---|---|---| | `TestGradientAllreduceOp` (18 variants) | 2 failed | 18/18 pass | | `TestGradientAllreduceOpTraining` (5) | crash in comparison helper | pass (3 locally + muon skips fp16) | | `TestUnmanagedGradientAccumulationOffload` + `InactiveParams` (6) | crash + numeric mismatch | 6/6 pass | | `TestZero3ParamPartitioningBase` / `TestZeroToFP32` (regression smoke) | — | pass | GPU behavior is unchanged by construction (same buckets, same order); GPU CI will re-confirm. ## Follow-up With the comparison helper fixed, `test_unmanaged_varying_backward_count[3]` (ZeRO-3) can now reach its numeric comparison and reveals a pre-existing unmanaged-vs-managed mismatch on CPU that was previously masked by the harness crash. Out of scope here. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

[Workflow] run modal GPU workflows only from the merge queue (#8412) ## Motivation Both modal workflows spend Modal GPU quota on every PR event (`review_requested` / `ready_for_review` / `synchronize`), while the merge queue entry — the commit that actually gates the merge — never runs them, because the workflows have no `merge_group` trigger. This PR moves the GPU spend to the merge queue only. ## Behavior change | Event | Before | After | |---|---|---| | `pull_request_target` (review_requested / ready_for_review / synchronize) | collect + **Modal run** (when tests are selected) | collect only (no secrets, no Modal) — `deploy` reports *skipped*, which still satisfies Required checks | | `merge_group` (checks_requested) | **nothing runs** | collect + Modal run on the merged tree, gating the merge | | `push` to master / `workflow_dispatch` | collect + Modal run | unchanged | ## Changes - `.github/workflows/modal-torch-latest.yml` - add a `merge_group` (`checks_requested`) trigger — required checks must be reported on queued entries or the merge fails; - the candidate-tree fetch and diff-based selection now also run for `merge_group` (the merge-group commit diffed against `merge_group.base_sha`), so queue runs keep subset selection instead of falling back to the full suite; - `deploy` gains `github.event_name != 'pull_request_target'`; the explicit failure-propagation path (a failing `collect-tests` still fails the check on every event) is preserved. - `.github/workflows/modal-accelerate.yml`: same trigger and the same `deploy` gating. `dorny/paths-filter@v4` (≥ v4.0.1) supports `merge_group` natively, so path filtering is unchanged. - Docs: `TEST_SELECTION.md` (job flow, deploy gating, a `merge_group` trust-context note in the security model) and `CONTRIBUTING.md`. Net effect: one Modal run per queued (batched) merge instead of one per PR event. The repo's merge queue is active (recent merges went through it), so this takes effect immediately once merged. ## Verification - YAML parse + trigger-structure check for both workflows. - Simulated the GitHub-expression fallbacks for all four event contexts: fork head/base SHAs resolve as before; for `merge_group` the candidate resolves to the base-repo merge-group SHA and the queue's base SHA. - Deploy-gating truth table (10 cases): PR events never create a Sandbox (even for `mode=all`); selector failures still fail the check on every event; queue entries run for `all`/`subset` and skip for `none`. - `bash -n` on the modified inline script; `python ci/test_tests_fetcher.py` 16/16; `pre-commit run --files` passes on all four files. ## Notes for reviewers - `merge_group` runs use the workflow/controller from the merged commit with access to secrets — inherent to GitHub's merge queue, documented in the security-model section of `TEST_SELECTION.md`; PRs touching `ci/*` should be reviewed with that in mind. - If the modal check names are added to the required-checks list, queue entries become gated by the modal results (PR-side `skipped` counts as success). No settings change is needed for this PR itself. ## Security hardening (follow-up commit) Follow-up `586a049a4` (from Codex review): on `merge_group` runs, every GitHub-side checkout now resolves to the **trusted base revision** (`github.event.merge_group.base_sha`) instead of the merge-group commit, so queued-PR code never runs next to the Modal/HF tokens: - `modal-torch-latest`: the selector and the controller come from trusted master; the merged candidate is still only ever fetched by exact SHA inside the no-secret Modal Sandbox. - `modal-accelerate`: the `modal run` launcher comes from trusted master; the merged candidate is fetched as validated git data (`ci/torch_latest.py checkout-candidate`) into `$RUNNER_TEMP/deepspeed-candidate`, and `ci/accelerate.py` uploads the tree under test from `DS_CI_CANDIDATE_ROOT` when set. Remotely the env var is unset and the module keeps resolving inside the image at `/root` as before. The residual exposure is a queued PR rewriting the workflow YAML itself — inherent to GitHub's merge queue and documented in `TEST_SELECTION.md`. Bootstrap note: this PR's own queue entry runs the pre-change trusted launcher once; the new launcher/candidate split is live from the next merge on. --------- Signed-off-by: Ma, Guokai <guokai.ma@intel.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

[AutoTP] Allow ZeRO stage 3 inference with tensor parallelism (#8167) Replace the unconditional zero_optimization_stage() <= 2 assert in _configure_tensor_parallel_states with a guard that only blocks AutoTP + ZeRO-3 when an optimizer is present. The ZeRO-Inference path (no optimizer → DummyOptim → DeepSpeedZeRoOffload) is now permitted; the training path (optimizer → DeepSpeedZeroOptimizer_Stage3) raises NotImplementedError. AutoTP with ZeRO-3 checkpoint saving should be tracked seperatly in a seperate PR. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Wait for all async all-gathers in coalesced ZeRO-3 parameter gather (#8539) ## Problem While investigating multi-rank CPU CI (deepspeedai/DeepSpeed#8381), `TestUnmanagedGradientAccumulation::test_unmanaged_varying_backward_count[3]` failed intermittently with parameters turned to NaN, and `test_zero_coalesce_grad_reduction` showed mismatches with denormal garbage values (`4.8e+37`, `1e-38`) — the signature of uninitialized memory. ## Root cause `_allgather_params_coalesced()` launches one async all_gather per parameter, but only waits on the **last** handle before pointing `param.data` at the `torch.empty` flat buffers: ```python launch_handles[-1].wait() param.data = gathered_tensor.narrow(...) if not get_accelerator().resolves_data_dependency(): get_accelerator().synchronize() ``` On CUDA this is safe by construction: the ops are enqueued from a single thread onto the same stream (FIFO), and `torch.cuda.synchronize()` waits for everything anyway. On gloo each async handle runs on an independent background thread with **no ordering between handles** — waiting for the last one says nothing about the earlier ones — and the CPU accelerator's `synchronize()` is a no-op. Reading `param.data` can therefore race the gather and expose uninitialized memory. The corruption is timing-dependent (locally 9/20 runs fail; #8382's comparison-helper fix made the numeric assertions reachable), which is why it was never caught on GPU CI. ## Fix Wait on every handle. On CUDA the earlier handles are guaranteed complete by stream ordering, so the extra waits are immediate status checks with no real waiting; the change makes correctness independent of that structural coincidence. ## Verification (CPU/gloo, world sizes 2–3) | Test | Before | After | |---|---|---| | `test_unmanaged_varying_backward_count[3]` | 9/20 failed | 0/20 failed | | `TestZero3ParamPartitioningBase` / `TestGradientAllreduceOp` / `TestUnmanagedGradientAccumulation` (regression, 57 cases) | — | 55 pass + 2 skip | GPU behavior is unchanged by construction (same buckets, same order, stream-FIFO + full synchronize); GPU CI will re-confirm. Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@intel.com>