huggingface / huggingface/peft
[Bug] Prefix tuning crashes for CPU-only device maps
- 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
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