#4257·accelerate

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