feat: enable meta device targets in TransformerBridge device_map

[hiro-nikaitou]

Part of #1280.

Enables device targets in TransformerBridge , building on the CPU support shipped in #1459.

Changes:

• Remove from so explicit meta targets are passed to HF instead of raising at validation time.
• Update / docstrings to describe meta support and its known limitations.
• Update post-load check in — meta targets now warn only at forward time via Accelerate (no hard block at boot).
• Update to .

Limitations (documented in resolve_device_map docstring):
Meta-offloaded parameters are skipped by but may still cause forward-pass errors in Bridge-wrapped components that access offloaded tensors without Accelerate . This is safe for config inspection () and for users who understand the constraints.

Closes the portion of #1280. Disk offload remains blocked pending the hook-routing fix described in #1459.

[hiro-nikaitou]

Hi @jlarson4 👋 Just a friendly ping — PR #1493 (meta device support for #1280) has been ready since July 8 with all CI green and no conflicts. Could you take a look when you get a chance? Once this is merged I can start investigating the disk offload blocker. Thanks!

[jlarson4]

Hi @hiro-nikaitou, thanks for the contribution! I ran the enabled path end-to-end and every meta device_map scenario I tested crashes at boot, so this can't merge as-is.

I ran the following (tiny-gpt2, CPU, this branch):

from transformer_lens.model_bridge import TransformerBridge

TransformerBridge.boot_transformers("sshleifer/tiny-gpt2", device_map={"": "meta"})
# Boot FAILS: NotImplementedError: aten::equal ... Meta tensors
#   — inside HF's from_pretrained (tie_weights -> torch.equal on meta), so any
#     weight-tied model (most decoder LMs) can't even finish loading.

TransformerBridge.boot_transformers("sshleifer/tiny-gpt2", device_map={"transformer.h.1": "meta", "": "cpu"})
# Boot FAILS: KeyError 'transformer.wte.weight'  in accelerate's offload hooks,
#   *inside* align_module_device (which boot already uses for the dtype cast).
#   Meta params have no backing store, so alignment cannot materialize them.

TransformerBridge.boot_transformers("sshleifer/tiny-gpt2", device_map={"": "meta"}, load_weights=False)
# the docstring's example
# Boots OK, but params land on CPU and hf_device_map is ABSENT: 
# from_config never receives the device_map, so the meta target is
# silently ignored. Identical result to omitting device_map.

The docstring's example scenario is genuinely enabled, but the map is discarded, not honored. And its failure-mode caveat is off: failures are at boot (not forward), and align_module_device is not the missing remedy. The partial-meta case fails inside alignment, because meta tensors have nothing to align from.

Net effect of merging as-is: users lose dev's clear fail-fast ValueError and instead get an obscure aten::equal error or a deep KeyError, while the only "working" case is one where the directive is silently ignored.

You have a couple paths forward for resolving this:

1. Scope the allowance to the real use case: accept meta in the validator only when load_weights=False, keep the fail-fast rejection otherwise, and reword the docstring ("the map is not applied; params load on CPU via from_config"). Small, honest, and preserves your intent.
2. Demonstrate one end-to-end scenario where a meta target is actually honored and boots, pinned by a boot-level acceptance test (this means confronting HF's tie_weights torch.equal probe and the offload-index gap).

I lean towards option 1, but if you'd like to dig deeper & tackle option 2, you are more than welcome to. Whichever route you take, make sure to add tests that validate your changes.