[Proposal] Add support for `cpu`, `meta`, and `disk` to TransformerBridge `device_map`

[jlarson4]

Proposal

Support cpu, meta, and disk for the device_map parameter of boot_transformers.

Motivation

Please outline the motivation for the proposal.
Is your feature request related to a problem? e.g.,"I'm always frustrated when [...]".
If this is related to another GitHub issue, please link here too.

Pitch

We want to allow cpu and other non-GPU devices to be used by torch for multi-device processing. At current, we error when any of the above types are requested in a multi-device setup.

Additional context

The dtype-cast loop in transformers.py:550-552 is the blocker for cpu support:

for param in hf_model.parameters():
    if param.is_floating_point() and param.dtype != dtype:
        param.data = param.data.to(dtype=dtype)

For meta/offloaded params, .to(dtype) either no-ops or breaks accelerate's bookkeeping. Two small fixes should unlock CPU offload:

• Skip params on meta device (offload placeholders)
• Use accelerate's align_module_device context manager when iterating offloaded params

Checklist

• I have checked that there is no similar issue in the repo (required)

[jlarson4]

See #1270 for details on the initial setup of device_map and #872 for the original discussion about this topic

[snakefood3232]

I see the blocker is the dtype-cast loop in transformers.py, lines 550-552, which affects CPU support. I'd handle this by skipping parameters on the meta device and using Accelerate's align_module_device for offloaded parameters. I can have a PR ready in 3 days. I've built similar device-mapping solutions in distributed systems before.

[jlarson4]

@snakefood3232 Excellent! I will assign this to you, and keep my eyes open for a PR. Thank you for your contribution!