Conversation
The MPS accelerator was a stub: memory queries returned None, no communication backend was set, fp16/bf16 were reported unsupported, and every op builder resolved to NotImplementedBuilder. Training with any ZeRO stage failed on an Apple Silicon machine. - accelerator/mps_accelerator.py: report real torch.mps memory stats, fp16/bf16 support, torch.mps.Event, gloo as the comm backend, and treat MPS as a synchronized device (single in-order command queue, no public streams). Unified memory makes pin_memory a no-op. - deepspeed/comm/torch.py: gloo cannot operate on MPS tensors, so collectives stage MPS tensors through CPU copies (stage_on_cpu). - accelerator/abstract_accelerator.py + runtime/zero: MPS has no fp64, so gradient-norm accumulation picks its dtype via the new is_fp64_supported()/get_norm_dtype() instead of hard-coded double(). - op_builder/mps: new backend package with a torch._foreach based FusedAdam that mirrors the CUDA multi_tensor_adam math; Metal kernels will later plug into the same builder classes. - tests/unit/common.py: MPS must use spawn (Metal's compiler service is lost in forked children) and reports its own device count. - tests/unit/ops/adam/test_adamw.py: check FusedAdam against torch.optim.Adam/AdamW on the active accelerator. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
|
Superseded by #8293 (same branch pushed directly to the upstream repo). |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1865963518
ℹ️ 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".
| if kwargs.get('async_op', False): | ||
| return StagedWork(work, copy_back) | ||
| copy_back() |
There was a problem hiding this comment.
On MPS, asynchronous all_reduce and all_reduce_coalesced calls copy the staged CPU tensor back before Gloo completes because deepspeed/comm/comm.py forwards async_op positionally at lines 658 and 682, while this wrapper checks only kwargs. Multi-rank callers using async_op=True can therefore consume stale or partially reduced tensors even after the returned raw work is waited on; bind the call arguments to the wrapped signature or otherwise detect the positional flag and return StagedWork.
Useful? React with 👍 / 👎.
| @staticmethod | ||
| @torch.no_grad() |
There was a problem hiding this comment.
When PyTorch is not preinstalled, the preceding ImportError handler leaves torch undefined, but this decorator is evaluated immediately and raises NameError. This breaks the repository's supported no-torch setup path: setup.py imports op_builder after explicitly tolerating a failed torch import, and op_builder/__init__.py eagerly imports every builder package, including this new MPS module. Define the decorated implementation only when torch is available or defer the torch reference until load()/execution.
Useful? React with 👍 / 👎.
| # MPS has no native collective backend; gloo is the only torch backend available on macOS. | ||
| self._communication_backend_name = "gloo" |
There was a problem hiding this comment.
A repo-wide search of docs/ and README.md finds no MPS or Apple Silicon documentation, although this commit introduces a user-facing accelerator backend, distributed staging behavior, supported dtypes, and important limitations. Add corresponding setup, configuration, and limitation guidance as required for new features by the workspace rules.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
Summary
The MPS accelerator was a stub: memory queries returned
None, no communication backend was set, fp16/bf16 were reported unsupported, and every op builder resolved toNotImplementedBuilder.deepspeed.initialize+ one training step failed for every ZeRO stage on an Apple Silicon machine.This PR is the first step of Apple Silicon support: make single-device training work end to end with pure-PyTorch ops. Metal kernels come later and plug into the
op_builder/mpsclasses added here.Changes
accelerator/mps_accelerator.py— realtorch.mpsmemory stats, fp16/bf16 support (bf16 gated on macOS 14+),torch.mps.Event,glooas the comm backend, andis_synchronized_device() = True(MPS has a single in-order command queue and no public streams orrecord_stream). Unified memory makespin_memorya no-op (torch'spin_memory()also raises under MPS).deepspeed/comm/torch.py— gloo cannot operate on MPS tensors (even at world size 1), so collectives stage MPS tensors through CPU copies via astage_on_cpudecorator; async ops copy back onwait().accelerator/abstract_accelerator.py+runtime/zero— MPS has no fp64. Gradient-norm accumulation now picks its dtype via a new concreteis_fp64_supported()(defaultTrue) andget_norm_dtype()instead of hard-coded.double().op_builder/mps/— new backend package (MPSOpBuilder,NotImplementedBuilder,FusedAdamBuilder).FusedAdamis implemented withtorch._foreach_*ops and mirrors the math incsrc/adam/multi_tensor_adam.cu, following the HPU precedent of Python-backed builders.tests/unit/common.py— MPS must usespawn(Metal's compiler service is lost inforkserverchildren, which hangs the harness) and reports its device count via the accelerator.tests/unit/ops/adam/test_adamw.py—test_fused_adam_matches_torchchecksFusedAdamagainsttorch.optim.Adam/AdamWon the active accelerator (fp32/bf16 × Adam/AdamW), so it also guards the CUDA kernel.Verified on an M5 Max (macOS 26.3, torch 2.13.0)
deepspeed.initialize(single process).FusedAdammatchestorch.optimto 2e-7 in fp32.DS_ACCELERATOR=mps pytest unit/runtime/test_ds_config_dict.py unit/runtime/test_ds_initialize.py unit/runtime/half_precision/test_fp16.py unit/runtime/half_precision/test_dynamic_loss_scale.py unit/runtime/zero/test_zero_grad_clip.py unit/runtime/zero/test_zero_context.py unit/checkpoint/test_zero_optimizer.py: 132 passed, 0 failed, 126 skipped (multi-device tests;device_count() == 1).Known limitations / follow-ups
FusedAdamdiffers from the CUDA kernel by ~1 bf16 ulp (CUDA computes in fp32 and stores bf16; the_foreachpath rounds in bf16).torch.mps.compile_shader, and an Apple Silicon tutorial page.