huggingface / huggingface/accelerate

`PrefixedDataset` violates the `Mapping` contract

Open Beginner friendly
#4,257 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.