From f298973e86b117088d92756b5183d5b805f3f6ad Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:33:08 -0700 Subject: [PATCH 01/11] deepspeed: fix Windows path/timer assumptions and related latent bugs --- deepspeed/checkpoint/ds_to_universal.py | 4 +++- deepspeed/compile/patch_fake_tensor.py | 8 ++++++++ deepspeed/env_report.py | 6 ++++++ deepspeed/launcher/constants.py | 6 +++++- deepspeed/profiling/flops_profiler/profiler.py | 12 +++++++++--- deepspeed/runtime/config.py | 6 ++++-- deepspeed/utils/timer.py | 10 ++++++---- 7 files changed, 41 insertions(+), 11 deletions(-) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index ec9c728db4a6..1abc68ea9b94 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -252,7 +252,9 @@ def _merge_zero_shards(param_base_path, state, tp_degree, slice_shapes=None): empty_tp_indices.append((tp_index, len(slices))) continue - pattern = re.compile(f"{prefix_path}\\.([0-9]+)") + # prefix_path is a filesystem path (os.path.join), so on Windows it contains backslashes, + # which regex would otherwise treat as escape sequences; re.escape neutralizes them. + pattern = re.compile(f"{re.escape(prefix_path)}\\.([0-9]+)") dp_indices = set() for p in paths: m = pattern.match(p) diff --git a/deepspeed/compile/patch_fake_tensor.py b/deepspeed/compile/patch_fake_tensor.py index 72a2195a83ce..599c01b37ec7 100644 --- a/deepspeed/compile/patch_fake_tensor.py +++ b/deepspeed/compile/patch_fake_tensor.py @@ -76,6 +76,13 @@ def _resolve_zero3_guarded_value(builder, guard, value): def patch_fake_tensor(): + # Idempotent: called both here (tests) and from init_z3.py on every ZeRO-3 + # compile engine init. Without this guard, re-invoking within the same + # process (e.g. a later engine/test reusing a worker that already patched) + # would wrap the previous wrapper again, stacking layers with no teardown. + if getattr(FakeTensorMode.from_tensor, "_ds_zero3_patched", False): + return + # dynamo tracer uses wrap_to_fake_tensor_and_record # Wrapping FakeTensorMode.from_tensor is not sufficient as dynamo generates SymbolicContext before calling from_tensor original_wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record @@ -132,4 +139,5 @@ def from_tensor_wrapper(self, t, *args, **kwargs): with unset_fake_temporarily(): return original_from_tensor(self, wrap_if_ds_param(t), *args, **kwargs) + from_tensor_wrapper._ds_zero3_patched = True FakeTensorMode.from_tensor = from_tensor_wrapper diff --git a/deepspeed/env_report.py b/deepspeed/env_report.py index 5d2b0acd637f..f4417ff5175d 100644 --- a/deepspeed/env_report.py +++ b/deepspeed/env_report.py @@ -4,6 +4,7 @@ # DeepSpeed Team import os +import sys import torch import deepspeed import subprocess @@ -102,6 +103,11 @@ def installed_cann_version(): def get_shm_size(): + if sys.platform == "win32": + # os.statvfs (and /dev/shm itself) don't exist on Windows; unlike the + # UNKNOWN case below, this isn't a detection failure, there is just no + # single size-capped shared-memory mount to report here. + return "N/A (Windows has no /dev/shm equivalent)", None try: shm_stats = os.statvfs('/dev/shm') except (OSError, FileNotFoundError, ValueError, AttributeError): diff --git a/deepspeed/launcher/constants.py b/deepspeed/launcher/constants.py index 366ae0b236f8..6f7ebbf63e56 100644 --- a/deepspeed/launcher/constants.py +++ b/deepspeed/launcher/constants.py @@ -3,6 +3,9 @@ # DeepSpeed Team +import os +import tempfile + PDSH_LAUNCHER = 'pdsh' PDSH_MAX_FAN_OUT = 1024 @@ -11,6 +14,7 @@ IMPI_LAUNCHER = 'impi' SLURM_LAUNCHER = 'slurm' MVAPICH_LAUNCHER = 'mvapich' -MVAPICH_TMP_HOSTFILE = '/tmp/deepspeed_mvapich_hostfile' +# /tmp doesn't exist on Windows; use the platform's actual temp directory instead. +MVAPICH_TMP_HOSTFILE = os.path.join(tempfile.gettempdir(), 'deepspeed_mvapich_hostfile') ELASTIC_TRAINING_ID_DEFAULT = "123456789" diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index a8caaa51fa7e..e3cc59be26f4 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/deepspeed/profiling/flops_profiler/profiler.py @@ -370,22 +370,28 @@ def print_model_profile(self, profile_step=1, module_depth=-1, top_modules=1, de if self.ds_engine and self.ds_engine.wall_clock_breakdown(): fwd_latency = self.ds_engine.timers(FORWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 print(line_fmt.format('fwd latency: ', duration_to_string(fwd_latency))) + # fwd_latency can measure as exactly 0 for tiny models on a high-resolution timer, so guard + # the FLOPS-per-second division instead of crashing with a ZeroDivisionError. + fwd_flops_per_second = total_flops / fwd_latency if fwd_latency > 0 else 0 print( line_fmt.format('fwd FLOPS per GPU = fwd flops per GPU / fwd latency: ', - flops_to_string(total_flops / fwd_latency))) + flops_to_string(fwd_flops_per_second))) if self.ds_engine and self.ds_engine.wall_clock_breakdown(): bwd_factor = 2 + self.recompute_fwd_factor bwd_latency = self.ds_engine.timers(BACKWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 step_latency = self.ds_engine.timers(STEP_GLOBAL_TIMER).elapsed(False) / 1000.0 print(line_fmt.format('bwd latency: ', duration_to_string(bwd_latency))) + bwd_flops_per_second = bwd_factor * total_flops / bwd_latency if bwd_latency > 0 else 0 print( line_fmt.format(f'bwd FLOPS per GPU = {bwd_factor:g} * fwd flops per GPU / bwd latency: ', - flops_to_string(bwd_factor * total_flops / bwd_latency))) + flops_to_string(bwd_flops_per_second))) + fwd_bwd_latency = fwd_latency + bwd_latency + fwd_bwd_flops_per_second = (bwd_factor + 1) * total_flops / fwd_bwd_latency if fwd_bwd_latency > 0 else 0 print( line_fmt.format( f'fwd+bwd FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / (fwd+bwd latency): ', - flops_to_string((bwd_factor + 1) * total_flops / (fwd_latency + bwd_latency)))) + flops_to_string(fwd_bwd_flops_per_second))) print(line_fmt.format('step latency: ', duration_to_string(step_latency))) diff --git a/deepspeed/runtime/config.py b/deepspeed/runtime/config.py index 24e1b365f987..10ecb08e3114 100755 --- a/deepspeed/runtime/config.py +++ b/deepspeed/runtime/config.py @@ -563,8 +563,10 @@ def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None): self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) else: # HF zero.init case where there is no mpu - if "sequence_parallel_size" in config: - self.world_size = dist.get_world_size() / config["sequence_parallel_size"] + # `config` (the constructor arg) may be a file path or base64 string rather than + # a dict; the parsed dict is always `self._param_dict`. + if "sequence_parallel_size" in self._param_dict: + self.world_size = dist.get_world_size() / self._param_dict["sequence_parallel_size"] else: self.world_size = dist.get_world_size() except (RuntimeError, AssertionError, AttributeError): diff --git a/deepspeed/utils/timer.py b/deepspeed/utils/timer.py index 65506cb6df7b..5e240954acd5 100644 --- a/deepspeed/utils/timer.py +++ b/deepspeed/utils/timer.py @@ -61,7 +61,9 @@ def start(self): """Start the timer.""" assert not self.started_, f"{self.name_} timer has already been started" if self.use_host_timer: - self.start_time = time.time() + # perf_counter is monotonic and higher-resolution than time.time(), which is + # coarse (~15ms) on Windows and can make short intervals measure as 0. + self.start_time = time.perf_counter() else: event_class = get_accelerator().Event self.start_event = event_class(enable_timing=True) @@ -73,7 +75,7 @@ def stop(self, reset=False, record=False): assert self.started_, "timer is not started" event_class = get_accelerator().Event if self.use_host_timer: - self.end_time = time.time() + self.end_time = time.perf_counter() self.event_timers.append(self.end_time - self.start_time) else: event_class = get_accelerator().Event @@ -236,7 +238,7 @@ def start(self): if self.global_step_count >= self.start_step: if self.config.synchronized: get_accelerator().synchronize() - self.start_time = time.time() + self.start_time = time.perf_counter() def _is_report_boundary(self): if self.steps_per_output is None: @@ -254,7 +256,7 @@ def stop(self, global_step=False, report_speed=True): if self.start_time > 0: if self.config.synchronized: get_accelerator().synchronize() - self.end_time = time.time() + self.end_time = time.perf_counter() duration = self.end_time - self.start_time self.total_elapsed_time += duration self.step_elapsed_time += duration From 46516c591ccdfc8ac0d15a5af71f18a887db3a2c Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:34:27 -0700 Subject: [PATCH 02/11] csrc/pin_memory: port pin_memory op to Windows --- csrc/pin_memory/deepspeed_pin_tensor.cpp | 44 +++++++++++++++++++++++- csrc/pin_memory/deepspeed_pin_tensor.h | 23 ++++++++++++- csrc/pin_memory/page_alloc.cpp | 31 +++++++++++++++++ 3 files changed, 96 insertions(+), 2 deletions(-) diff --git a/csrc/pin_memory/deepspeed_pin_tensor.cpp b/csrc/pin_memory/deepspeed_pin_tensor.cpp index f97600f6cf7c..fd30c684b75d 100644 --- a/csrc/pin_memory/deepspeed_pin_tensor.cpp +++ b/csrc/pin_memory/deepspeed_pin_tensor.cpp @@ -11,15 +11,25 @@ Functionality for managing CPU tensors occupying page-locked memory. #include +#if defined(_WIN32) +#define NOMINMAX +#include +#else #include +#endif using namespace std; deepspeed_pin_tensor_t::~deepspeed_pin_tensor_t() { for (auto iter = _locked_tensors.begin(); iter != _locked_tensors.end(); ++iter) { +#if defined(_WIN32) + VirtualUnlock(iter->first, iter->second); + _aligned_free((void*)iter->first); +#else munlock(iter->first, iter->second); std::free((void*)iter->first); +#endif } _locked_tensors.clear(); } @@ -30,7 +40,7 @@ std::shared_ptr deepspeed_pin_tensor_t::shared() return mgr; } -extern "C" void* deepspeed_pin_tensor_mgr_holder() +extern "C" DS_PIN_TENSOR_EXPORT void* deepspeed_pin_tensor_mgr_holder() { // Heap-allocate the shared_ptr so its control block outlives any transient // copies made by other extensions that resolve this symbol via dlsym. @@ -71,8 +81,13 @@ bool deepspeed_pin_tensor_t::free(void* addr) std::lock_guard guard(_mutex); auto iter = _locked_tensors.find(addr); if (iter != _locked_tensors.end()) { +#if defined(_WIN32) + VirtualUnlock(addr, iter->second); + _aligned_free(addr); +#else munlock(addr, iter->second); std::free(addr); +#endif _locked_tensors.erase(iter); return true; } @@ -96,3 +111,30 @@ bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer) } return false; }; + +#if defined(_WIN32) +extern "C" DS_PIN_TENSOR_EXPORT torch::Tensor deepspeed_pin_tensor_alloc_by_scalartype( + const int64_t num_elem, + const at::ScalarType elem_type) +{ + return deepspeed_pin_tensor_t::shared()->alloc(num_elem, elem_type); +} + +extern "C" DS_PIN_TENSOR_EXPORT torch::Tensor deepspeed_pin_tensor_alloc_by_options( + const int64_t num_elem, + const torch::TensorOptions& options) +{ + return deepspeed_pin_tensor_t::shared()->alloc(num_elem, options); +} + +extern "C" DS_PIN_TENSOR_EXPORT bool deepspeed_pin_tensor_free_tensor(torch::Tensor& locked_tensor) +{ + return deepspeed_pin_tensor_t::shared()->free(locked_tensor); +} + +extern "C" DS_PIN_TENSOR_EXPORT bool deepspeed_pin_tensor_check_is_managed( + const torch::Tensor& buffer) +{ + return deepspeed_pin_tensor_t::shared()->is_managed(buffer); +} +#endif diff --git a/csrc/pin_memory/deepspeed_pin_tensor.h b/csrc/pin_memory/deepspeed_pin_tensor.h index 043e96e32142..7514b45cb573 100644 --- a/csrc/pin_memory/deepspeed_pin_tensor.h +++ b/csrc/pin_memory/deepspeed_pin_tensor.h @@ -16,6 +16,15 @@ TODO: Implement a full-featured manager that #include #include +// Windows only exposes symbols that a DLL explicitly exports, unlike POSIX shared objects which +// export everything by default; async_io/gds (see deepspeed_pin_tensor_client.cpp) need these +// visible to find them across the .pyd boundary via GetProcAddress, the Windows analogue of dlsym. +#if defined(_WIN32) +#define DS_PIN_TENSOR_EXPORT __declspec(dllexport) +#else +#define DS_PIN_TENSOR_EXPORT +#endif + struct deepspeed_pin_tensor_t { std::map _locked_tensors; std::mutex _mutex; @@ -43,4 +52,16 @@ struct deepspeed_pin_tensor_t { }; // Exported so async_io/gds can resolve the same manager across .so boundaries. -extern "C" void* deepspeed_pin_tensor_mgr_holder(); +extern "C" DS_PIN_TENSOR_EXPORT void* deepspeed_pin_tensor_mgr_holder(); + +#if defined(_WIN32) +extern "C" DS_PIN_TENSOR_EXPORT torch::Tensor deepspeed_pin_tensor_alloc_by_scalartype( + const int64_t num_elem, + const at::ScalarType elem_type); +extern "C" DS_PIN_TENSOR_EXPORT torch::Tensor deepspeed_pin_tensor_alloc_by_options( + const int64_t num_elem, + const torch::TensorOptions& options); +extern "C" DS_PIN_TENSOR_EXPORT bool deepspeed_pin_tensor_free_tensor(torch::Tensor& locked_tensor); +extern "C" DS_PIN_TENSOR_EXPORT bool deepspeed_pin_tensor_check_is_managed( + const torch::Tensor& buffer); +#endif diff --git a/csrc/pin_memory/page_alloc.cpp b/csrc/pin_memory/page_alloc.cpp index ff2f769b0626..c7f370977f4c 100644 --- a/csrc/pin_memory/page_alloc.cpp +++ b/csrc/pin_memory/page_alloc.cpp @@ -4,8 +4,13 @@ #include "page_alloc.h" +#if defined(_WIN32) +#define NOMINMAX +#include +#else #include #include +#endif #include #include #include @@ -13,6 +18,31 @@ void* ds_page_aligned_alloc(const int64_t size, const bool lock) { +#if defined(_WIN32) + SYSTEM_INFO sys_info; + GetSystemInfo(&sys_info); + void* ptr = _aligned_malloc(static_cast(size), sys_info.dwPageSize); + if (ptr == nullptr) { return nullptr; } + + if (lock == false) { return ptr; } + + // VirtualLock only succeeds up to the process's working-set quota, so that quota must be + // raised by the allocation size first -- unlike POSIX mlock, which has no such per-process cap. + SIZE_T min_ws = 0, max_ws = 0; + GetProcessWorkingSetSize(GetCurrentProcess(), &min_ws, &max_ws); + SetProcessWorkingSetSize(GetCurrentProcess(), + min_ws + static_cast(size), + max_ws + static_cast(size)); + + if (!VirtualLock(ptr, static_cast(size))) { + std::cerr << "VirtualLock failed to allocate " << size << " bytes with error no " + << GetLastError() << std::endl; + _aligned_free(ptr); + return nullptr; + } + + return ptr; +#else void* ptr; int retval; @@ -31,4 +61,5 @@ void* ds_page_aligned_alloc(const int64_t size, const bool lock) } return ptr; +#endif } From 785329857b3e2d39ec89202689ef2ca02bd3a169 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:35:31 -0700 Subject: [PATCH 03/11] csrc/aio: port async_io op to Windows --- csrc/aio/common/deepspeed_aio_common.cpp | 118 ++++++++++++++---- csrc/aio/common/deepspeed_aio_common.h | 4 +- csrc/aio/common/deepspeed_aio_types.cpp | 8 ++ csrc/aio/common/deepspeed_aio_types.h | 26 +++- csrc/aio/common/deepspeed_aio_utils.cpp | 40 +++++- csrc/aio/common/deepspeed_aio_utils.h | 14 ++- csrc/aio/py_lib/deepspeed_aio_op_desc.cpp | 2 +- csrc/aio/py_lib/deepspeed_aio_op_desc.h | 4 +- csrc/aio/py_lib/deepspeed_cpu_op.cpp | 2 +- csrc/aio/py_lib/deepspeed_cpu_op.h | 2 +- .../py_lib/deepspeed_pin_tensor_client.cpp | 86 +++++++++++-- csrc/aio/py_lib/deepspeed_py_aio.cpp | 14 ++- csrc/aio/py_lib/deepspeed_py_copy.cpp | 6 +- csrc/aio/py_lib/deepspeed_py_io_handle.cpp | 47 +++++-- csrc/aio/py_lib/deepspeed_py_io_handle.h | 6 +- 15 files changed, 312 insertions(+), 67 deletions(-) diff --git a/csrc/aio/common/deepspeed_aio_common.cpp b/csrc/aio/common/deepspeed_aio_common.cpp index 9d7ff5093017..77dde91a1a73 100644 --- a/csrc/aio/common/deepspeed_aio_common.cpp +++ b/csrc/aio/common/deepspeed_aio_common.cpp @@ -11,17 +11,20 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. #include #include +#if !defined(_WIN32) #include #include #include #include #include -#include #include +#endif +#include #include #include #include +#include #include #include #include @@ -40,9 +43,11 @@ using namespace std::chrono; static const std::string c_library_name = "deepspeed_aio"; +#if !defined(_WIN32) static void _report_aio_statistics(const char* tag, const std::vector>& latencies) __attribute__((unused)); +#endif static void _report_aio_statistics(const char* tag, const std::vector>& latencies) @@ -68,6 +73,38 @@ static void _get_aio_latencies(std::vector>& raw_l std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size(); } +#if defined(_WIN32) +// Windows has no io_submit/io_pgetevents equivalent, so "submit" here actually performs a +// blocking ReadFile/WriteFile immediately; completion is then a no-op (see _do_io_complete). +static void _win_submit_one(io_request_t* req) +{ + char* buf = static_cast(req->_buf); + int64_t offset = req->_offset; + int64_t remaining = req->_nbytes; + + while (remaining > 0) { + OVERLAPPED ov = {}; + ov.Offset = static_cast(offset & 0xffffffff); + ov.OffsetHigh = static_cast(offset >> 32); + DWORD bytes_transferred = 0; + const BOOL ok = + req->_read_op + ? ReadFile(req->_fd, buf, static_cast(remaining), &bytes_transferred, &ov) + : WriteFile(req->_fd, buf, static_cast(remaining), &bytes_transferred, &ov); + if (!ok || bytes_transferred == 0) { + const auto error_code = GetLastError(); + report_file_error( + "", req->_read_op ? "ReadFile" : "WriteFile", static_cast(error_code)); + assert(ok && bytes_transferred > 0); + return; + } + buf += bytes_transferred; + offset += bytes_transferred; + remaining -= bytes_transferred; + } +} +#endif + static void _do_io_submit_singles(const int64_t n_iocbs, const int64_t iocb_index, std::unique_ptr& aio_ctxt, @@ -75,9 +112,14 @@ static void _do_io_submit_singles(const int64_t n_iocbs, { for (auto i = 0; i < n_iocbs; ++i) { const auto st = std::chrono::high_resolution_clock::now(); +#if defined(_WIN32) + _win_submit_one(aio_ctxt->_iocbs[i]); +#else const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, 1, aio_ctxt->_iocbs.data() + i); + assert(submit_ret > 0); +#endif submit_times.push_back(std::chrono::high_resolution_clock::now() - st); -#if DEBUG_DS_AIO_SUBMIT_PERF +#if DEBUG_DS_AIO_SUBMIT_PERF && !defined(_WIN32) printf("submit(usec) %f io_index=%lld buf=%p len=%lu off=%llu \n", submit_times.back().count() * 1e6, iocb_index, @@ -85,7 +127,6 @@ static void _do_io_submit_singles(const int64_t n_iocbs, aio_ctxt->_iocbs[i]->u.c.nbytes, aio_ctxt->_iocbs[i]->u.c.offset); #endif - assert(submit_ret > 0); } } @@ -95,9 +136,14 @@ static void _do_io_submit_block(const int64_t n_iocbs, std::vector>& submit_times) { const auto st = std::chrono::high_resolution_clock::now(); +#if defined(_WIN32) + for (auto i = 0; i < n_iocbs; ++i) { _win_submit_one(aio_ctxt->_iocbs[i]); } +#else const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, n_iocbs, aio_ctxt->_iocbs.data()); + assert(submit_ret > 0); +#endif submit_times.push_back(std::chrono::high_resolution_clock::now() - st); -#if DEBUG_DS_AIO_SUBMIT_PERF +#if DEBUG_DS_AIO_SUBMIT_PERF && !defined(_WIN32) printf("submit(usec) %f io_index=%lld nr=%lld buf=%p len=%lu off=%llu \n", submit_times.back().count() * 1e6, iocb_index, @@ -106,7 +152,6 @@ static void _do_io_submit_block(const int64_t n_iocbs, aio_ctxt->_iocbs[0]->u.c.nbytes, aio_ctxt->_iocbs[0]->u.c.offset); #endif - assert(submit_ret > 0); } static int _do_io_complete(const int64_t min_completes, @@ -115,12 +160,18 @@ static int _do_io_complete(const int64_t min_completes, std::vector>& reap_times) { const auto start_time = std::chrono::high_resolution_clock::now(); +#if defined(_WIN32) + // _win_submit_one already blocked until each request finished, so every + // request handed in was completed by the time we get here. + const int64_t n_completes = max_completes; +#else int64_t n_completes = io_pgetevents(aio_ctxt->_io_ctxt, min_completes, max_completes, aio_ctxt->_io_events.data(), nullptr, nullptr); +#endif reap_times.push_back(std::chrono::high_resolution_clock::now() - start_time); assert(n_completes >= min_completes); return n_completes; @@ -264,8 +315,23 @@ void report_file_error(const char* filename, const std::string file_op, const in std::cerr << c_library_name << ": " << err_msg << std::endl; } -int open_file(const char* filename, const bool read_op) +aio_fd_t open_file(const char* filename, const bool read_op) { +#if defined(_WIN32) + // No FILE_FLAG_NO_BUFFERING (Windows' O_DIRECT analogue): it would force sector-aligned + // offsets/sizes/buffers, which the synchronous path in _win_submit_one doesn't guarantee. + const DWORD access = read_op ? GENERIC_READ : GENERIC_WRITE; + const DWORD disposition = read_op ? OPEN_EXISTING : OPEN_ALWAYS; + const auto fd = CreateFileA( + filename, access, FILE_SHARE_READ, nullptr, disposition, FILE_ATTRIBUTE_NORMAL, nullptr); + if (fd == INVALID_HANDLE_VALUE) { + const auto error_code = GetLastError(); + const auto error_msg = read_op ? " open for read " : " open for write "; + report_file_error(filename, error_msg, static_cast(error_code)); + return AIO_INVALID_FD; + } + return fd; +#else const int flags = read_op ? (O_RDONLY | O_DIRECT) : (O_WRONLY | O_CREAT | O_DIRECT); #if defined(__ENABLE_CANN__) int* flags_ptr = (int*)&flags; @@ -280,32 +346,38 @@ int open_file(const char* filename, const bool read_op) return -1; } return fd; +#endif +} + +void close_file(const aio_fd_t fd) +{ +#if defined(_WIN32) + CloseHandle(fd); +#else + close(fd); +#endif } int regular_read(const char* filename, std::vector& buffer) { - const auto fd = open(filename, O_RDONLY, 0600); - assert(fd != -1); - struct stat fs; - const auto result = fstat(fd, &fs); - assert(result != -1); - int64_t num_bytes = fs.st_size; + // Uses the C stdio API (rather than POSIX open/read) so this helper works unchanged on + // Windows, which has no open()/read() pair with the same semantics. + auto* file = fopen(filename, "rb"); + assert(file != nullptr); + assert(fseek(file, 0, SEEK_END) == 0); + const int64_t num_bytes = ftell(file); + assert(num_bytes >= 0); + assert(fseek(file, 0, SEEK_SET) == 0); buffer.resize(num_bytes); - int64_t read_bytes = 0; - auto r = 0; - do { - const auto buffer_ptr = buffer.data() + read_bytes; - const auto bytes_to_read = num_bytes - read_bytes; - r = read(fd, buffer_ptr, bytes_to_read); - read_bytes += r; - } while (r > 0); + const int64_t read_bytes = + num_bytes == 0 ? 0 : static_cast(fread(buffer.data(), 1, num_bytes, file)); if (read_bytes != num_bytes) { - std::cerr << "read error " << " read_bytes (read) = " << read_bytes - << " num_bytes (fstat) = " << num_bytes << std::endl; + std::cerr << "read error " << " read_bytes (fread) = " << read_bytes + << " num_bytes (ftell) = " << num_bytes << std::endl; } assert(read_bytes == num_bytes); - close(fd); + fclose(file); return 0; } diff --git a/csrc/aio/common/deepspeed_aio_common.h b/csrc/aio/common/deepspeed_aio_common.h index aa4e49f4f4ed..3e5761c4c545 100644 --- a/csrc/aio/common/deepspeed_aio_common.h +++ b/csrc/aio/common/deepspeed_aio_common.h @@ -26,7 +26,9 @@ void do_aio_operation_overlap(const bool read_op, deepspeed_aio_config_t* config, deepspeed_aio_perf_t* perf); -int open_file(const char* filename, const bool read_op); +aio_fd_t open_file(const char* filename, const bool read_op); + +void close_file(const aio_fd_t fd); void report_file_error(const char* filename, const std::string file_op, const int error_code); diff --git a/csrc/aio/common/deepspeed_aio_types.cpp b/csrc/aio/common/deepspeed_aio_types.cpp index 5e34a61065d4..bfb5eb4de65e 100644 --- a/csrc/aio/common/deepspeed_aio_types.cpp +++ b/csrc/aio/common/deepspeed_aio_types.cpp @@ -61,16 +61,24 @@ aio_context::aio_context(const int block_size, const int queue_depth) { _block_size = block_size; _queue_depth = queue_depth; +#if defined(_WIN32) + for (auto i = 0; i < queue_depth; ++i) { _iocbs.push_back(new io_request_t()); } +#else for (auto i = 0; i < queue_depth; ++i) { _iocbs.push_back((struct iocb*)calloc(1, sizeof(struct iocb))); } _io_events.resize(queue_depth); io_queue_init(queue_depth, &_io_ctxt); +#endif } aio_context::~aio_context() { +#if defined(_WIN32) + for (auto& req : _iocbs) { delete req; } +#else for (auto& iocb : _iocbs) { free(iocb); } _io_events.resize(0); io_queue_release(_io_ctxt); +#endif } diff --git a/csrc/aio/common/deepspeed_aio_types.h b/csrc/aio/common/deepspeed_aio_types.h index ce6a4e5cdfa7..16ca7be7334d 100644 --- a/csrc/aio/common/deepspeed_aio_types.h +++ b/csrc/aio/common/deepspeed_aio_types.h @@ -7,7 +7,12 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. */ +#if defined(_WIN32) +#define NOMINMAX +#include +#else #include +#endif #include #include @@ -15,6 +20,23 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. using namespace std; +#if defined(_WIN32) +using aio_fd_t = HANDLE; +const aio_fd_t AIO_INVALID_FD = INVALID_HANDLE_VALUE; +struct win_aio_request { + aio_fd_t _fd; + void* _buf; + size_t _nbytes; + int64_t _offset; + bool _read_op; +}; +using io_request_t = win_aio_request; +#else +using aio_fd_t = int; +const aio_fd_t AIO_INVALID_FD = -1; +using io_request_t = struct iocb; +#endif + struct deepspeed_aio_latency_t { double _min_usec; double _max_usec; @@ -48,9 +70,11 @@ struct deepspeed_aio_config_t { }; struct aio_context { +#if !defined(_WIN32) io_context_t _io_ctxt; std::vector _io_events; - std::vector _iocbs; +#endif + std::vector _iocbs; int _block_size; int _queue_depth; diff --git a/csrc/aio/common/deepspeed_aio_utils.cpp b/csrc/aio/common/deepspeed_aio_utils.cpp index 8c85e1c9c116..693457911007 100644 --- a/csrc/aio/common/deepspeed_aio_utils.cpp +++ b/csrc/aio/common/deepspeed_aio_utils.cpp @@ -17,7 +17,7 @@ using namespace std; const int c_block_size = 128 * 1024; const int c_io_queue_depth = 8; -io_xfer_ctxt::io_xfer_ctxt(const int fd, +io_xfer_ctxt::io_xfer_ctxt(const aio_fd_t fd, const int64_t file_offset, const int64_t buffer_offset, const int64_t num_bytes, @@ -33,7 +33,7 @@ io_xfer_ctxt::io_xfer_ctxt(const int fd, io_prep_context::io_prep_context(const bool read_op, const std::unique_ptr& xfer_ctxt, const size_t block_size, - const std::vector* iocbs) + const std::vector* iocbs) : _read_op(read_op), _xfer_ctxt(xfer_ctxt), _block_size(block_size), _iocbs(iocbs) { } @@ -52,11 +52,20 @@ void io_prep_context::prep_iocbs(const int n_iocbs, if ((shift + _block_size) > num_bytes) { byte_count = num_bytes - shift; } +#if defined(_WIN32) + auto* req = _iocbs->at(i); + req->_fd = _xfer_ctxt->_fd; + req->_buf = xfer_buffer; + req->_nbytes = byte_count; + req->_offset = xfer_offset; + req->_read_op = _read_op; +#else if (_read_op) { io_prep_pread(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); } else { io_prep_pwrite(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); } +#endif } } @@ -74,7 +83,7 @@ io_prep_generator::io_prep_generator(const bool read_op, _remaining_io_blocks = _num_io_blocks; } -int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* iocbs) +int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* iocbs) { if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) { assert(static_cast(_remaining_bytes) == _remaining_io_blocks); @@ -89,11 +98,20 @@ int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* (_next_iocb_index * _block_size); const auto xfer_offset = _xfer_ctxt->_file_base_offset + (_next_iocb_index * _block_size); const auto num_bytes = min(static_cast(_block_size), _remaining_bytes); +#if defined(_WIN32) + auto* req = iocbs->at(i); + req->_fd = _xfer_ctxt->_fd; + req->_buf = xfer_buffer; + req->_nbytes = num_bytes; + req->_offset = xfer_offset; + req->_read_op = _read_op; +#else if (_read_op) { io_prep_pread(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); } else { io_prep_pwrite(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); } +#endif _remaining_bytes -= num_bytes; } _remaining_io_blocks -= actual_n_iocbs; @@ -103,16 +121,30 @@ int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* int64_t get_file_size(const char* filename, int64_t& size) { +#if defined(_WIN32) + WIN32_FILE_ATTRIBUTE_DATA attrs; + if (!GetFileAttributesExA(filename, GetFileExInfoStandard, &attrs)) { return -1; } + size = (static_cast(attrs.nFileSizeHigh) << 32) | attrs.nFileSizeLow; + return 0; +#else struct stat st; if (stat(filename, &st) == -1) { return -1; } size = st.st_size; return 0; +#endif } -int64_t get_fd_file_size(const int fd, int64_t& size) +int64_t get_fd_file_size(const aio_fd_t fd, int64_t& size) { +#if defined(_WIN32) + LARGE_INTEGER file_size; + if (!GetFileSizeEx(fd, &file_size)) { return -1; } + size = file_size.QuadPart; + return 0; +#else struct stat st; if (fstat(fd, &st) == -1) { return -1; } size = st.st_size; return 0; +#endif } diff --git a/csrc/aio/common/deepspeed_aio_utils.h b/csrc/aio/common/deepspeed_aio_utils.h index d951bbd9f158..1cca74f1bae0 100644 --- a/csrc/aio/common/deepspeed_aio_utils.h +++ b/csrc/aio/common/deepspeed_aio_utils.h @@ -13,12 +13,14 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. #include #include +#if !defined(_WIN32) #include #include #include #include #include #include +#endif #include #include @@ -29,13 +31,13 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. #include struct io_xfer_ctxt { - const int _fd; + const aio_fd_t _fd; const int64_t _file_base_offset; const int64_t _buffer_base_offset; const void* _mem_buffer; const int64_t _num_bytes; - io_xfer_ctxt(const int fd, + io_xfer_ctxt(const aio_fd_t fd, const int64_t file_offset, const int64_t buffer_offset, const int64_t num_bytes, @@ -46,12 +48,12 @@ struct io_prep_context { const bool _read_op; const std::unique_ptr& _xfer_ctxt; const size_t _block_size; - const std::vector* _iocbs; + const std::vector* _iocbs; io_prep_context(const bool read_op, const std::unique_ptr& xfer_ctxt, const size_t block_size, - const std::vector* iocbs); + const std::vector* iocbs); void prep_iocbs(const int n_iocbs, const size_t num_bytes, @@ -73,8 +75,8 @@ struct io_prep_generator { const std::unique_ptr& xfer_ctxt, const size_t block_size); - int prep_iocbs(const int n_iocbs, std::vector* iocbs); + int prep_iocbs(const int n_iocbs, std::vector* iocbs); }; int64_t get_file_size(const char* filename, int64_t& size); -int64_t get_fd_file_size(const int fd, int64_t& size); +int64_t get_fd_file_size(const aio_fd_t fd, int64_t& size); diff --git a/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp b/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp index c0d801a1ca11..72095980dea0 100644 --- a/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp +++ b/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp @@ -49,7 +49,7 @@ void warn_consumer_ssd_writes() noexcept io_op_desc_t::io_op_desc_t(const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const int intra_op_parallelism, const bool validate, diff --git a/csrc/aio/py_lib/deepspeed_aio_op_desc.h b/csrc/aio/py_lib/deepspeed_aio_op_desc.h index 20b4d7813e2c..dae0bb7bc585 100644 --- a/csrc/aio/py_lib/deepspeed_aio_op_desc.h +++ b/csrc/aio/py_lib/deepspeed_aio_op_desc.h @@ -14,7 +14,7 @@ void warn_consumer_ssd_writes() noexcept; struct io_op_desc_t { const bool _read_op; torch::Tensor _buffer; - int _fd; + aio_fd_t _fd; std::string _filename; const int _intra_op_parallelism; const int64_t _num_bytes_per_thread; @@ -24,7 +24,7 @@ struct io_op_desc_t { io_op_desc_t(const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const int intra_op_parallelism, const bool validate, diff --git a/csrc/aio/py_lib/deepspeed_cpu_op.cpp b/csrc/aio/py_lib/deepspeed_cpu_op.cpp index 6fb9d1b9a1f1..e636407a7ef8 100644 --- a/csrc/aio/py_lib/deepspeed_cpu_op.cpp +++ b/csrc/aio/py_lib/deepspeed_cpu_op.cpp @@ -12,7 +12,7 @@ cpu_op_desc_t::cpu_op_desc_t( const std::shared_ptr& pinned_tensor_mgr, const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const int intra_op_parallelism, const bool validate, diff --git a/csrc/aio/py_lib/deepspeed_cpu_op.h b/csrc/aio/py_lib/deepspeed_cpu_op.h index 5490e38bbf3e..694a5a3cb536 100644 --- a/csrc/aio/py_lib/deepspeed_cpu_op.h +++ b/csrc/aio/py_lib/deepspeed_cpu_op.h @@ -16,7 +16,7 @@ struct cpu_op_desc_t : io_op_desc_t { cpu_op_desc_t(const std::shared_ptr& pinned_tensor_mgr, const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const int intra_op_parallelism, const bool validate, diff --git a/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp b/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp index fb9cd05aa852..b072a5391875 100644 --- a/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp +++ b/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp @@ -5,24 +5,94 @@ /* Resolve the process-wide pin-tensor manager exported by the pin_memory op. The manager is compiled only into pin_memory; async_io/gds must load that op -first (see AsyncIOBuilder.load) so this symbol is visible via RTLD_GLOBAL. +first (see AsyncIOBuilder.load) so this symbol can be found here -- via +RTLD_GLOBAL/dlsym on POSIX, or by scanning loaded modules on Windows. */ #include "deepspeed_pin_tensor.h" -#include #include +#if defined(_WIN32) +#define NOMINMAX +#define PSAPI_VERSION 2 +#include +#include +#include +#else +#include +#endif + +using holder_fn_t = void* (*)(); + +namespace { +const char* const kMissingSymbolMessage = + "DeepSpeed pin_memory op must be loaded before async_io/gds (missing exported pin_memory " + "symbol). Load PinMemoryBuilder first."; +} // namespace + +#if defined(_WIN32) +namespace { +FARPROC find_exported_symbol(const char* name) +{ + DWORD needed = 0; + EnumProcessModules(GetCurrentProcess(), nullptr, 0, &needed); + std::vector modules(needed / sizeof(HMODULE)); + if (!EnumProcessModules(GetCurrentProcess(), modules.data(), needed, &needed)) { + return nullptr; + } + for (auto mod : modules) { + if (auto* addr = GetProcAddress(mod, name)) { return addr; } + } + return nullptr; +} + +template +FnPtr resolve(const char* name) +{ + auto* addr = find_exported_symbol(name); + if (addr == nullptr) { throw std::runtime_error(kMissingSymbolMessage); } + return reinterpret_cast(addr); +} +} // namespace +#endif + std::shared_ptr deepspeed_pin_tensor_t::shared() { - using holder_fn_t = void* (*)(); +#if defined(_WIN32) + auto* fn = resolve("deepspeed_pin_tensor_mgr_holder"); +#else auto* fn = reinterpret_cast(dlsym(RTLD_DEFAULT, "deepspeed_pin_tensor_mgr_holder")); - if (fn == nullptr) { - throw std::runtime_error( - "DeepSpeed pin_memory op must be loaded before async_io/gds (missing " - "deepspeed_pin_tensor_mgr_holder). Load PinMemoryBuilder first."); - } + if (fn == nullptr) { throw std::runtime_error(kMissingSymbolMessage); } +#endif auto* holder = static_cast*>(fn()); return *holder; } + +#if defined(_WIN32) +torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type) +{ + using fn_t = torch::Tensor (*)(const int64_t, const at::ScalarType); + return resolve("deepspeed_pin_tensor_alloc_by_scalartype")(num_elem, elem_type); +} + +torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, + const torch::TensorOptions& options) +{ + using fn_t = torch::Tensor (*)(const int64_t, const torch::TensorOptions&); + return resolve("deepspeed_pin_tensor_alloc_by_options")(num_elem, options); +} + +bool deepspeed_pin_tensor_t::free(torch::Tensor& locked_tensor) +{ + using fn_t = bool (*)(torch::Tensor&); + return resolve("deepspeed_pin_tensor_free_tensor")(locked_tensor); +} + +bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer) +{ + using fn_t = bool (*)(const torch::Tensor&); + return resolve("deepspeed_pin_tensor_check_is_managed")(buffer); +} +#endif diff --git a/csrc/aio/py_lib/deepspeed_py_aio.cpp b/csrc/aio/py_lib/deepspeed_py_aio.cpp index 8dfef8c3def3..60e2e1d82b7b 100644 --- a/csrc/aio/py_lib/deepspeed_py_aio.cpp +++ b/csrc/aio/py_lib/deepspeed_py_aio.cpp @@ -11,11 +11,13 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. #include #include +#if !defined(_WIN32) #include #include #include #include #include +#endif #include #include @@ -49,7 +51,7 @@ int deepspeed_py_aio_write(const torch::Tensor& buffer, deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } warn_consumer_ssd_writes(); auto write_buffer = (char*)buffer.data_ptr(); @@ -67,7 +69,7 @@ int deepspeed_py_aio_write(const torch::Tensor& buffer, const std::chrono::duration aio_time = std::chrono::high_resolution_clock::now() - start_time; - close(fd); + close_file(fd); if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } @@ -89,14 +91,18 @@ int deepspeed_py_aio_read(torch::Tensor& buffer, const auto start_time = std::chrono::high_resolution_clock::now(); int64_t num_file_bytes; if (-1 == get_file_size(filename, num_file_bytes)) { +#if defined(_WIN32) + const auto error_code = static_cast(GetLastError()); +#else const auto error_code = errno; +#endif report_file_error(filename, " fstat for read", error_code); return -1; } deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } auto read_buffer = (char*)buffer.data_ptr(); assert(static_cast(buffer.nbytes()) == num_file_bytes); @@ -113,7 +119,7 @@ int deepspeed_py_aio_read(torch::Tensor& buffer, const std::chrono::duration aio_time = std::chrono::high_resolution_clock::now() - start_time; - close(fd); + close_file(fd); if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); } diff --git a/csrc/aio/py_lib/deepspeed_py_copy.cpp b/csrc/aio/py_lib/deepspeed_py_copy.cpp index f5480e9d9d83..97ff915786c0 100644 --- a/csrc/aio/py_lib/deepspeed_py_copy.cpp +++ b/csrc/aio/py_lib/deepspeed_py_copy.cpp @@ -46,8 +46,12 @@ static void helper_memcpy_1(float* dest, float* src, size_t param_size) #endif if (param_size > rounded_size) { + // MSVC's OpenMP loop index must be signed; GCC/Clang accept either. + const auto tail_end = static_cast(param_size); #pragma omp parallel for - for (size_t k = rounded_size; k < param_size; k++) { dest[k] = src[k]; } + for (int64_t k = static_cast(rounded_size); k < tail_end; k++) { + dest[k] = src[k]; + } } } diff --git a/csrc/aio/py_lib/deepspeed_py_io_handle.cpp b/csrc/aio/py_lib/deepspeed_py_io_handle.cpp index 2630822339a1..20ce6a979f02 100644 --- a/csrc/aio/py_lib/deepspeed_py_io_handle.cpp +++ b/csrc/aio/py_lib/deepspeed_py_io_handle.cpp @@ -9,9 +9,17 @@ Functionality for swapping optimizer tensors to/from (NVMe) storage devices. #include "deepspeed_py_io_handle.h" #include +#if defined(_WIN32) +#include +#endif #include "deepspeed_aio_op_desc.h" +// Windows disk sector sizes are commonly 4096 bytes vs. the 512-byte default assumed on Linux. +#if defined(_WIN32) +#define O_DIRECT_ALIGNMENT 4096 +#else #define O_DIRECT_ALIGNMENT 512 +#endif using namespace std; @@ -23,7 +31,11 @@ static bool is_valid_bytes_to_read(const char* filename, { int64_t num_file_bytes; if (-1 == get_file_size(filename, num_file_bytes)) { +#if defined(_WIN32) + const auto error_code = static_cast(GetLastError()); +#else const auto error_code = errno; +#endif report_file_error(filename, " fstat for read", error_code); return false; } @@ -95,14 +107,18 @@ int deepspeed_io_handle_t::read(torch::Tensor& buffer, int64_t num_file_bytes; if (-1 == get_file_size(filename, num_file_bytes)) { +#if defined(_WIN32) + const auto error_code = static_cast(GetLastError()); +#else const auto error_code = errno; +#endif report_file_error(filename, " fstat for read", error_code); return -1; } assert(static_cast(buffer.nbytes()) == num_file_bytes); const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } auto read_buffer = (char*)buffer.data_ptr(); std::unique_ptr xfer_ctxt( @@ -114,7 +130,7 @@ int deepspeed_io_handle_t::read(torch::Tensor& buffer, do_aio_operation_sequential(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); } - close(fd); + close_file(fd); const std::chrono::duration aio_time = std::chrono::high_resolution_clock::now() - start_time; @@ -136,7 +152,7 @@ int deepspeed_io_handle_t::write(const torch::Tensor& buffer, const auto start_time = std::chrono::high_resolution_clock::now(); const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } warn_consumer_ssd_writes(); auto write_buffer = (char*)buffer.data_ptr(); @@ -152,7 +168,7 @@ int deepspeed_io_handle_t::write(const torch::Tensor& buffer, const std::chrono::duration aio_time = std::chrono::high_resolution_clock::now() - start_time; - close(fd); + close_file(fd); if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } @@ -220,7 +236,7 @@ int deepspeed_io_handle_t::_wait_locked() completed_op->finish(); - if (!completed_op->_filename.empty()) { close(completed_op->_fd); } + if (!completed_op->_filename.empty()) { close_file(completed_op->_fd); } --_num_pending_ops; ++num_completed_ops; @@ -245,7 +261,7 @@ bool deepspeed_io_handle_t::_is_valid_parallel_aio_op(const bool read_op, const std::shared_ptr deepspeed_io_handle_t::_create_io_op_desc( const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const int64_t file_offset) @@ -261,7 +277,7 @@ std::shared_ptr deepspeed_io_handle_t::_create_io_op_desc( } int deepspeed_io_handle_t::_pread(const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const bool async, @@ -289,13 +305,13 @@ int deepspeed_io_handle_t::pread(const torch::Tensor& buffer, if (!_is_valid_parallel_aio_op(true, buffer_bytes)) { return -1; } const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } return _pread(buffer, fd, filename, validate, async, file_offset); } int deepspeed_io_handle_t::_pwrite(const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const bool async, @@ -321,7 +337,7 @@ int deepspeed_io_handle_t::pwrite(const torch::Tensor& buffer, if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } + if (fd == AIO_INVALID_FD) { return -1; } return _pwrite(buffer, fd, filename, validate, async, file_offset); } @@ -361,7 +377,16 @@ int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer, const auto num_write_bytes = static_cast(buffer.nbytes()); if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } - return _pwrite(buffer, fd, nullptr, false, true, file_offset); +#if defined(_WIN32) + // Callers pass a Python-level (CRT) fd from os.open(), but aio_fd_t is a Win32 HANDLE on this + // platform; _get_osfhandle recovers the HANDLE backing that CRT fd. + const auto handle = reinterpret_cast(_get_osfhandle(fd)); + if (handle == AIO_INVALID_FD) { return -1; } +#else + const auto handle = fd; +#endif + + return _pwrite(buffer, handle, nullptr, false, true, file_offset); } at::Tensor deepspeed_io_handle_t::new_cpu_locked_tensor(const int64_t num_elem, diff --git a/csrc/aio/py_lib/deepspeed_py_io_handle.h b/csrc/aio/py_lib/deepspeed_py_io_handle.h index a6a9f17e10f8..c459416e9d30 100644 --- a/csrc/aio/py_lib/deepspeed_py_io_handle.h +++ b/csrc/aio/py_lib/deepspeed_py_io_handle.h @@ -93,14 +93,14 @@ struct deepspeed_io_handle_t { bool _is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes); int _pread(const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const bool async, const int64_t file_offset); int _pwrite(const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const bool async, @@ -108,7 +108,7 @@ struct deepspeed_io_handle_t { virtual std::shared_ptr _create_io_op_desc(const bool read_op, const torch::Tensor& buffer, - const int fd, + const aio_fd_t fd, const char* filename, const bool validate, const int64_t file_offset); From 152393db7d7115c1655f886c881ee9703068931b Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:36:26 -0700 Subject: [PATCH 04/11] op_builder: support building CPU ops with MSVC on Windows --- op_builder/builder.py | 8 +++++++- op_builder/cpu/async_io.py | 13 ++++++++++++- op_builder/cpu/builder.py | 3 +++ op_builder/cpu/comm.py | 11 +++++++++-- op_builder/cpu/no_impl.py | 6 ++++++ op_builder/pin_memory.py | 10 ---------- op_builder/pin_memory_load.py | 2 +- 7 files changed, 38 insertions(+), 15 deletions(-) diff --git a/op_builder/builder.py b/op_builder/builder.py index 465aa4399a90..c94c247c6a0c 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -567,7 +567,13 @@ def load(self, verbose=False): if torch.cuda.is_available() and isinstance(self, CUDAOpBuilder): self.validate_torch_op_version(torch_info) - op_module = importlib.import_module(self.absolute_name()) + try: + op_module = importlib.import_module(self.absolute_name()) + except ImportError: + # installed_ops can say an op was prebuilt even though its extension module + # isn't actually importable here (e.g. a Windows wheel missing that .pyd); + # fall back to compiling it on the fly instead of raising. + return self.jit_load(verbose) __class__._loaded_ops[self.name] = op_module return op_module else: diff --git a/op_builder/cpu/async_io.py b/op_builder/cpu/async_io.py index e6f0afca698d..2f4ae0542154 100644 --- a/op_builder/cpu/async_io.py +++ b/op_builder/cpu/async_io.py @@ -5,6 +5,7 @@ import shutil import subprocess +import sys from .builder import CPUOpBuilder @@ -37,8 +38,10 @@ def include_paths(self): return ['csrc/aio/py_lib', 'csrc/aio/common', 'csrc/pin_memory'] def cxx_args(self): - # -O0 for improved debugging, since performance is bound by I/O args = super().cxx_args() + if sys.platform == "win32": + return args + # -O0 for improved debugging, since performance is bound by I/O import torch TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): @@ -48,6 +51,10 @@ def cxx_args(self): return args def extra_ldflags(self): + if sys.platform == "win32": + # -laio has no meaning on Windows (no libaio); MSVC OpenMP is enabled via + # the /openmp compile flag in cxx_args(), not a linker flag. + return [] return ['-laio', '-fopenmp'] def check_for_libaio_pkg(self): @@ -79,6 +86,10 @@ def load(self, verbose=False): return super().load(verbose=verbose) def is_compatible(self, verbose=False): + if sys.platform == "win32": + # No libaio on Windows; the op is built on native Win32 file I/O instead. + return super().is_compatible(verbose) + # Check for the existence of libaio by using distutils # to compile and link a test program that calls io_submit, # which is a function provided by libaio that is used in the async_io op. diff --git a/op_builder/cpu/builder.py b/op_builder/cpu/builder.py index d881842ad0b1..d4fa4aee906d 100644 --- a/op_builder/cpu/builder.py +++ b/op_builder/cpu/builder.py @@ -4,6 +4,7 @@ # DeepSpeed Team import os +import sys try: # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed @@ -30,6 +31,8 @@ def builder(self): return cpp_ext def cxx_args(self): + if sys.platform == "win32": + return ['/O2', '/openmp', '/EHsc', '/W3', '/std:c++20'] args = ['-O3', '-g', '-Wno-reorder'] CPU_ARCH = self.cpu_arch() SIMD_WIDTH = self.simd_width() diff --git a/op_builder/cpu/comm.py b/op_builder/cpu/comm.py index bd1a3213b0b5..19504c5d59b2 100644 --- a/op_builder/cpu/comm.py +++ b/op_builder/cpu/comm.py @@ -30,6 +30,10 @@ def cxx_args(self): return ['-O2', '-fopenmp'] def is_compatible(self, verbose=False): + if sys.platform == "win32": + if verbose: + self.warning(f"{self.NAME} requires oneCCL, which is not verified on Windows.") + return False # TODO: add soft compatibility check for private binary release. # a soft check, as in we know it can be trivially changed. return super().is_compatible(verbose) @@ -63,10 +67,13 @@ def include_paths(self): return includes def cxx_args(self): + if sys.platform == "win32": + return super().cxx_args() return ['-O2', '-fopenmp'] def is_compatible(self, verbose=False): - # The shared-memory kernels use Linux-only APIs, so let other platforms fall back to gloo. - if sys.platform != 'linux': + # The shared-memory kernels have Linux and Windows implementations; other + # platforms (e.g. macOS) fall back to gloo. + if sys.platform not in ('linux', 'win32'): return False return super().is_compatible(verbose) diff --git a/op_builder/cpu/no_impl.py b/op_builder/cpu/no_impl.py index 69d114a9f1c0..6d467882b316 100644 --- a/op_builder/cpu/no_impl.py +++ b/op_builder/cpu/no_impl.py @@ -17,6 +17,12 @@ def __init__(self, name=None): def absolute_name(self): return f'deepspeed.ops.comm.{self.NAME}_op' + def is_compatible(self, verbose=False): + # Has no sources to build; only exists so callers get a clear error from + # load() instead of NoneType, so it must never be picked up for a real + # (pre-)compile -- an empty ext_modules entry crashes MSVC's linker. + return False + def load(self, verbose=True): raise ValueError("This op had not been implemented on CPU backend.") diff --git a/op_builder/pin_memory.py b/op_builder/pin_memory.py index 339dccffdbee..50365e6e6ca0 100644 --- a/op_builder/pin_memory.py +++ b/op_builder/pin_memory.py @@ -2,8 +2,6 @@ # DeepSpeed Team -import sys - from .builder import OpBuilder from .pin_memory_load import load_pin_memory_module @@ -18,14 +16,6 @@ def __init__(self): def absolute_name(self): return f'deepspeed.ops.pin_memory.{self.NAME}_op' - def is_compatible(self, verbose=False): - # The allocator relies on POSIX mlock/posix_memalign, which are unavailable on Windows. - if sys.platform == "win32": - if verbose: - self.warning(f"{self.NAME} is only supported on POSIX platforms, not Windows.") - return False - return super().is_compatible(verbose) - def sources(self): return [ 'csrc/pin_memory/page_alloc.cpp', diff --git a/op_builder/pin_memory_load.py b/op_builder/pin_memory_load.py index 179b4b38ab3e..b52e00037bbf 100644 --- a/op_builder/pin_memory_load.py +++ b/op_builder/pin_memory_load.py @@ -17,6 +17,6 @@ def load_pin_memory_module(builder, verbose=False): module = super(type(builder), builder).load(verbose=verbose) setattr(sys, _SYS_CACHE_ATTR, module) so_path = getattr(module, "__file__", None) - if so_path: + if so_path and sys.platform != "win32": ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL) return module From 05b87c52af0627009b229ccd9d4521d4ce887926 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:37:45 -0700 Subject: [PATCH 05/11] csrc/cpu/comm: port shared-memory comm backend to Windows --- csrc/cpu/comm/arm64/shm.h | 7 +++ csrc/cpu/comm/shm.cpp | 101 +++++++++++++++++++++++++++++--- csrc/cpu/comm/shm_interface.cpp | 4 +- csrc/cpu/comm/x86_64/shm.h | 24 +++++--- 4 files changed, 118 insertions(+), 18 deletions(-) diff --git a/csrc/cpu/comm/arm64/shm.h b/csrc/cpu/comm/arm64/shm.h index f6bdc41c6d43..4f76d3e0b520 100644 --- a/csrc/cpu/comm/arm64/shm.h +++ b/csrc/cpu/comm/arm64/shm.h @@ -13,6 +13,13 @@ #include #include +#if defined(_MSC_VER) && !defined(__clang__) +// MSVC's arm64_neon.h defines the float16x4_t/float16x8_t vector types but, +// unlike GCC/Clang's arm_neon.h, never typedefs the scalar float16_t that +// vld1_f16/vst1_f16 below cast a pointer to. +typedef unsigned short float16_t; +#endif + // 128 bits = 16 bytes -> fits 8 fp16/bf16 or 4 fp32 elements. static int vector_length_in_bytes = 16; // When widening fp16/bf16 -> fp32, 4 elements fit in one 128-bit register. diff --git a/csrc/cpu/comm/shm.cpp b/csrc/cpu/comm/shm.cpp index 976c1f91c463..091c7aaa5bee 100644 --- a/csrc/cpu/comm/shm.cpp +++ b/csrc/cpu/comm/shm.cpp @@ -6,15 +6,28 @@ #include #include +#include +#include +#include +#include "shm.h" + +#ifdef _WIN32 +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#else #include -#include #include -#include "shm.h" +#endif #if defined(__riscv) #define TARGET_RISCV 1 #include "riscv64/shm.h" -#elif defined(__aarch64__) +#elif defined(__aarch64__) || defined(_M_ARM64) #define TARGET_ARM 1 #include "arm64/shm.h" #else @@ -40,13 +53,69 @@ enum coll_state { }; // SHM building blocks +#ifdef _WIN32 +#define SHM_INVALID_DESCRIPTOR NULL +#else +#define SHM_INVALID_DESCRIPTOR (-1) +#endif + struct SharedData { const char* name; +#ifdef _WIN32 + HANDLE descriptor; +#else int descriptor; +#endif void* bytes; size_t nbytes; }; +#ifdef _WIN32 +void shared_open(SharedData* data, const char* name, size_t nbytes) +{ + HANDLE h = OpenFileMappingA(FILE_MAP_ALL_ACCESS, FALSE, name); + if (h != NULL) { + void* bytes = MapViewOfFile(h, FILE_MAP_ALL_ACCESS, 0, 0, nbytes); + data->name = name; + data->descriptor = h; + data->bytes = bytes; + data->nbytes = nbytes; + } else { + if (GetLastError() != ERROR_FILE_NOT_FOUND) { + // don't print if shm can not be found because we want to loop over from + // caller again until the other ranks created the shm + printf("shared_open %s failed, error=%lu\n", name, GetLastError()); + } + errno = ENOENT; + data->descriptor = SHM_INVALID_DESCRIPTOR; + } +} + +void shared_create(SharedData* data, const char* name, void* bytes, size_t nbytes) +{ + HANDLE h = + CreateFileMappingA(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, (DWORD)nbytes, name); + if (h != NULL) { + void* mapped = MapViewOfFile(h, FILE_MAP_ALL_ACCESS, 0, 0, nbytes); + memcpy(mapped, bytes, nbytes); + data->name = name; + data->descriptor = h; + data->bytes = mapped; + data->nbytes = nbytes; + } else { + printf("shared_create %s failed\n", name); + data->descriptor = SHM_INVALID_DESCRIPTOR; + } +} + +void shared_close(SharedData* data) +{ + if (data->descriptor != SHM_INVALID_DESCRIPTOR) { + UnmapViewOfFile(data->bytes); + CloseHandle(data->descriptor); + } +} +#else void shared_open(SharedData* data, const char* name, size_t nbytes) { int d = shm_open(name, O_RDWR, S_IRUSR | S_IWUSR); @@ -62,7 +131,7 @@ void shared_open(SharedData* data, const char* name, size_t nbytes) // caller again until the other ranks created the shm printf("shared_open %s failed, errno=%d\n", name, errno); } - data->descriptor = -1; + data->descriptor = SHM_INVALID_DESCRIPTOR; } } @@ -78,11 +147,26 @@ void shared_create(SharedData* data, const char* name, void* bytes, size_t nbyte void shared_close(SharedData* data) { - if (data->descriptor != -1) { + if (data->descriptor != SHM_INVALID_DESCRIPTOR) { munmap(data->bytes, data->nbytes); shm_unlink(data->name); } } +#endif + +static std::string shm_owner_id() +{ +#ifdef _WIN32 + // USERNAME is queried instead of GetUserNameA() so this doesn't pull in an + // extra link dependency on Advapi32.lib just to namespace shm names by user. + char name[256]; + DWORD len = GetEnvironmentVariableA("USERNAME", name, sizeof(name)); + if (len == 0 || len >= sizeof(name)) { return "unknown"; } + return std::string(name, len); +#else + return std::to_string(getuid()); +#endif +} static int world_size; @@ -347,11 +431,12 @@ void shm_initialize(int size, int rank, char* addr_string, char* port_string) char shm_name_prefix[NAME_BUF_SIZE]; char shm_name[NAME_BUF_SIZE]; + auto owner_id = shm_owner_id(); snprintf(shm_name_prefix, NAME_BUF_SIZE, - "%s_%d_%s_%s", + "%s_%s_%s_%s", SHM_BUFFER_NAME, - getuid(), + owner_id.c_str(), addr_string, port_string); // create shared workspace for SHM based allreduce @@ -380,7 +465,7 @@ void shm_initialize(int size, int rank, char* addr_string, char* port_string) // printf("open %s, %d\n", shm_name, rank); do { shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace)); - } while (allreduce_buffer.descriptor == -1 && errno == ENOENT); + } while (allreduce_buffer.descriptor == SHM_INVALID_DESCRIPTOR && errno == ENOENT); workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes; workspace[i] = workspace_buf_other; } else { diff --git a/csrc/cpu/comm/shm_interface.cpp b/csrc/cpu/comm/shm_interface.cpp index 5be5cb799a7b..010f52d3747d 100644 --- a/csrc/cpu/comm/shm_interface.cpp +++ b/csrc/cpu/comm/shm_interface.cpp @@ -39,9 +39,9 @@ void initialize(int size, int rank) is_initialized = 1; auto addr_string = std::getenv("MASTER_ADDR"); - if (addr_string == NULL) { addr_string = ""; } + if (addr_string == NULL) { addr_string = const_cast(""); } auto port_string = std::getenv("MASTER_PORT"); - if (port_string == NULL) { port_string = ""; } + if (port_string == NULL) { port_string = const_cast(""); } if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } } diff --git a/csrc/cpu/comm/x86_64/shm.h b/csrc/cpu/comm/x86_64/shm.h index 9b02eb3779cd..62316b4ff404 100644 --- a/csrc/cpu/comm/x86_64/shm.h +++ b/csrc/cpu/comm/x86_64/shm.h @@ -5,14 +5,22 @@ #include -inline __m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +// MSVC has no equivalent of GCC/Clang's function-multiversioning target attribute: it exposes all +// AVX-512 intrinsics regardless of /arch, so the attribute is simply unneeded there. +#if defined(_MSC_VER) && !defined(__clang__) +#define DS_ATTRIBUTE_TARGET_AVX512BW +#else +#define DS_ATTRIBUTE_TARGET_AVX512BW __attribute__((target("avx512bw"))) +#endif + +inline __m512 cvt_bf16_to_fp32(const __m256i src) DS_ATTRIBUTE_TARGET_AVX512BW; inline __m512 cvt_bf16_to_fp32(const __m256i src) { auto y = _mm512_cvtepu16_epi32(src); return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); } -inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_bf16(const __m512 src) DS_ATTRIBUTE_TARGET_AVX512BW; inline __m256i cvt_fp32_to_bf16(const __m512 src) { __m512i value = _mm512_castps_si512(src); @@ -33,10 +41,10 @@ inline __m256i cvt_fp32_to_bf16(const __m512 src) return _mm512_cvtusepi32_epi16(t_value); } -inline __m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_fp16_to_fp32(const __m256i src) DS_ATTRIBUTE_TARGET_AVX512BW; inline __m512 cvt_fp16_to_fp32(const __m256i src) { return _mm512_cvtph_ps(src); } -inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_fp16(const __m512 src) DS_ATTRIBUTE_TARGET_AVX512BW; inline __m256i cvt_fp32_to_fp16(const __m512 src) { return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); @@ -49,13 +57,13 @@ inline __m256i cvt_fp32_to_fp16(const __m512 src) static int vector_length_in_bytes = 32; void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); + DS_ATTRIBUTE_TARGET_AVX512BW; void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); + DS_ATTRIBUTE_TARGET_AVX512BW; void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); + DS_ATTRIBUTE_TARGET_AVX512BW; -void parallel_memcpy(void* to, void* from, size_t n_bytes) __attribute__((target("avx512bw"))); +void parallel_memcpy(void* to, void* from, size_t n_bytes) DS_ATTRIBUTE_TARGET_AVX512BW; #define VLOAD_U8(X) _mm256_loadu_si256((__m256i*)(X)) #define VLOAD_U16(X) _mm256_loadu_si256((__m256i*)(X)) From 397faa47b37752b10b337b54270368683a360339 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:39:14 -0700 Subject: [PATCH 06/11] tests: fix cross-test state leaks and a fragile dynamo backend assumption --- tests/unit/compile/test_zero3_grad_dtype.py | 15 ++- tests/unit/utils/test_pin_memory_tracker.py | 20 ++-- tests/unit/v1/compile/util.py | 114 ++++++++++---------- 3 files changed, 85 insertions(+), 64 deletions(-) diff --git a/tests/unit/compile/test_zero3_grad_dtype.py b/tests/unit/compile/test_zero3_grad_dtype.py index 02915fd0d1e1..4e5c8caf8935 100644 --- a/tests/unit/compile/test_zero3_grad_dtype.py +++ b/tests/unit/compile/test_zero3_grad_dtype.py @@ -128,7 +128,20 @@ def forward(self, value): def backend(_gm, _inputs): backend_calls.append(None) param.data = torch.ones(2, 2, device=param.device) - return lambda _weight, value: (value, ) + + # dynamo doesn't guarantee the compiled graph is called with a fixed (weight, value) + # signature -- e.g. a constant-folded weight may be dropped from args entirely -- so + # the fake backend must locate the actual forward input rather than assume its position. + def compiled_forward(*args): + tensor_args = [arg for arg in args if torch.is_tensor(arg)] + for arg in tensor_args: + if arg.shape == value.shape: + return (arg, ) + if len(tensor_args) == 1: + return (tensor_args[0], ) + raise AssertionError(f"compiled callable invoked without the forward input: {args}") + + return compiled_forward patch_fake_tensor() compiled = torch.compile(module, backend=backend) diff --git a/tests/unit/utils/test_pin_memory_tracker.py b/tests/unit/utils/test_pin_memory_tracker.py index 2946801fd603..1cdd05e1944b 100644 --- a/tests/unit/utils/test_pin_memory_tracker.py +++ b/tests/unit/utils/test_pin_memory_tracker.py @@ -3,6 +3,7 @@ import logging +import pytest import torch from deepspeed.utils.pin_memory_tracker import ( @@ -13,8 +14,18 @@ ) -def test_track_accumulates_and_resets(): +@pytest.fixture(autouse=True) +def _reset_tracker(): + # _tracker is a module-level singleton shared across every test here, so + # without a guaranteed reset on both sides, a test running out of file + # order (or a prior test failing before its own cleanup) leaks byte counts + # into whatever runs next. + _tracker.reset() + yield _tracker.reset() + + +def test_track_accumulates_and_resets(): track_pinned_memory(100) track_pinned_memory(2**30) assert _tracker._bytes == 100 + 2**30 @@ -24,10 +35,8 @@ def test_track_accumulates_and_resets(): def test_summary_does_not_raise(): - _tracker.reset() track_pinned_memory(2**30) pinned_memory_summary("unit-test") - _tracker.reset() def test_fmt_bytes(): @@ -41,11 +50,9 @@ def test_torch_tensor_nbytes_is_consistent(): t = torch.zeros(1024, dtype=torch.float32) track_pinned_memory(t.nbytes) assert _tracker._bytes == 4096 - _tracker.reset() def test_checkpoint_thresholds_double_from_32gb(): - _tracker.reset() gb = 1024**3 assert _tracker._next_checkpoint == 32 * gb track_pinned_memory(30 * gb) # below 32 GB -> no crossing @@ -54,13 +61,11 @@ def test_checkpoint_thresholds_double_from_32gb(): assert _tracker._next_checkpoint == 64 * gb track_pinned_memory(100 * gb) # 140 GB -> crosses 64 and 128 in one call assert _tracker._next_checkpoint == 256 * gb - _tracker.reset() def test_checkpoint_emits_info(caplog): # The DeepSpeed logger does not propagate, so flip propagation so caplog # (root-based) can observe the checkpoint INFO records. - _tracker.reset() ds_logger = logging.getLogger("DeepSpeed") old_prop = ds_logger.propagate ds_logger.propagate = True @@ -74,4 +79,3 @@ def test_checkpoint_emits_info(caplog): assert "32" in checkpoints[0] finally: ds_logger.propagate = old_prop - _tracker.reset() diff --git a/tests/unit/v1/compile/util.py b/tests/unit/v1/compile/util.py index c61554091a1c..cc81a813fd39 100644 --- a/tests/unit/v1/compile/util.py +++ b/tests/unit/v1/compile/util.py @@ -72,25 +72,28 @@ def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None, rto ys = [torch.randn_like(x) for x in xs] target_losses = [] - for x, y in zip(xs, ys): - baseline_loss = baseline_engine(x, y) - target_loss = target_engine(x, y) - target_losses.append(target_loss.detach().float().item()) + # Always destroy the engines, even if a loss/parameter assertion fails, otherwise their + # process groups leak into whatever test runs next. + try: + for x, y in zip(xs, ys): + baseline_loss = baseline_engine(x, y) + target_loss = target_engine(x, y) + target_losses.append(target_loss.detach().float().item()) - allclose_on_all_ranks(baseline_loss, target_loss, "Loss values are not close.", rtol=RTOL, atol=ATOL) + allclose_on_all_ranks(baseline_loss, target_loss, "Loss values are not close.", rtol=RTOL, atol=ATOL) - baseline_engine.backward(baseline_loss) - target_engine.backward(target_loss) + baseline_engine.backward(baseline_loss) + target_engine.backward(target_loss) - baseline_engine.step() - target_engine.step() + baseline_engine.step() + target_engine.step() - with GatheredParameters(target_engine.parameters()): - for p1, p2 in zip(baseline_engine.parameters(), target_engine.parameters()): - allclose_on_all_ranks(p1, p2, "Parameters are not equal.", rtol=RTOL, atol=ATOL) - - baseline_engine.destroy() - target_engine.destroy() + with GatheredParameters(target_engine.parameters()): + for p1, p2 in zip(baseline_engine.parameters(), target_engine.parameters()): + allclose_on_all_ranks(p1, p2, "Parameters are not equal.", rtol=RTOL, atol=ATOL) + finally: + baseline_engine.destroy() + target_engine.destroy() # Returned so callers can compare two compiled configurations far more tightly. return target_losses @@ -184,49 +187,50 @@ def _ulysses_attn_forward(module, # Train both engines in lockstep; compare the losses at the final step. ul_loss = autosp_loss = None - for i in range(iterations): - torch.manual_seed(42 + i) - full_ids = torch.randint(0, vocab_size, (1, seq_length), device=device) - - # Ulysses: each rank processes its own shard. - shard_ids = full_ids[:, sp_rank * chunk:(sp_rank + 1) * chunk] - shard_pos = torch.arange(sp_rank * chunk, (sp_rank + 1) * chunk, device=device).unsqueeze(0) - shard_mask = torch.ones(1, chunk, device=device, dtype=torch.long) - ul_out = ulysses_engine(input_ids=shard_ids, - labels=shard_ids, - position_ids=shard_pos, - attention_mask=shard_mask) - # Average per-shard losses across SP ranks to get the full-sequence loss. - ul_loss = ul_out.loss.clone() - dist.all_reduce(ul_loss, group=sp_group) - ul_loss = ul_loss / sp_size - - # AutoSP: full sequence. dynamic=True makes all shapes symbolic, so mark_dynamic - # is not needed; only the tag attributes that the autosp pass uses are set here. - autosp_ids = full_ids.clone() - autosp_lbl = autosp_ids.clone() - autosp_pos = torch.arange(seq_length, device=device).unsqueeze(0) - autosp_msk = torch.ones(1, seq_length, device=device, dtype=torch.long) - autosp_ids.tag = autosp_constants.AUTOSP_INPUT_ID_KEY - autosp_lbl.tag = autosp_constants.AUTOSP_LABEL_ID_KEY - autosp_pos.tag = autosp_constants.AUTOSP_POSITION_ID_KEY - autosp_out = autosp_engine(input_ids=autosp_ids, - labels=autosp_lbl, - position_ids=autosp_pos, - attention_mask=autosp_msk) - autosp_loss = autosp_out.loss - - ulysses_engine.backward(ul_out.loss) - ulysses_engine.step() - autosp_engine.backward(autosp_loss) - autosp_engine.step() + try: + for i in range(iterations): + torch.manual_seed(42 + i) + full_ids = torch.randint(0, vocab_size, (1, seq_length), device=device) + + # Ulysses: each rank processes its own shard. + shard_ids = full_ids[:, sp_rank * chunk:(sp_rank + 1) * chunk] + shard_pos = torch.arange(sp_rank * chunk, (sp_rank + 1) * chunk, device=device).unsqueeze(0) + shard_mask = torch.ones(1, chunk, device=device, dtype=torch.long) + ul_out = ulysses_engine(input_ids=shard_ids, + labels=shard_ids, + position_ids=shard_pos, + attention_mask=shard_mask) + # Average per-shard losses across SP ranks to get the full-sequence loss. + ul_loss = ul_out.loss.clone() + dist.all_reduce(ul_loss, group=sp_group) + ul_loss = ul_loss / sp_size + + # AutoSP: full sequence. dynamic=True makes all shapes symbolic, so mark_dynamic + # is not needed; only the tag attributes that the autosp pass uses are set here. + autosp_ids = full_ids.clone() + autosp_lbl = autosp_ids.clone() + autosp_pos = torch.arange(seq_length, device=device).unsqueeze(0) + autosp_msk = torch.ones(1, seq_length, device=device, dtype=torch.long) + autosp_ids.tag = autosp_constants.AUTOSP_INPUT_ID_KEY + autosp_lbl.tag = autosp_constants.AUTOSP_LABEL_ID_KEY + autosp_pos.tag = autosp_constants.AUTOSP_POSITION_ID_KEY + autosp_out = autosp_engine(input_ids=autosp_ids, + labels=autosp_lbl, + position_ids=autosp_pos, + attention_mask=autosp_msk) + autosp_loss = autosp_out.loss + + ulysses_engine.backward(ul_out.loss) + ulysses_engine.step() + autosp_engine.backward(autosp_loss) + autosp_engine.step() + finally: + ulysses_engine.destroy() + del ALL_ATTENTION_FUNCTIONS["ulyssess"] + autosp_engine.destroy() allclose_on_all_ranks(autosp_loss, ul_loss, "AutoSP and Ulysses losses are not close.", rtol=RTOL, atol=ATOL) - ulysses_engine.destroy() - del ALL_ATTENTION_FUNCTIONS["ulyssess"] - autosp_engine.destroy() - def create_gm_nodes(batch_size: int = 1, seq_len: int = 16): """ From 0b663916b3159674b9c9567f471741f0643f67ff Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:40:25 -0700 Subject: [PATCH 07/11] tests: make the unit test harness runnable on Windows --- .../checkpoint/test_autotp_uc_checkpoint.py | 3 +- tests/unit/comm/test_dist.py | 3 +- tests/unit/common.py | 35 ++++++++++++++----- tests/unit/launcher/test_user_args.py | 11 +++++- tests/unit/runtime/test_ds_initialize.py | 2 ++ tests/unit/v1/nvme/test_aio.py | 1 + 6 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py index 58bbe065c6e8..e82a8e599ae1 100644 --- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py +++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py @@ -905,7 +905,8 @@ def _launch_procs(self, num_procs, init_method): # CPU/gloo test: the number of processes is not bound to the accelerator's # device_count() (CPU sockets), so bypass the base class's per-device gate # that would otherwise skip a 4-process test on a single-socket CPU box. - torch.multiprocessing.set_start_method('forkserver', force=True) + start_method = 'forkserver' if 'forkserver' in torch.multiprocessing.get_all_start_methods() else 'spawn' + torch.multiprocessing.set_start_method(start_method, force=True) self._launch_daemonic_procs(num_procs, init_method) def test(self, tmpdir): diff --git a/tests/unit/comm/test_dist.py b/tests/unit/comm/test_dist.py index 0f6ca25947d6..d893dc736241 100644 --- a/tests/unit/comm/test_dist.py +++ b/tests/unit/comm/test_dist.py @@ -279,7 +279,8 @@ def _launch_procs(self, num_procs, init_method): self.non_daemonic_procs = True self.reuse_dist_env = False return self._launch_non_daemonic_procs(num_procs, init_method) - torch.multiprocessing.set_start_method('forkserver', force=True) + start_method = 'forkserver' if 'forkserver' in torch.multiprocessing.get_all_start_methods() else 'spawn' + torch.multiprocessing.set_start_method(start_method, force=True) self._launch_daemonic_procs(num_procs, init_method) def test(self): diff --git a/tests/unit/common.py b/tests/unit/common.py index de63f2d183e7..84cd18c39015 100644 --- a/tests/unit/common.py +++ b/tests/unit/common.py @@ -4,6 +4,7 @@ # DeepSpeed Team import os +import platform import re import time import inspect @@ -65,6 +66,13 @@ def get_master_port(base_port=29500, port_range_size=1000): def _get_cpu_socket_count(): + # /proc/cpuinfo (and the cat/grep/sort/wc pipeline below) is Linux-only, so + # Windows queries the physical socket count via WMI instead. + if platform.system() == "Windows": + return int( + subprocess.check_output( + ["powershell", "-Command", + "(Get-CimInstance Win32_ComputerSystem).NumberOfProcessors"]).decode().strip()) import shlex p1 = subprocess.Popen(shlex.split("cat /proc/cpuinfo"), stdout=subprocess.PIPE) p2 = subprocess.Popen(["grep", "physical id"], stdin=p1.stdout, stdout=subprocess.PIPE) @@ -204,13 +212,21 @@ def _launch_daemonic_procs(self, num_procs, init_method): try: skip_msgs = skip_msgs_async.get(self.exec_timeout) except mp.TimeoutError: - # Shortcut to exit pytest in the case of a hanged test. This - # usually means an environment error and the rest of tests will - # hang (causing super long unit test runtimes) - pytest.exit("Test hanged, exiting", returncode=1) - finally: - # Regardless of the outcome, ensure proper teardown - # Tear down distributed environment and close process pools + # A hung worker can't respond to the graceful _dist_destroy RPC + # that _close_pool relies on, so terminate the pool directly here + # instead of exiting the whole session: under xdist, pytest.exit() + # kills this worker's channel to the controller in a way that + # surfaces as an INTERNALERROR for the entire run, even though + # only this one test actually hung. + pool.terminate() + pool.join() + if self.reuse_dist_env: + self._pool_cache.pop(num_procs, None) + pytest.fail("Test hanged and was terminated after exceeding the execution timeout") + except BaseException: + self._close_pool(pool, num_procs) + raise + else: self._close_pool(pool, num_procs) # If we skipped a test, propagate that to this process @@ -287,8 +303,9 @@ def _launch_procs(self, num_procs, init_method): if os.environ.get('DS_DISABLE_REUSE_DIST_ENV', '0') == '1': self.reuse_dist_env = False - # Set start method to `forkserver` (or `fork`) - mp.set_start_method('forkserver', force=True) + # Set start method to `forkserver` (or `fork`). Windows only supports `spawn`. + start_method = 'forkserver' if 'forkserver' in mp.get_all_start_methods() else 'spawn' + mp.set_start_method(start_method, force=True) if self.non_daemonic_procs: self._launch_non_daemonic_procs(num_procs, init_method) diff --git a/tests/unit/launcher/test_user_args.py b/tests/unit/launcher/test_user_args.py index fd1489803812..d00e87ffa211 100644 --- a/tests/unit/launcher/test_user_args.py +++ b/tests/unit/launcher/test_user_args.py @@ -4,7 +4,9 @@ # DeepSpeed Team import pytest +import shutil import subprocess +import sys from types import SimpleNamespace @@ -65,12 +67,19 @@ def dummy_runner(): def test_user_args(cmd, multi_node): if multi_node and get_accelerator().device_name() == "cpu": pytest.skip("CPU accelerator does not support this test yet") - p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + # On Windows the deepspeed console script is a .bat file, which CreateProcess + # cannot launch directly (unlike POSIX exec) -- it needs cmd.exe as a shell. + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=(sys.platform == "win32")) out, err = p.communicate() assert "ARG PARSE SUCCESS" in out.decode("utf-8"), f"User args not parsed correctly: {err.decode('utf-8')}" def test_bash_string_args(tmpdir, user_script_fp): + if shutil.which("bash") is None: + pytest.skip("bash is not available on this system") + if sys.platform == "win32": + pytest.skip("bare 'deepspeed' is not invocable from bash on Windows (entry point is a .bat file)") + bash_script = f""" ARGS="--prompt 'DeepSpeed is the best'" echo ${{ARGS}}|xargs deepspeed --num_nodes 1 --num_gpus 1 {user_script_fp} diff --git a/tests/unit/runtime/test_ds_initialize.py b/tests/unit/runtime/test_ds_initialize.py index 80fd622d5534..a287710eac69 100644 --- a/tests/unit/runtime/test_ds_initialize.py +++ b/tests/unit/runtime/test_ds_initialize.py @@ -28,6 +28,8 @@ @pytest.mark.parametrize('method', ['spawn', 'fork', 'forkserver']) def test_start_method_safety(method): import torch.multiprocessing as mp + if method not in mp.get_all_start_methods(): + pytest.skip(f"start method '{method}' is not available on this platform") mp.set_start_method(method, force=True) diff --git a/tests/unit/v1/nvme/test_aio.py b/tests/unit/v1/nvme/test_aio.py index 6ceff7c289f4..d0b32cc44b8e 100644 --- a/tests/unit/v1/nvme/test_aio.py +++ b/tests/unit/v1/nvme/test_aio.py @@ -330,6 +330,7 @@ def test_write(self, tmpdir, use_cuda_pinned_tensor, async_queue, use_unpinned_t class TestAsyncFileOffset(DistributedTest): world_size = 1 + @pytest.mark.skipif(not hasattr(os, "O_DIRECT"), reason="O_DIRECT is not supported on this platform") @pytest.mark.parametrize('use_fd', [False, True]) def test_offset_write(self, tmpdir, file_partitions, use_cuda_pinned_tensor, use_fd): From 41e1fbc0bf35bee732f3c519151bad403844b201 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:41:15 -0700 Subject: [PATCH 08/11] csrc: use signed loop indices in OpenMP for-loops for MSVC --- csrc/adam/cpu_adam_impl.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index 6c81f301ef5e..48974ebe109d 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/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(t); + const auto offset_signed = static_cast(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(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; From 77db904233a3521cc6a80f8876aae293d8e2c991 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:41:50 -0700 Subject: [PATCH 09/11] build_win.bat: enable AIO and pin_memory op builds --- build_win.bat | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/build_win.bat b/build_win.bat index 9a35cc51b6a7..676ffe015863 100644 --- a/build_win.bat +++ b/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 From ab8a48d0edb338e2188665fb92978b1f940a8d47 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:42:14 -0700 Subject: [PATCH 10/11] Add Windows wheel builds to the release workflow and a Windows CI job --- .github/workflows/release.yml | 97 +++++++++++++++++++-- .github/workflows/setup-win-venv/action.yml | 64 ++++++++++++++ .github/workflows/windows-torch-latest.yml | 49 +++++++++++ README.md | 7 ++ requirements/requirements-win-dev.txt | 8 ++ setup.py | 1 + 6 files changed, 220 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/setup-win-venv/action.yml create mode 100644 .github/workflows/windows-torch-latest.yml create mode 100644 requirements/requirements-win-dev.txt diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 00736d1c93fb..1dda9d5e8b5d 100644 --- a/.github/workflows/release.yml +++ b/.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 diff --git a/.github/workflows/setup-win-venv/action.yml b/.github/workflows/setup-win-venv/action.yml new file mode 100644 index 000000000000..1b57ce4216d8 --- /dev/null +++ b/.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 diff --git a/.github/workflows/windows-torch-latest.yml b/.github/workflows/windows-torch-latest.yml new file mode 100644 index 000000000000..f5918580afbf --- /dev/null +++ b/.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 diff --git a/README.md b/README.md index fb60b992dd8b..fdf5a2b94a61 100755 --- a/README.md +++ b/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. diff --git a/requirements/requirements-win-dev.txt b/requirements/requirements-win-dev.txt new file mode 100644 index 000000000000..9edb80fa4076 --- /dev/null +++ b/requirements/requirements-win-dev.txt @@ -0,0 +1,8 @@ +accelerate +mup +pre-commit>=3.2.0 +pytest>=7.2.0,<8.4.0 +pytest-forked +pytest-randomly +pytest-xdist +transformers>=4.51.3 diff --git a/setup.py b/setup.py index d5fd60f842b9..98ad95a1c30c 100755 --- a/setup.py +++ b/setup.py @@ -84,6 +84,7 @@ def get_env_if_set(key, default: typing.Any = ""): '1bit_mpi': fetch_requirements('requirements/requirements-1bit-mpi.txt'), 'readthedocs': fetch_requirements('requirements/requirements-readthedocs.txt'), 'dev': fetch_requirements('requirements/requirements-dev.txt'), + 'win-dev': fetch_requirements('requirements/requirements-win-dev.txt'), 'autotuning': fetch_requirements('requirements/requirements-autotuning.txt'), 'autotuning_ml': fetch_requirements('requirements/requirements-autotuning-ml.txt'), 'sparse': fetch_requirements('requirements/requirements-sparse_pruning.txt'), From 244dd3e463774ab66dc779582ba91f633554caf4 Mon Sep 17 00:00:00 2001 From: nithin-aashik-mcw Date: Mon, 14 Sep 2026 22:48:21 -0700 Subject: [PATCH 11/11] csrc: fix psapi.h/windows.h include order in pin_tensor_client --- csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp b/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp index b072a5391875..c55d0e1ab807 100644 --- a/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp +++ b/csrc/aio/py_lib/deepspeed_pin_tensor_client.cpp @@ -16,8 +16,8 @@ RTLD_GLOBAL/dlsym on POSIX, or by scanning loaded modules on Windows. #if defined(_WIN32) #define NOMINMAX #define PSAPI_VERSION 2 -#include #include +#include #include #else #include