Add Native Windows Support

.github/workflows/release.yml

@@ -6,9 +6,8 @@ on:
       - 'v*.*.*'
 
 jobs:
-  deploy:
+  build-sdist:
     runs-on: ubuntu-24.04
-    environment: release-env
 
     steps:
     - uses: actions/checkout@v7
@@ -28,13 +27,99 @@ jobs:
         pip install setuptools
         pip install build
         DS_BUILD_STRING=" " python -m build --sdist
-    - name: Publish to PyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
+    - name: Upload sdist
+      uses: actions/upload-artifact@v4
+      with:
+        name: dist-sdist
+        path: dist/*.tar.gz
+
+  build-wheels:
+    name: windows-release / build wheel (${{ matrix.arch }}, py${{ matrix.python-version }})
+    runs-on: ${{ matrix.runs-on }}
+    concurrency:
+      group: windows-release-${{ github.ref }}-${{ matrix.arch }}-${{ matrix.python-version }}
+      cancel-in-progress: true
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: [arm64, x64]
+        python-version: ['3.11', '3.12', '3.13']
+        include:
+          - arch: arm64
+            runs-on: windows-11-arm
+            msvc-arch: arm64
+          - arch: x64
+            runs-on: windows-latest
+            msvc-arch: amd64
+
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Setup virtual environment
+        uses: ./.github/workflows/setup-win-venv
+        with:
+          python-version: ${{ matrix.python-version }}
+          msvc-arch: ${{ matrix.msvc-arch }}
+
+      - name: Get release version from tag
+        shell: bash
+        run: |
+          echo "RELEASE_VERSION=${GITHUB_REF#refs/*/v}" >> $GITHUB_ENV
+
+      - name: Check release version
+        shell: bash
+        run: |
+          pip install packaging
+          python release/check_release_version.py --release_version ${{ env.RELEASE_VERSION }}
+
+      - name: Build wheel
+        shell: pwsh
+        run: |
+          $env:DS_BUILD_STRING = " "
+          python -m build --wheel --no-isolation
+
+      - name: Upload wheel
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist-wheel-${{ matrix.arch }}-py${{ matrix.python-version }}
+          path: dist/*.whl
+
+  publish-pypi:
+    name: release / publish to PyPI
+    needs: [build-sdist, build-wheels]
+    runs-on: ubuntu-24.04
+    environment: release-env
+
+    steps:
+      - uses: actions/checkout@v7
+      - name: Download all distributions
+        uses: actions/download-artifact@v4
+        with:
+          pattern: dist-*
+          path: dist
+          merge-multiple: true
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}
+          repository-url: https://upload.pypi.org/legacy/
+
+  bump-version:
+    name: release / bump version
+    needs: [publish-pypi]
+    runs-on: ubuntu-24.04
+    environment: release-env
+
+    steps:
+    - uses: actions/checkout@v7
       with:
-        password: ${{ secrets.PYPI_API_TOKEN }}
-        repository-url: https://upload.pypi.org/legacy/
+        ref: "master"
+    - name: Get release version from tag
+      run: |
+        echo "RELEASE_VERSION=${GITHUB_REF#refs/*/v}" >> $GITHUB_ENV
     - name: Bump version
       run: |
+        pip install packaging
         python release/bump_patch_version.py --current_version ${{ env.RELEASE_VERSION }}
     - name: Create Pull Request
       uses: peter-evans/create-pull-request@v8

.github/workflows/setup-win-venv/action.yml

