`PrefixedDataset` violates the `Mapping` contract
Author: betacatslingCreated Sep 14, 2026Updated Sep 14, 2026
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 aPrefixedDatasettherefore produces keys that raiseKeyErrorwhen passed back to__getitem__, anddict(prefixed)/prefixed[key]/key in prefixedare unusable.__len__returns the length of the whole underlying mapping, not the number of entries under the prefix.
Reproduction
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
accelerateversion: 1.16.0.dev0 (main, f13f7c1)- Python: 3.11
Source: huggingface/accelerate