Error Loading contriever

[Reza-esfandiarpoor]

Tevatron raises an exception when trying to load facebook/contriever-msmarco.

tevatron version: installed from commit 4ad4804faf9b2bb59408528ec8ce6a4684cd596a

Command:

EMBEDDING_OUTPUT_DIR='./query_embeddings'
python -m tevatron.retriever.driver.encode \
  --output_dir=temp \
  --model_name_or_path 'facebook/contriever-msmarco' \
  --bf16 \
  --normalize \
  --encode_is_query \
  --per_device_eval_batch_size 128 \
  --query_max_len 32 \
  --passage_max_len 156 \
  --dataset_name Tevatron/msmarco-passage \
  --dataset_split dev \
  --encode_output_path $EMBEDDING_OUTPUT_DIR/query-dev.pkl

Traceback

Traceback (most recent call last):
  File "/oscar/data/superlab/projects/sydir_collection/sydir_env/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/oscar/data/superlab/projects/sydir_collection/sydir_env/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/oscar/data/superlab/projects/sydir_collection/custom_dev_package_installs/tevatron/src/tevatron/retriever/driver/encode.py", line 114, in <module>
    main()
  File "/oscar/data/superlab/projects/sydir_collection/custom_dev_package_installs/tevatron/src/tevatron/retriever/driver/encode.py", line 63, in main
    model = DenseModel.load(
  File "/oscar/data/superlab/projects/sydir_collection/custom_dev_package_installs/tevatron/src/tevatron/retriever/modeling/encoder.py", line 161, in load
    base_model = cls.TRANSFORMER_CLS.from_pretrained(model_name_or_path, device_map='auto', **hf_kwargs)
  File "/oscar/data/superlab/projects/sydir_collection/sydir_env/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 563, in from_pretrained
    return model_class.from_pretrained(
  File "/oscar/data/superlab/projects/sydir_collection/sydir_env/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3587, in from_pretrained
    no_split_modules = model._get_no_split_modules(device_map)
  File "/oscar/data/superlab/projects/sydir_collection/sydir_env/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1864, in _get_no_split_modules
    raise ValueError(
ValueError: BertModel does not support `device_map='auto'`. To implement support, the model class needs to implement the `_no_split_modules` attribute.

[MXueguang]

@ArvinZhuang this seems like a bug by the auto device map.

@Reza-esfandiarpoor remove the device auto map in the encoder.py should be able to fix it.

[ArvinZhuang]

True.. didn't realize old models are not supported..