@@ -0,0 +1,64 @@
+name: Create Virtual Environment
+description: Set up Python, the MSVC toolchain, and install PyTorch for Windows CI
+
+inputs:
+  python-version:
+    description: Python version to set up
+    required: true
+  msvc-arch:
+    description: MSVC architecture to configure (e.g. amd64, arm64)
+    required: true
+  torch-version:
+    description: PyTorch version to install
+    required: false
+    default: '2.14.0'
+  torchvision-version:
+    description: torchvision version to install (skipped if empty)
+    required: false
+    default: ''
+
+runs:
+  using: "composite"
+  steps:
+    - name: Configure DeepSpeed build flags
+      shell: pwsh
+      run: |
+        "DISTUTILS_USE_SDK=1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+        "DS_BUILD_OPS=1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+        "DS_BUILD_AIO=1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+        "DS_BUILD_PIN_MEMORY=1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+        "DS_ENABLE_NINJA=1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+
+    - uses: actions/setup-python@v6
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Setup MSVC environment
+      shell: pwsh
+      run: |
+        $vswhere = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe"
+        $vsPath = & $vswhere -latest -products * -property installationPath
+        $vcvarsall = Join-Path $vsPath 'VC\Auxiliary\Build\vcvarsall.bat'
+        $envDump = cmd /c "`"$vcvarsall`" ${{ inputs.msvc-arch }} >nul && set"
+        foreach ($line in $envDump) {
+          if ($line -match '^([^=]+)=(.*)$') {
+            "$($matches[1])=$($matches[2])" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+          }
+        }
+
+    - name: Verify MSVC toolchain
+      shell: pwsh
+      run: where cl
+
+    - name: Install build tools
+      shell: pwsh
+      run: pip install build wheel setuptools
+
+    - name: Install PyTorch
+      shell: pwsh
+      run: pip install torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cpu
+
+    - name: Install torchvision
+      if: ${{ inputs.torchvision-version != '' }}
+      shell: pwsh
+      run: pip install torchvision==${{ inputs.torchvision-version }} --index-url https://download.pytorch.org/whl/cpu

.github/workflows/windows-torch-latest.yml

@@ -0,0 +1,49 @@
+name: windows-torch-latest
+
+on:
+  workflow_dispatch:
+  pull_request:
+  merge_group:
+    branches: [ master ]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    name: windows-torch-latest / unit tests (${{ matrix.arch }}, py${{ matrix.python-version }})
+    runs-on: ${{ matrix.runs-on }}
+    strategy:
+      fail-fast: false
+      matrix:
+        arch: [arm64, x64]
+        python-version: ['3.11', '3.12', '3.13']
+        include:
+          - arch: arm64
+            runs-on: windows-11-arm
+            msvc-arch: arm64
+            torchvision-version: 0.29.0a0
+          - arch: x64
+            runs-on: windows-latest
+            msvc-arch: amd64
+            torchvision-version: 0.29.0
+
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Setup virtual environment
+        uses: ./.github/workflows/setup-win-venv
+        with:
+          python-version: ${{ matrix.python-version }}
+          msvc-arch: ${{ matrix.msvc-arch }}
+          torchvision-version: ${{ matrix.torchvision-version }}
+
+      - name: Install deepspeed
+        run: pip install .[win-dev,autotuning] --pre --no-build-isolation
+
+      - name: Unit tests
+        run: |
+          pytest --maxfail=100 --color=yes --durations=0 --verbose -rF -n 4 unit\
+          pytest --maxfail=100 --color=yes --durations=0 --verbose -rF -m sequential unit\
+        working-directory: tests

README.md

@@ -152,6 +152,13 @@ We regularly push releases to [PyPI](https://pypi.org/project/deepspeed/) and en
 pip install deepspeed
 ```
 
+On Windows ARM64, PyPI does not host a matching PyTorch wheel, so install
+PyTorch from PyTorch's own index first:
+
+```bash
+pip install deepspeed --extra-index-url https://download.pytorch.org/whl/cpu
+```
+
 After installation, you can validate your install and see which extensions/ops
 your machine is compatible with via the DeepSpeed environment report.
 

build_win.bat

@@ -1,15 +1,10 @@
 @echo off
 
-set CUDA_HOME=%CUDA_PATH%
 set DISTUTILS_USE_SDK=1
-
-set DS_BUILD_AIO=0
-set DS_BUILD_CUTLASS_OPS=0
-set DS_BUILD_EVOFORMER_ATTN=0
-set DS_BUILD_FP_QUANTIZER=0
-set DS_BUILD_GDS=0
-set DS_BUILD_RAGGED_DEVICE_OPS=0
-set DS_BUILD_DEEP_COMPILE=0
+set DS_BUILD_OPS=1
+set DS_BUILD_AIO=1
+set DS_BUILD_PIN_MEMORY=1
+set DS_ENABLE_NINJA=1
 
 python -m build --wheel --no-isolation
 

csrc/adam/cpu_adam_impl.cpp

@@ -57,8 +57,11 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params,
             size_t copy_size = TILE;
             if ((t + TILE) > _param_size) copy_size = _param_size - t;
             size_t offset = copy_size + t;
+            // MSVC's OpenMP loop index must be signed; GCC/Clang accept either.
+            const auto t_signed = static_cast<int64_t>(t);
+            const auto offset_signed = static_cast<int64_t>(offset);
 #pragma omp parallel for if (parallel)
-            for (size_t k = t; k < offset; k++) {
+            for (int64_t k = t_signed; k < offset_signed; k++) {
                 float grad = (float)grads[k];
                 float param = (float)_params[k];
                 float momentum = _exp_avg[k];
@@ -289,8 +292,10 @@ void adamw_rollback_inplace(float* params,
     const float lr_lambda = lr * lambda;
     const float one_minus_lr_lambda = 1.0f - lr_lambda;
 
+    // MSVC's OpenMP loop index must be signed; GCC/Clang accept either.
+    const auto param_size_signed = static_cast<int64_t>(param_size);
 #pragma omp parallel for
-    for (size_t i = 0; i < param_size; ++i) {
+    for (int64_t i = 0; i < param_size_signed; ++i) {
         const float bias_correction1 = 1.0f - beta1_pow;
         const float bias_correction2 = 1.0f - beta2_pow;