huggingface / huggingface/accelerate
`PrefixedDataset` violates the `Mapping` contract
- Dominant language
- Python
- Stars
- 9.9k
- Forks
- 1.5k
- Avg merge
- 5d 2h
- Merged PRs (30d)
- 27
Description
## Description
`accelerate.utils.offload.PrefixedDataset` wraps a mapping and exposes only the
keys under a given prefix. Two of its `Mapping` methods are inconsistent with
`__getitem__`:
- `__getitem__` expects *unprefixed* keys (`dataset[f"{prefix}{key}"]`), but
`__iter__` yields the raw, still-prefixed keys of the underlying mapping.
Iterating a `PrefixedDataset` therefore produces keys that raise `KeyError`
when passed back to `__getitem__`, and `dict(prefixed)` / `prefixed[key]` /
`key in prefixed` are unusable.
- `__len__` returns the length of the whole underlying mapping, not the number
of entries under the prefix.
## Reproduction
```python
from accelerate.utils import PrefixedDataset
dataset = {"block1.weight": 0, "block1.bias": 1, "block2.weight": 2}
prefixed = PrefixedDataset(dataset, "block1.")
len(prefixed) # -> 3, expected 2
list(prefixed) # -> ["block1.weight", "block1.bias"], expected ["weight", "bias"]
dict(prefixed) # -> KeyError: 'block1.block1.weight'
```
## Expected behavior
`__iter__` should yield the keys with the prefix stripped (so each yielded key
can be passed to `__getitem__`), and `__len__` should count only the entries
matching the prefix, making `PrefixedDataset` behave as a proper `Mapping`
view of the prefixed sub-dictionary.
Note: the current single use site (`attach_align_device_hook`) only calls
`weights_map[name]`, so this is a latent correctness bug rather than an active
crash, but any `Mapping` consumer (`dict(...)`, `.keys()`, `in`, `len()`)
misbehaves today.
## Environment
- `accelerate` version: 1.16.0.dev0 (main, f13f7c1)
- Python: 3.11
Contributor guide
Research direction
Start in accelerate/utils/offload.py at PrefixedDataset and inspect its __getitem__, __iter__, and __len__ methods, then check the attach_align_device_hook use site. Make iteration yield unprefixed keys and make len count only matching prefixed entries; verify that dict(), key lookup, membership, and length behave consistently with the Mapping contract.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100