Repository navigation
Conversation
Implement register_host_memory on the XPU accelerator using Level Zero's zeMemAllocHost + external memory mapping (the L0 equivalent of cudaHostRegister), with posix_memalign + mlock for the native pin path. An earlier attempt that registered through the SYCL runtime did not work and is replaced by this direct L0 route. Also adds a ZeRO-3 offload e2e benchmark (registered vs unregistered) and XPU coverage in the pin-memory unit tests. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
be9da92 to
d3fbb9a
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d3fbb9afc0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if handles is None: | ||
| return None | ||
| ze, context = handles | ||
| ze.zeMemFree(context, address) |
There was a problem hiding this comment.
When Level Zero returns a non-success status from this call, unregister_host_memory still appears to succeed because the result is discarded. In that scenario NativePinnedMemory._unregister_device discards the tracked registration and unpin frees the CPU allocation (deepspeed/utils/pin_memory.py:109-118), leaving any mapping that the driver retained pointing at host pages that can be recycled; check the result and raise so the caller keeps the allocation alive.
Useful? React with 👍 / 👎.
| "Level Zero returned a non-identity mapping for host registration " | ||
| "(%#x != %#x); falling back to mlock-only pinned memory.", mapped.value, address) | ||
| return False | ||
| return rc == _ZE_RESULT_SUCCESS |
There was a problem hiding this comment.
When loader and extension discovery succeeds but zeMemAllocHost rejects a range or runs out of resources, this return silently degrades DS_PIN_MEMORY_REGISTER_DEVICE=1 to mlock-only. The false return does not enter the exception warning in NativePinnedMemory._new_locked (deepspeed/utils/pin_memory.py:73-77), unlike the missing-extension and nonidentity branches here, so users can unknowingly benchmark or train without device registration; emit a warning_once with the Level Zero result before returning false.
Useful? React with 👍 / 👎.
|
|
||
| class _FakeLoader: | ||
|
|
||
| def zeMemAllocHost(self, context, host_desc, size, alignment, out): |
There was a problem hiding this comment.
For actual XPU execution, this fake loader only writes the requested pointer back, so it cannot establish the key production assumption that a mapping created in this private Level Zero context is recognized by torch's XPU context or actually enables an asynchronous H2D/D2H copy. I checked .github/workflows/xpu-max1100.yml:67-84; that hardware workflow does not run tests/unit/v1/pin_memory, while the commit reports only monkeypatched tests, so add and execute an end-to-end native-pinned transfer or offload test on real XPU hardware before relying on this path.
AGENTS.md reference: AGENTS.md:L81-L82
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
this PR is validated on a real XPU machine. So this fake loader is only a mock. From test point of view, this v1 test better be run on each of the accelerators, rather than contain tests for different accelerators. This may require a separate sweep on our current test structure for multi-accelerator and solve together. I think this comments should not be blocking
- unregister_host_memory raises on zeMemFree failure (matching the CUDA path's check_error) so NativePinnedMemory.unpin keeps the CPU allocation alive instead of freeing pages the driver may still have mapped - register_host_memory warns (warning_once) with the Level Zero result code when zeMemAllocHost rejects a range, instead of silently degrading to mlock-only - tests: add alloc-failure degradation and unregister-failure raise cases Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
45623d1 to
c2e29e5
Compare
This comment has been minimized.
This comment has been minimized.
|
@delock Do we want to label this PR with the roadmap items? |
Level Zero maps an ext_memmap_sysmem range only for the first device of the context that creates it; device-initiated copies from any other device fail with ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY (reproduced on Arc Pro B60/B70 with a raw Level Zero program, independent of torch). DeepSpeed ranks each pick local_rank, so every rank but the first of a node crashed on the first registered copy. - register in a context restricted to the target device (zeContextCreateEx), cached per torch device index - match the L0 device by UUID instead of assuming zeDeviceGet order - track registrations (address -> loader+context) so unregister frees in the context that created the mapping Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
|
Hi @jinyouzhi , let me know if the latest commit works for you, thanks! |
c643a9b I tested this commit and it works for me. LGTM, Thanks! @delock |
Motivation
DeepSpeed's native pin-memory path (
DS_PIN_MEMORY_BACKEND=native) currently has no device registration on XPU —register_host_memory()returnsFalse, so pinned host buffers are mlock-only and device access to them does not go through device page-table mappings. This PR implements the L0 equivalent ofcudaHostRegisteron the XPU accelerator.Implementation
XPU_Accelerator.register_host_memory(): map existing host memory into the device's page tables via Level Zero's external host-memory mapping (zeMemAllocHostwith theZE_extension_external_memmap_sysmemdescriptor chain), created in our own L0 context — device page tables are shared across the driver's contexts, so buffers registered here are accessible from torch's runtime as well.unregister_host_memory()releases the mapping.Falsewhen the L0 loader or the extension is unavailable.Tests
Three monkeypatched unit tests in
tests/unit/v1/pin_memory/test_pin_memory.py(same pattern as the existing CUDA/CPU/NPU tests — no XPU hardware needed in CI):pNext-> external memmap desc), address and size passthroughFalserather than raisingOperational notes
numactl --cpunodebind=N --membind=N) is recommended when registering large ranges.ulimit -l unlimited(or an equivalent memlock limit) is required for large registrations; the DeepSpeed pin-memory path already checks and reports this.