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

[aswanth-07]

System Info

• PEFT main: 9c16ee66cd4c58bd9cdf2d8b4e06c1cf8e8f8efe
• The minimal reproduction is configuration-only and requires no model download or accelerator.

Who can help?

@BenjaminBossan

Reproduction

get_layer_device_map assumes that every hf_device_map contains at least one device other than "cpu" or "disk":

from types import SimpleNamespace

from peft.utils.integrations import get_layer_device_map

model = SimpleNamespace(
    hf_device_map={"": "cpu"},
    config=SimpleNamespace(num_hidden_layers=2),
)

print(get_layer_device_map(model))

Current result:

StopIteration

The exception comes from selecting the first non-CPU, non-disk device before handling the valid single-root map:

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

This helper is reached while PEFT constructs and places the prefix-tuning cache. A model loaded with device_map="cpu", or an automatically generated map that falls back entirely to CPU and disk, therefore fails before the prefix can be used.

Hugging Face documents that a device map may place modules on "cpu" or "disk", and that supplying one device such as "cpu" maps the complete model to it:

• https://huggingface.co/docs/transformers/main_classes/model
• https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference

I searched open and closed PEFT issues and pull requests for get_layer_device_map, CPU-only prefix tuning, hf_device_map, and StopIteration. I found the original multi-device implementation in #2189 but no report or active change for an all-CPU map.

Expected behavior

A root map of {"": "cpu"} with two transformer layers should resolve to:

{0: "cpu", 1: "cpu"}

CPU and disk placements need a CPU execution fallback when no accelerator exists. The proposed scope is limited to get_layer_device_map plus focused unit coverage for CPU-only root and per-layer maps. No cache or tuner behavior would change for maps that already contain an accelerator.

Would you like me to prepare that focused PR? I will wait for explicit approval before changing code.

AI assistance was used during the source audit and duplicate search. I reviewed the reported path and reproduced the StopIteration directly from the current function.

[Cunninger]

I'll take this. get_layer_device_map() currently does next(d for d in hf_device_map.values() if d not in ["cpu", "disk"]) before handling the valid CPU-only single-root map, which raises StopIteration. I'll guard the CPU/disk-only case (or pass a default to next) and add a regression test.

[Cunninger]

PR: #3621