Skip to content

FIX get_layer_device_map StopIteration on CPU-only device maps - #3621

Open
Cunninger wants to merge 1 commit into
huggingface:mainfrom
Cunninger:fix/cpu-only-device-map-3619
Open

FIX get_layer_device_map StopIteration on CPU-only device maps#3621
Cunninger wants to merge 1 commit into
huggingface:mainfrom
Cunninger:fix/cpu-only-device-map-3619

Conversation

@Cunninger

Copy link
Copy Markdown

Fixes #3619

get_layer_device_map() assumed every hf_device_map contains at least one device that is not "cpu" or "disk":

main_device = next(d for d in model.hf_device_map.values() if d not in ["cpu", "disk"])

A valid CPU-only map such as {"": "cpu"} (typical for CPU-only prefix-tuning) raises StopIteration before the existing single-root handler can run.

This keeps the GPU/accelerator path unchanged, and falls back to the first non-disk device (default "cpu") when the map is CPU/disk-only.

Tests

pytest tests/test_integrations.py::TestGetLayerDeviceMap -v

6 related tests in that class: 3 passed locally (CPU-only single-root map, per-layer GPU map, CPU offload mixed with GPU).

AI assistance

AI assistance was used to draft this change. I reviewed every line, reproduced the StopIteration from #3619, and ran the tests above. Coordination comment: #3619 (comment)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Prefix tuning crashes for CPU-only device maps

1 participant