diff --git a/tests/unit/model_bridge/test_boot_device.py b/tests/unit/model_bridge/test_boot_device.py new file mode 100644 index 000000000..b5247b9ce --- /dev/null +++ b/tests/unit/model_bridge/test_boot_device.py @@ -0,0 +1,22 @@ +"""boot()'s device for a pre-loaded ``hf_model``.""" + +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +from transformer_lens.model_bridge import TransformerBridge + + +def test_device_defaults_to_hf_model_device(monkeypatch): + # boot never moves hf_model, so cfg.device must be the model's device, not get_device()'s. + monkeypatch.setattr( + "transformer_lens.model_bridge.sources.transformers.source.get_device", lambda: "cuda" + ) + hf_config = AutoConfig.for_model( + "gpt_neox", hidden_size=32, num_hidden_layers=1, num_attention_heads=4, intermediate_size=64 + ) + with torch.device("meta"): + hf_model = AutoModelForCausalLM.from_config(hf_config) + + bridge = TransformerBridge.boot_transformers("EleutherAI/pythia-70m", hf_model=hf_model) + + assert bridge.cfg.device == "meta" diff --git a/transformer_lens/model_bridge/sources/transformers/source.py b/transformer_lens/model_bridge/sources/transformers/source.py index 54e77fa7a..bbb72634c 100644 --- a/transformer_lens/model_bridge/sources/transformers/source.py +++ b/transformer_lens/model_bridge/sources/transformers/source.py @@ -259,9 +259,10 @@ def boot( n_devices, device_map, device, max_memory ) if resolved_device_map is None: - if device is None: + # hf_model is never moved; without device= cfg.device is read from it below. + if device is None and hf_model is None: device = get_device() - adapter.cfg.device = str(device) + adapter.cfg.device = None if device is None else str(device) else: # cfg.device set from hf_device_map after the model is loaded; provisionally None. adapter.cfg.device = None