diff --git a/tests/acceptance/model_bridge/test_multi_gpu_bridge.py b/tests/acceptance/model_bridge/test_multi_gpu_bridge.py index dac8d3203..3119d9b74 100644 --- a/tests/acceptance/model_bridge/test_multi_gpu_bridge.py +++ b/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: diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index bce60d96f..ca06ca75e 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/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 ( diff --git a/transformer_lens/utilities/multi_gpu.py b/transformer_lens/utilities/multi_gpu.py index ccb334e84..2a88c81bd 100644 --- a/transformer_lens/utilities/multi_gpu.py +++ b/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,7 +195,9 @@ 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(): @@ -199,7 +205,7 @@ def _validate_device_map_values( 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." )