feat: enable meta device targets in TransformerBridge device_map

tests/acceptance/model_bridge/test_multi_gpu_bridge.py

@@ -92,10 +92,11 @@ def test_disk_value_in_device_map_rejected(self):
         with pytest.raises(ValueError, match="not supported yet"):
             resolve_device_map(None, bad, None)
 
-    def test_meta_value_in_device_map_rejected(self):
+    def test_meta_value_in_device_map_allowed(self):
         bad: Dict[str, Union[str, int]] = {"transformer.h.0": "meta"}
-        with pytest.raises(ValueError, match="not supported yet"):
-            resolve_device_map(None, bad, None)
+        dm, mm = resolve_device_map(None, bad, None)
+        assert dm is bad
+        assert mm is None
 
 
 class TestFindEmbeddingDevice:

transformer_lens/model_bridge/sources/transformers.py

@@ -676,14 +676,15 @@ def boot(
     #   - pre-loaded hf_model that the caller dispatched themselves (e.g., device_map="auto")
     hf_device_map_post = getattr(hf_model, "hf_device_map", None)
     if hf_device_map_post:
-        # CPU placement is supported. Disk / meta offload still needs a separate Bridge
-        # hook-routing pass because wrapped subcomponents can bypass Accelerate hooks.
+        # CPU and meta device targets are supported. Disk offload still needs a
+        # separate Bridge hook-routing pass because wrapped subcomponents can bypass
+        # Accelerate's offload hooks.
         offload_values = {str(v).lower() for v in hf_device_map_post.values() if isinstance(v, str)}
-        unsupported = offload_values & {"disk", "meta"}
+        unsupported = offload_values & {"disk"}
         if unsupported:
             raise ValueError(
                 f"hf_device_map contains unsupported offload targets: {sorted(unsupported)}. "
-                "TransformerBridge currently supports CPU device_map targets, but disk / meta "
+                "TransformerBridge currently supports CPU and meta device_map targets, but disk "
                 "offload can bypass Accelerate hooks inside wrapped Bridge components."
             )
         if (

transformer_lens/utilities/multi_gpu.py

@@ -17,7 +17,7 @@
 else:
     ConfigType = Any
 
-_UNSUPPORTED_OFFLOAD_DEVICE_MAP_VALUES = {"disk", "meta"}
+_UNSUPPORTED_OFFLOAD_DEVICE_MAP_VALUES = {"disk"}
 
 AvailableDeviceMemory = list[tuple[int, int]]
 """
@@ -161,9 +161,13 @@ def resolve_device_map(
 
     Semantics:
         - Explicit ``device_map`` wins and is passed through unchanged (user-provided
-          ``max_memory`` is passed through too). CPU targets are supported; disk / meta
-          offload targets are still rejected because Bridge component wrappers can bypass
-          Accelerate's offload hooks during forward passes.
+          ``max_memory`` is passed through too). CPU and meta targets are supported;
+          disk offload targets are still rejected because Bridge component wrappers
+          can bypass Accelerate's offload hooks during forward passes.
+        - Meta device targets are unvalidated: Bridge-wrapped components may fail during
+          forward if they access offloaded params without Accelerate's
+          ``align_module_device``. Use ``meta`` only when you understand the limitations
+          (e.g. config inspection with ``load_weights=False``).
         - ``n_devices=None`` or ``1``: returns ``(None, None)`` — single-device path.
         - ``n_devices > 1``: returns ``("balanced", {0: "auto", ..., n-1: "auto"})``.
           ``"balanced"`` is accelerate's string directive for balanced layer dispatch;
@@ -191,15 +195,17 @@ def resolve_device_map(
 def _validate_device_map_values(
     device_map: Union[str, Dict[str, Union[str, int]]],
 ) -> None:
-    """Reject explicit disk / meta values in a user-supplied device_map dict."""
+    """Reject explicit disk values in a user-supplied device_map dict.
+    Meta values are allowed (unvalidated — see :func:`resolve_device_map`).
+    """
     if isinstance(device_map, str):
         return
     for key, value in device_map.items():
         normalized = str(value).lower() if isinstance(value, str) else None
         if normalized in _UNSUPPORTED_OFFLOAD_DEVICE_MAP_VALUES:
             raise ValueError(
                 f"device_map[{key!r}]={value!r} is not supported yet. TransformerBridge "
-                "currently supports CPU device_map targets, but disk / meta offload can "
+                "currently supports CPU and meta device_map targets, but disk offload can "
                 "bypass Accelerate hooks inside wrapped Bridge components."
             )