Zhipeng Wang
64fcec6ba7
## Summary This PR is the first step (phase 0) of enabling Apple Silicon support for DeepSpeed: make single-device training work end to end with pure-PyTorch ops. [**To-Do in phase 1**] Metal kernels will come later and plug into the `op_builder/mps` classes added here. 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`. `deepspeed.initialize` + one training step failed for every ZeRO stage on an Apple Silicon machine. This PR aims on enabling capabilities. ### Changes - **`accelerator/mps_accelerator.py`** — real `torch.mps` memory stats, fp16/bf16 support (bf16 gated on macOS 14+), `torch.mps.Event`, `gloo` as the comm backend, and `is_synchronized_device() = True` (PyTorch's MPS backend effectively exposes a single in-order execution stream and currently provides no public CUDA-style stream API or `record_stream` mechanism.). Unified memory makes `pin_memory` a no-op (torch's `pin_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 a `stage_on_cpu` decorator; async ops copy back on `wait()`. - **`accelerator/abstract_accelerator.py` + `runtime/zero`** — MPS has no fp64. Gradient-norm accumulation now picks its dtype via a new concrete `is_fp64_supported()` (default `True`) and `get_norm_dtype()` instead of hard-coded `.double()`. - **`op_builder/mps/`** — new backend package (`MPSOpBuilder`, `NotImplementedBuilder`, `FusedAdamBuilder`). `FusedAdam` is implemented with `torch._foreach_*` ops and mirrors the math in `csrc/adam/multi_tensor_adam.cu`, following the HPU precedent of Python-backed builders. - **`tests/unit/common.py`** — MPS must use `spawn` (Metal's compiler service is lost in `forkserver` children, which hangs the harness) and reports its device count via the accelerator. - **`tests/unit/ops/adam/test_adamw.py`** — `test_fused_adam_matches_torch` checks `FusedAdam` against `torch.optim.Adam/AdamW` on 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) - ZeRO 1/2/3 × fp32/bf16/fp16 train end to end with `deepspeed.initialize` (single process). Also tested on ZeRO stage 0 which disables ZeRO completely and falling back to standard data parallelism. - MPS `FusedAdam` matches `torch.optim` to 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 - bf16 `FusedAdam` differs from the CUDA kernel by ~1 bf16 ulp (CUDA computes in fp32 and stores bf16; the `_foreach` path rounds in bf16). - The CPU-staged gloo path is only exercised at world size 1 here; multi-Mac runs are untested. - Follow-ups: macOS arm64 CI workflow, arm64 build of CPU Adam for ZeRO-Offload, Metal kernels via `torch.mps.compile_shader`, and an Apple Silicon tutorial page. --------- Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
DeepSpeed Documentation
This directory includes the source code for the website and documentation of DeepSpeed. The code-docs/ directory is used to build deepspeed.readthedocs.io.
deepspeed.ai is the recommended way to read all DeepSpeed documentation. Directly viewing the Markdown files in this directory will not include images and other features.
Building the documentation locally
You can serve the DeepSpeed website locally. This is especially useful for development.
Prerequisites
The DeepSpeed website relies on Jekyll. There are several guides for installation. The instructions below assume you are in an Ubuntu environment and have been tested on WSL.
First ensure that you have the necessary packages (e.g., make and zlib).
sudo apt-get install build-essential zlib1g-dev ruby-full
Add these lines to your .bashrc or equivalent to ensure you have permissions to install Ruby packages without sudo.
export GEM_HOME="$HOME/gems"
export PATH="$HOME/gems/bin:$PATH"
Don't forget to source ~/.bashrc afterward 😊.
Now we can install Jekyll and Bundler:
gem install jekyll bundler
Start a local webserver
We now need to install the required Ruby packages for the website.
NOTE: you should change to this folder (i.e., docs) before running the installation command to avoid this error:
Could not locate Gemfile
NOTE: This step frequently hangs when connected to a VPN (including MSVPN). Simply disconnect for the package installation.
bundle install
Depending on your environment, you may need to add webrick to avoid the following error:
gems/gems/jekyll-3.9.5/lib/jekyll/commands/serve/servlet.rb:3:in `require': cannot load such file -- webrick (LoadError)
bundle add webrick
You can now start a local webserver via:
bundle exec jekyll serve
The website should now be accessible at http://localhost:4000
Update the Readthedocs.io API documentation
Use the following steps to update the public API documentation.
- Make your documentation changes and push them to the rtd-staging branch. This will rebuild the docs in the staging branch. NOTE: It is acceptable to force push to this branch to overwrite previous changes.
- View the result of the result of the build here
- Once the build is complete view the newly modified API documentation here
- Once you are satisfied with the changes create a new branch off of rtd-staging to push into master.