Conversation
…erator Nothing upstream exercised the MPS paths, so regressions could only be caught on a developer's Mac. The new workflow runs the MPS-green unit test subset (fused/CPU Adam, comm staging, config-driven training) on GitHub's arm64 macOS runners, which expose a working MPS device. It cannot interfere with the existing CI: PR triggers are scoped to MPS-relevant paths, the job runs in its own workflow with a hard timeout and concurrency cancellation, and it is not a required check. MPS_Accelerator now fails at construction with a clear message on torch older than 2.3, where the torch.mps memory queries ZeRO depends on do not exist, instead of an AttributeError deep inside the runtime. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e794b17f2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Spawn-based DistributedTest cases cost 60-75s each on the shared arm64 runners, which pushed the sequential suite past the job budget before pytest could even print its summary. Three workers bring it inside the limit, and a 300s per-launch timeout fails hung tests fast. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
DistributedFixture registers itself through pytest internals (_pytestfixturefunction) that stopped working in newer pytest, which is why requirements-dev.txt pins pytest<8.4; the unpinned install on the macOS runner picked a newer version and the fixture-based tests errored with 'fixture not found' instead of skipping. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
PKUWZP
requested review from
delock,
sfc-gh-truwase and
tohtana
and
a lite review from Copilot
August 31, 2026 01:46
There was a problem hiding this comment.
Pull request overview
This PR extends Apple Silicon (MPS) support by adding a dedicated macOS CI workflow to exercise MPS-gated test coverage, and by adding an early, clearer construction-time guard in MPS_Accelerator for missing torch.mps memory APIs needed by ZeRO.
Changes:
- Add a new
mps-torch-latestGitHub Actions workflow that runs a targeted MPS unit test subset onmacos-15arm64 runners. - Add an
MPS_Acceleratorconstructor check fortorch.mps.recommended_max_memoryto fail fast with an actionable error message on unsupported torch builds.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| accelerator/mps_accelerator.py | Adds a torch/MPS capability guard intended to prevent deep runtime AttributeErrors in ZeRO. |
| .github/workflows/mps-torch-latest.yml | Introduces a macOS arm64 CI workflow to run MPS-green unit tests under DS_ACCELERATOR=mps. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…to 2.3 The constructor guard was only manually validated; it now has unit tests covering both successful construction and the clear error when torch.mps.recommended_max_memory is absent, runnable on any platform since torch>=2.3 exposes the attribute everywhere. The setup guide's "2.4 or newer" is aligned to the enforced 2.3 floor, noting 2.7+ for the Metal fused Adam kernel. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
getattr-chain the torch.mps floor check so torch builds without a torch.mps module still get the clear ValueError instead of an AttributeError from the check itself. The merge_group trigger had no path scoping, so the macOS job would have run for every merge queue group despite the workflow's scoped PR triggers; MPS coverage there adds runner cost without being a required check, so drop it. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
delock
reviewed
Aug 31, 2026
recommended_max_memory is absent from the v2.3.0 and v2.4.0 tags and first appears in v2.5.0, so the guard's message, its tests, and the setup guide all advertised too low a floor. The check itself is feature-detected and unaffected. Pointed out by @delock in review. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
delock
approved these changes
Aug 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes the remaining gap in the Apple Silicon support series (#8293, #8300, #8303, #8307): none of the MPS paths were exercised by CI — every MPS-gated test skips on Linux runners, so regressions could only be caught on a developer's Mac.
macOS CI workflow (
mps-torch-latest.yml)Runs the MPS-green unit test subset on GitHub's arm64 macOS runners (
macos-15), which expose a working MPS device:unit/ops/adam/test_adamw.py— Metal/foreach FusedAdam vs fp32-math reference, CPU Adam configs incl. ZeRO-Offloadunit/comm/test_dist.py— gloo CPU-staging for collectives and P2P (TestMpsStagedP2P)unit/runtime/test_ds_config_dict.py— config-drivendeepspeed.initialize+ training stepsDesigned not to interfere with existing CI:
paths:to MPS-relevant files (accelerator/**,op_builder/mps/**,csrc/mps/**,deepspeed/comm/**, the two test dirs, and the workflow itself) — the check does not even appear on unrelated PRs.timeout-minutes: 45.schedule+workflow_dispatchfor coverage between touching PRs.torch floor check
MPS_Accelerator.__init__now fails with a clear message on torch older than 2.3, where thetorch.mpsmemory queries ZeRO depends on (recommended_max_memory) do not exist — previously this surfaced as a bareAttributeErrordeep inside ZeRO's flatten logic. Feature-detected rather than version-parsed. (The Metal FusedAdam kernel already degrades gracefully on torch withoutcompile_shader.)Validation
recommended_max_memoryhidden, construction raises the explicitValueError.