Load model `all-mpnet-base-v2` with `map_device="auto"`

[maayansharon10]

Hi,

Trying to load model all-mpnet-base-v2 with map_device="auto".

Following this closed issue I understand that it is possible to load a model with map_device="auto", but I'm not sure what's the syntax when using the example from the doc and not using .from_pretrained(). Couldn't find an example in the documentation.

My current code follows the documentation:

    model = SentenceTransformer(
        model_name_or_path="sentence-transformers/all-mpnet-base-v2",
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            model_name=model_name,
        ),
        similarity_fn_name="cosine",
        device="cuda",
    )

When tried the following:

    model = SentenceTransformer(
        model_name_or_path="sentence-transformers/all-mpnet-base-v2",
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            model_name=model_name,
        ),
        similarity_fn_name="cosine",
        device="cuda",
        model_kwargs={'device_map': "auto"},
    )

Got this error -

  File ".../model_informationRetrievalEvaluator.py", line 46, in main
    model = SentenceTransformer(
  File "..Venv/lib/python3.9/site-packages/sentence_transformers/SentenceTransformer.py", line 287, in __init__
    modules = self._load_sbert_model(
  File "...Venv/lib/python3.9/site-packages/sentence_transformers/SentenceTransformer.py", line 1487, in _load_sbert_model
    module = Transformer(model_name_or_path, cache_dir=cache_folder, **kwargs)
  File "...Venv/lib/python3.9/site-packages/sentence_transformers/models/Transformer.py", line 54, in __init__
    self._load_model(model_name_or_path, config, cache_dir, **model_args)
  File "...Venv/lib/python3.9/site-packages/sentence_transformers/models/Transformer.py", line 85, in _load_model
    self.auto_model = AutoModel.from_pretrained(
  File "...Venv/lib/python3.9/site-packages/transformers/models/auto/auto_factory.py", line 563, in from_pretrained
    return model_class.from_pretrained(
  File "...Venv/lib/python3.9/site-packages/transformers/modeling_utils.py", line 3669, in from_pretrained
    no_split_modules = model._get_no_split_modules(device_map)
  File "...Venv/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1892, in _get_no_split_modules
    raise ValueError(
ValueError: MPNetModel does not support `device_map='auto'`. To implement support, the model class needs to implement the `_no_split_modules` attribute.

Again, circling back to the face that in the issue above, and I might misunderstood, there's an option to add map_device. If not possible, is there any other way to tackle this?

P.s This issue might be similar, not sure if I should have commented there, sorry in advance if I should have.

[tomaarsen]

Hello!

Your code is correct,

        model_kwargs={'device_map': "auto"},

indeed works with most transformers architectures. However, it seems that modeling_mpnet hasn't been updated to add the _no_split_modules behaviour. Sadly, I can't find much info/docs on how to set _no_split_modules beyond https://github.com/huggingface/transformers/issues/29786, so I'm not sure what it should be defined as.

Perhaps we can make a fresh issue on transformers to add _no_split_modules to MPNetModel. cc @amyeroberts

• Tom Aarsen