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>

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>

fix: AutoTP partition_config uses full hierarchical module path (#8088) ## Problem When using custom patterns with AutoTP, built from the immediate parent only instead of the accumulated hierarchical path. This meant patterns like never matched because the name passed was just (missing prefix). **Impact**: Custom patterns are silently ignored — parameters are not TP-sharded, causing OOM on multi-GPU setups with large models. ## Fix Two changes in : 1. **Line 574**: Build `full_name` from `class_name` (accumulated hierarchical path) instead of `prev_name` (immediate parent only). This ensures patterns see the complete module path. 2. **Line 591**: Pass `name` instead of `full_name` to the recursive `_replace_module` call, preventing path duplication at deeper nesting levels. Without this, `class_name` would accumulate the full prefix twice (e.g., `model.layers.0.model.layers.0.self_attn`). ## Note This bug only affects the `partition_config` code path (custom patterns). The default `linear_policies` and HuggingFace `tp_plan` paths are unaffected. Signed-off-by: Guokai Ma <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>

Do not pin DDP device_ids for CPU reference models (#8399) ## Problem `tests/unit/v1/zero/test_zero_user_backward.py` builds torch DDP **reference** models (the known-good baseline that ZeRO results are compared against) at five sites: ```python model_ddp = DDP(model_ddp, device_ids=[rank], output_device=rank) ``` `device_ids=[rank]` assumes rank ↔ GPU index. torch's DDP contract only allows `device_ids`/`output_device` for single-device GPU modules; **CPU modules live on one shared device and must omit them**, so multi-rank CPU runs died inside the DDP constructor with: ``` ValueError: DistributedDataParallel device_ids and output_device arguments only work with single-device/multiple-device GPU modules or CPU modules, ... ``` ## Change Route all five sites through one helper: ```python def wrap_ddp_reference(model, device, rank): # Only indexed devices take device_ids/output_device; CPU modules live on one shared device. if torch.device(device).type == 'cpu': return DDP(model) return DDP(model, device_ids=[rank], output_device=rank) ``` Design notes: - the condition restates the exact precondition torch's own DDP constructor enforces, in torch's device vocabulary — it follows the model's actual device rather than the global accelerator configuration; - a single helper means new reference-model sites cannot forget the branch (the first fix round in #8381 missed 4 of the 5 sites for exactly this reason); - GPU behavior is unchanged. ## Validation Validated as part of the multi-rank CPU CI experiment in #8381: all DDP-constructor failures were eliminated (the file's few remaining failures there are unrelated — see the triage table in that PR), zero regressions vs the same-commit baseline. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Consolidate LR scheduler tensor-LR tests into a contract suite (#8206) ## What Replaces the three per-scheduler `*_preserves_tensor_lr` tests (`WarmupLR`, `WarmupCosineLR`, `OneCycle`) with a single parametrized contract test driven by a `TENSOR_LR_CONTRACTS` table, plus a registry-coverage assertion that fails collection when a scheduler listed in `VALID_LR_SCHEDULES` has no contract row. ## Why #8205 recently fixed `OneCycle._initialize_lr` bypassing `update_lr()`, which silently replaced a caller-supplied tensor LR with a Python scalar on construction. WarmupLR and WarmupCosineLR had the same class of bug fixed in #8202. Each fix was caught only because someone happened to write a hand-rolled regression test for that specific scheduler. This PR turns that pattern into a contract: every scheduler in `VALID_LR_SCHEDULES` is pinned by one table row asserting tensor identity / shape / dtype at init and after `step(1)`. Two failure modes are covered: 1. **Registered but buggy** (writes `group['lr'] = scalar` in its constructor) — the parametrized contract case fails on the identity assertion. 2. **Forgotten** (a new scheduler is added to `VALID_LR_SCHEDULES` without a contract row) — `test_all_schedulers_covered_by_tensor_lr_contract` fails at collection with the exact missing class name. Both guards were verified empirically against this branch: - Applied the contract test on top of `pr-8202` (i.e. before #8205): the two `OneCycle` cases fail with `assert 0.01 is tensor(0.1000, dtype=torch.float64)`, exactly the symptom #8205 fixes. The other 8 cases pass, so no false positives. - Removed the `OneCycle` row from `TENSOR_LR_CONTRACTS`: the coverage test fails with `missing tensor-LR contract for: ['OneCycle']`. ## Test footprint CPU-only, no GPU/distributed runner required: ``` pytest tests/unit/runtime/test_lr_schedulers.py -k "preserves_tensor or covered_by_tensor" # 11 passed (5 schedulers x 2 shapes + 1 coverage) ``` `pre-commit run --files tests/unit/runtime/test_lr_schedulers.py` passes (yapf, flake8, check-torchdist, check-license, codespell). --- Authored by an AI coding agent (opencode, model: glm-5.2) under interactive human direction and review. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

docs: add test discipline rules to agent guidelines (#8372) ## Motivation AI coding agents increasingly submit tests alongside features, and a recurring failure mode of such contributions is tests coupled to the implementation rather than the contract: - asserting private method return values or exact internal strings, - re-implementing the logic under test as the test's own reference, - mocks that stand in for internals instead of a documented collaborator contract. Such tests pass trivially, break on harmless refactors (or worse, survive bugs they should catch), and give false coverage confidence. ## What this PR adds A `### Test discipline` section to the agent guidelines (`AGENTS.md` + `CLAUDE.md`, kept in sync as the file header requires) with operational rules: 1. A test is well-formed only if a different correct implementation of the same contract passes it. 2. Name the concrete incorrect behavior a test would catch before writing it; if none can be named, do not write it. 3. Assert observable outcomes through public/stable interfaces; white-box pinning only for specific fixed bugs, justified in a comment. 4. Anchor to an external oracle or an independently derived reference instead of re-implementing the logic under test. 5. Mocks must stand in for a collaborator's documented contract (schema, protocol), never for internals of the module under test. The rules are deliberately example-free to keep the guide lean. --------- Signed-off-by: Ma, Guokai <guokai.ma@intel.com>