[distributed][checkpoint] AssertionError: Booleans mismatch: False is not True

🐛 Describe the bug

Cases:
xpu_distributed,test.distributed.checkpoint.test_state_dict_stager.TestStateDictStager,test_tensor_pinned_and_shared

please get wheels from https://github.com/intel/torch-xpu-ops/actions/runs/19563892858 or use gh download

gh run download 19563892858 --repo intel/torch-xpu-ops --name Torch-XPU-Wheel-1826-19563892858-1 --dir path --pattern "*.zip"
git clone -b distributed_2.10 https://github.com/daisyden/pytorch.git
cd pytorch
pip install -r requirements.txt
pip install pytest expecttest
pytest -v test/distributed/checkpoint/test_state_dict_stager.py::TestStateDictStager::test_tensor_pinned_and_shared

=================================== FAILURES ===================================
_ TestStateDictStager.test_tensor_pinned_and_shared (pin_memory=True, share_memory=False) _
Traceback (most recent call last):
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/unittest/case.py", line 59, in testPartExecutor
    yield
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/unittest/case.py", line 498, in subTest
    yield
  File "/home/jenkins/actions-runner/_work/torch-xpu-ops/torch-xpu-ops/pytorch/test/distributed/checkpoint/test_state_dict_stager.py", line 766, in test_tensor_pinned_and_shared
    self.assertEqual(
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 4291, in assertEqual
    raise error_metas.pop()[0].to_error(  # type: ignore[index]
AssertionError: Booleans mismatch: False is not True
Tensor pinned status should be True
_ TestStateDictStager.test_tensor_pinned_and_shared (pin_memory=True, share_memory=True) _
Traceback (most recent call last):
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/unittest/case.py", line 59, in testPartExecutor
    yield
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/unittest/case.py", line 498, in subTest
    yield
  File "/home/jenkins/actions-runner/_work/torch-xpu-ops/torch-xpu-ops/pytorch/test/distributed/checkpoint/test_state_dict_stager.py", line 766, in test_tensor_pinned_and_shared
    self.assertEqual(
  File "/tmp/xpu-tool/Python/3.10.19/x64/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 4291, in assertEqual
    raise error_metas.pop()[0].to_error(  # type: ignore[index]
AssertionError: Booleans mismatch: False is not True

Versions

https://github.com/daisyden/pytorch/tree/distributed_2.10

@madhumitha0102

Looked into the issue briefly. The test case randomly initializes two tensors on the device and then converts into cpu_state_dict based on different pin_memory and shared_memory config. The test case fails for multi-rank config, even for single rank. The single-rank case should be resolved first. @madhumitha0102 can you please confirm the failure for single-rank?

I ran it on single rank and I see this failure for single rank as well

I created a single-rank reproducer. Can you please open an internal PyTorch JIRA to further investigate?

ut2404.py

@riverliuintel

@riverliuintel we see this issue on single rank as well. Please assign it to appropriate person. Thanks

@yucai-intel take it back from distributed, if this issue also exist in single rank.

@astachowiczhabana

Hi @astachowiczhabana , I just tested the functionality of torchxpubot fix on those not yet PR created issues, not mean to take this issue, please keep your follow up on it. Thanks!

@ZhaoqiongZ

Sure @ZhaoqiongZ, no problem. I'm removing the debug comments then (my own agent session get's confused by them :)

@torchxpubot

@torchxpubot fix

Starting automated fix on this issue. Follow progress in the workflow run.

Stage 1 - Triage

single-bug, reproducer present. Target repo is upstream pytorch/pytorch
(torch/distributed/checkpoint/_state_dict_stager.py plus the XPU pinned-memory
plumbing in ATen), not this repo.

Stage 2 - Reproduce: REPRODUCED

Reproduced on the nightly XPU wheel 2.16.0.dev20261027+xpu using the upstream test
file (no third-party fork needed), on BMG (Intel(R) Arc(TM) Pro B60 Graphics),
oneAPI 2026.1.1:

python -m pytest -v test/distributed/checkpoint/test_state_dict_stager.py::TestStateDictStager::test_tensor_pinned_and_shared

SUBFAILED (pin_memory=True, share_memory=False) - Booleans mismatch: False is not True
SUBFAILED (pin_memory=True, share_memory=True)  - Booleans mismatch: False is not True
UserWarning: Ignoring pin_memory flag for checkpoint staging as pinning memory
requires CUDA, but CUDA is not available.

Stage 3 - Root cause

Two layers, matching the analysis already in this thread:

StateDictStager.__init__ gates pinning on torch.cuda.is_available() and pins
through torch.cuda._pin_memory_utils, which calls cudaHostRegister. On an
XPU-only box the gate silently drops pin_memory, so is_pinned() is False
for both failing sub-tests.
There is no XPU equivalent of cudaHostRegister in PyTorch. XPU's
isPinnedPtr() only answers sycl::get_pointer_type(...) == usm::alloc::host,
which holds only for memory SYCL itself allocated, so in-place pinning of an
already-allocated (and, for share_memory=True, shm-backed) buffer cannot be
observed.

I measured the SYCL side directly rather than assuming (standalone DPC++ program,
oneAPI 2026.1.1, BMG):

probe	result
`prepare_for_device_copy` on a `MAP_SHARED` mapping	succeeds
... on an unaligned `malloc` pointer	succeeds
`get_pointer_type` after registering	still `unknown`
device-to-host `queue.memcpy` into the registered range	correct data
`release_from_device_copy` after the host buffer is freed	succeeds
two disjoint ranges registered inside one page	succeeds

So the page-locking primitive already exists and already works on shared memory;
the only missing piece is that SYCL's USM table cannot be told about it.

Stage 4 - Fix

Have PyTorch keep its own record of host-registered ranges rather than wait for
SYCL to expose one - which is in effect what CUDA does too, since the table behind
cudaPointerGetAttributes belongs to the CUDA runtime, not to any allocator:

at::xpu::hostRegister/hostUnregister/isHostRegistered wrap
prepare_for_device_copy/release_from_device_copy and keep an ordered map of
the registered ranges (ordered, so an interior address answers too, as it does on
CUDA - XPU's copy fast path in this repo queries isPinnedPtr() with offset
pointers). XPUHooks::isPinnedPtr() consults it after the USM-type query.
exposed as torch._C._xpu_hostRegister/_xpu_hostUnregister, with
torch/xpu/_pin_memory_utils.py mirroring the CUDA module of the same name.
the stager resolves _pin_memory_utils from torch.accelerator instead of
hardcoding CUDA, which is what staging.py next to it already does. A backend
shipping no such module keeps the old warn-and-disable behaviour.

This needs no compiler update, so pin=True, share=True is fixed now rather
than waiting on CMPLRLLVM-70326. 7 files, ~115 lines, all in pytorch/pytorch;
nothing in this repo, so there is no second patch to land.

Stage 5 - Verify: PASSED

Source build torch 2.16.0a0+gitc5c9a00, TORCH_XPU_ARCH_LIST=bmg,pvc,
oneAPI 2026.1.1, BMG:

python -m pytest -v test/distributed/checkpoint/test_state_dict_stager.py

	whole file	`test_tensor_pinned_and_shared`
before	`2 failed, 16 passed`	2 of 4 sub-tests fail
after	`16 passed, 16 subtests passed`	all 4 sub-tests pass

Both runs are on the same build - the "before" column was produced by reverting
only the two Python files, so collection and binaries are identical and the
comparison is apples to apples. lintrunner -a reports no lint issues on the
changed files.

Summary

Both failing sub-tests come from one bug with two layers: checkpoint staging only
pinned memory when CUDA was available, and PyTorch had no XPU counterpart of
cudaHostRegister underneath it. The fix makes the stager resolve its pinning
helpers from torch.accelerator, and adds at::xpu::hostRegister on top of
sycl::ext::oneapi::experimental::prepare_for_device_copy with PyTorch keeping
its own ordered map of registered ranges, which isPinnedPtr() consults - the
one piece SYCL cannot provide, and the reason this thread was waiting on a new
SYCL API. Taking that route fixes the share_memory=True case too, on the
shipping 2026.1 compiler rather than a future one, so the fix no longer depends on
CMPLRLLVM-70326. Verified on a source build: the whole file goes from
2 failed, 16 passed to 16 passed, 16 subtests passed, with the before run
taken on the same binaries so nothing else moved.

One related spot was deliberately left alone as out of scope and is worth a
human's eye: torch/distributed/_state_dict_utils.py::_create_cpu_state_dict
still imports torch.cuda._pin_memory_utils directly for its
share-plus-pin path, which the same accelerator-generic treatment would fix.

Branch: agent/fix-issue-2404 (in a pytorch/pytorch checkout)
Base: c5c9a00715c54fdf9e3498c529cd8e7fe714d636
Changed files: aten/src/ATen/xpu/CachingHostAllocator.cpp,
aten/src/ATen/xpu/CachingHostAllocator.h,
aten/src/ATen/xpu/detail/XPUHooks.cpp, torch/_C/__init__.pyi.in,
torch/csrc/xpu/Module.cpp,
torch/distributed/checkpoint/_state_dict_stager.py,
torch/xpu/_pin_memory_utils.py
Nothing was pushed. The patch is exported as a workflow artifact for a human
to review and apply.

Generated by fix job.

Review

Copy into a new comment, keep ONE verdict, add a note. Two verdicts in
one comment are dropped as ambiguous.

**strong-accepted** <!-- review: strong-accepted -->
**weak-accepted** <!-- review: weak-accepted -->
**weak-rejected** <!-- review: weak-rejected -->
**strong-rejected** <!-- review: strong-rejected -->
**issue-drop** <!-- review: issue-drop -->

notes:

On a batch run add item=<pytorch issue> to say which sub-item the
verdict is about, in the bold text and in the marker -- the marker
is invisible, so the bold line is all a reader sees. One comment per
sub-item, since a bare verdict on a batch does not say which test it
means and is dropped. run=<id> only when judging a run other than the
newest, in both places too.

A CI DISABLED sub-item, item= from the fan-out table:

**strong-accepted item=197334** <!-- review: strong-accepted item=197334 -->
notes: same TensorAssign overload as upstream CUDA, fp8 dispatch list matches

Another sub-item of the same batch, judged on an earlier run:

**weak-rejected run=35171979678 item=197335** <!-- review: weak-rejected run=35171979678 item=197335 -->
notes: green only because the test is skipped, the kernel still lacks fp8

An [xpu-alignment] issue that should not have been filed -- one issue
per difference, so no item=:

**issue-drop** <!-- review: issue-drop -->
notes: accepted XPU/CUDA numerics difference, nothing to repair

[DCP] Support pin_memory in StateDictStager without CUDA

label-issue: #2404

Root cause: checkpoint stager forces pin_memory=False on XPU because it gates pinning on torch.cuda.is_available() (pytorch torch/distributed/checkpoint/_state_dict_stager.py:36), so is_pinned() stays False.

axis	value	reason
`type`	`Bug`	Native Type unset; reported AssertionError defect matches Bug heuristic.
`test`	`test: ut`	Repro runs `pytest test/distributed/checkpoint/test_state_dict_stager.py`.
`module`	`module: distributed`	Defect in DCP checkpoint stager (torch/distributed/checkpoint); matches dup #5738.
`priority`	`Medium`	1 UT case, AssertionError without crash.
`dtype`	—	No dtype in failure signature.
`dependency component`	`dependency component: oneAPI`	Preserved from existing human label (see concern: looks misattributed).
symptom	`functionality`	Functional defect: XPU tensor not pinned, assert fails.
triage	`duplicate`	Duplicate of #5738 (HIGH, close_as_duplicate): same test + same XPU-not-pinned root cause (5738 is the newer dup).

Warning

apply_labels.py flagged items that need a human:

warning: dropping label(s) not defined in labels.json: ['—']

Applied to this issue. Generated by a local label-issue run.

Navigation Menu

Repository navigation

[distributed][checkpoint] AssertionError: Booleans mismatch: False is not True #2404

🐛 Describe the bug

Versions

18 remaining items

Stage 1 - Triage

Stage 2 - Reproduce: REPRODUCED

Stage 3 - Root cause

Stage 4 - Fix

Stage 5 - Verify: PASSED

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[distributed][checkpoint] AssertionError: Booleans mismatch: False is not True #2404

Description

🐛 Describe the bug

Versions

Activity

newtdms commented on Jan 16, 2026

madhumitha0102 commented on Jan 21, 2026

newtdms commented on Jan 22, 2026

madhumitha0102 commented on Jan 26, 2026

riverliuintel commented on Feb 4, 2026

18 remaining items

ZhaoqiongZ commented on Sep 8, 2026

astachowiczhabana commented on Sep 8, 2026

ZhaoqiongZ commented on Oct 28, 2026

torchxpubot commented on Oct 28, 2026

Stage 1 - Triage

Stage 2 - Reproduce: REPRODUCED

Stage 3 - Root cause

Stage 4 - Fix

Stage 5 - Verify: PASSED

daisyden commented on Nov 11, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions