huggingface / huggingface/peft

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

Open
#3,619 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
21.7k
Forks
2.5k
Avg merge
4d 12h
Merged PRs (30d)
59

Description

### 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"`:

```python
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:

```text
StopIteration
```

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

```python
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:

```python
{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.

Contributor guide

Open the contributing guide

Research direction

Start at get_layer_device_map in peft.utils.integrations and run the configuration-only reproduction with hf_device_map={"": "cpu"}. Add focused unit coverage for CPU-only root and per-layer maps; done means the two-layer root map resolves to {0: "cpu", 1: "cpu"} without changing accelerator-backed behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.