Transformers test_cpu_offload tests fail with KeyError: 'xpu:0'

[dvrogozh]

With:

• (accelerate) 8039158
• huggingface/transformers@31bb662
• Pytorch XPU backend (pytorch 2.7 xpu nightly)
• Note: no IPEX

On:

• Intel Data Center GPU Max (single device)

FAILED tests/models/blip/test_modeling_blip.py::BlipVQAModelTest::test_cpu_offload - KeyError: 'xpu:0'
FAILED tests/models/blip/test_modeling_blip.py::BlipVQAModelTest::test_disk_offload_bin - KeyError: 'xpu:0'
FAILED tests/models/blip/test_modeling_blip.py::BlipVQAModelTest::test_disk_offload_safetensors - KeyError: 'xpu:0'
FAILED tests/models/blip/test_modeling_blip.py::BlipTextImageModelTest::test_cpu_offload - KeyError: 'xpu:0'
FAILED tests/models/blip/test_modeling_blip.py::BlipTextImageModelTest::test_disk_offload_bin - KeyError: 'xpu:0'
FAILED tests/models/blip/test_modeling_blip.py::BlipTextImageModelTest::test_disk_offload_safetensors - KeyError: 'xpu:0'
FAILED tests/models/dab_detr/test_modeling_dab_detr.py::DabDetrModelTest::test_cpu_offload - KeyError: 'xpu:0'
FAILED tests/models/dab_detr/test_modeling_dab_detr.py::DabDetrModelTest::test_disk_offload_bin - KeyError: 'xpu:0'
FAILED tests/models/dab_detr/test_modeling_dab_detr.py::DabDetrModelTest::test_disk_offload_safetensors - KeyError: 'xpu:0'
FAILED tests/models/vilt/test_modeling_vilt.py::ViltModelTest::test_cpu_offload - KeyError: 'xpu:0'
FAILED tests/models/vilt/test_modeling_vilt.py::ViltModelTest::test_disk_offload_bin - KeyError: 'xpu:0'
FAILED tests/models/vilt/test_modeling_vilt.py::ViltModelTest::test_disk_offload_safetensors - KeyError: 'xpu:0'

Transformers test_cpu_offload and few other tests fail for blip, dab_detr, roberta and vilt models running with XPU backend. Note: I can't reproduce this issue with CUDA on A10.

# ZE_AFFINITY_MASK=0 TRANSFORMERS_TEST_DEVICE_SPEC=spec.py python3 -m pytest tests/models/blip/test_modeling_blip.py::BlipVQAModelTest::test_cpu_offload
...
                    elif is_xpu_available():
                        device = f"xpu:{device}"
>               del self.tied_params_map[value_pointer][device]
E               KeyError: 'xpu:0'

../accelerate/src/accelerate/hooks.py:399: KeyError

Issue happens here:

accelerate/src/accelerate/hooks.py

Lines 397 to 399 in 8039158

	elif is_xpu_available():
	device = f"xpu:{device}"
	del self.tied_params_map[value_pointer][device]

When issue happens self.tied_params_map[value_pointer] set is empty. The trivial if condition to check whether it's not empty allows to avoid the issue and tests pass. As I noted, I don't see this issue happening for CUDA. I also see that self.tied_pointers_to_remove is populated twice with the same values for XPU and then post_forward() is also called twice in a row with the issue happening on the second pass.

CC: @SunMarc @faaany @zucchini-nlp

[dvrogozh]

The trivial if condition to check whether it's not empty allows to avoid the issue and tests pass.

See #3403 with such a fix. if condition helps to avoid KeyError.... However, I am not sure why this situation happens. I afraid that I might not have addressed actual issue and just fixed symptom.... Can someone help suggest a better fix or explain why this fix would be correct one?