Support M1 GPU in FARMReader

[mathislucka]

Is your feature request related to a problem? Please describe.
Since haystack v1.6 we have support for pytorch 1.12 which also means support for the M1 GPU. However, we currently initialize the device to be either cpu or cuda depending on availability and if the user passes in the use_gpu=True parameter. For GPU use on the M1, pytorch actually uses the mps backend. See: https://pytorch.org/docs/stable/notes/mps.html

If we could allow the users to pass in the actual device into the FARMReader then this might support of GPU training and inference on the M1 possible.

Describe the solution you'd like
Allow the user to pass in devices=[<device>] into FARMReader.__init__ and use these devices in initialize_device_settings. We could make this non-breaking by making this an optional argument to the reader init and the device initialization.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

[mathislucka]

It is actually already there :D

[mathislucka]

Reopening this, as the device is not used for the inferencer. See:

haystack/haystack/modeling/infer.py

Line 229 in 632cd1c

return cls(

[mathislucka]

Additionally, currently transformers does not support pytorch 1.12 (see huggingface/transformers#17971 (comment)). When changing the code in inferencer to pass on the mps device. An error is raised during prediction:

Inferencing Samples:   0%|          | 0/1 [00:00<?, ? Batches/s]
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/haystack/modeling/infer.py", line 520, in _get_predictions_and_aggregate
    logits = self.model.forward(**batch)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/haystack/modeling/model/adaptive_model.py", line 477, in forward
    output_tuple = self.language_model.forward(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/haystack/modeling/model/language_model.py", line 700, in forward
    output_tuple = self.model(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/transformers/models/roberta/modeling_roberta.py", line 841, in forward
    embedding_output = self.embeddings(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/transformers/models/roberta/modeling_roberta.py", line 105, in forward
    position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/fanal/lib/python3.9/site-packages/transformers/models/roberta/modeling_roberta.py", line 1574, in create_position_ids_from_input_ids
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
NotImplementedError: The operator 'aten::cumsum.out' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
python-BaseException

[mathislucka]

Also see this for the current state of covered ops for the mps backend:

pytorch/pytorch#77764

Hey,

Thanks for sharing this information! I am new to haystack and wondering how to enable GPU in Mac Pro M1? I have PyTorch set up already with torch.backends.mps.is_available() = True. However, I still don't know how to activate it. Can you provide a bit more information?

Best

[sjrl]

Hey, @yli223 we do not currently support the M1 GPU. We would need to implement the changes explained by @mathislucka above in Haystack. In addition, we need also need to wait for HuggingFace transformers to support PyTorch 1.12 which is required for the M1 GPU to work (more info here huggingface/transformers#17925).

[vblagoje]

Update: the HF PR has been merged to main. Therefore, we can use this feature as soon as we support HF v4.21.2 release (as soon as it gets released). Do we need to add devices optional parameter anywhere else except infer.py @mathislucka @sjrl ?

[sjrl]

That's great! I would say that anywhere the user passes an option to initialize_device_settings should have the option of passing a list of devices instead. Similar to what is already done in this load function for the Inferencer

haystack/haystack/modeling/infer.py

Lines 175 to 176 in be127e5

	if devices is None:
	devices, n_gpu = initialize_device_settings(use_cuda=gpu, multi_gpu=False)

where devices is of type

haystack/haystack/modeling/infer.py

Line 128 in be127e5

devices: Optional[List[torch.device]] = None,

So what is inconsistent at the moment is that the devices option is only supported in some places in Haystack. And I think we should support it everywhere where the user can pass in the use_gpu boolean.