Prevent splitting of ModifiedMistralDecoderLayers

[hatzel]

Previously the layers could be split across devices, e.g. when using device_map='auto'. This would result errors like this one:

RuntimeError:
Expected all tensors to be on the same device,
but found at least two devices, cuda:0 and cuda:1!

The remote code I got from huggingface has a different name MistralEncoderModel rather than MistralBiModel but I assume that this is the correct file to change. I tested this in locally by editing the remote files in my ~/.cache/ directory.

[hatzel]

I suppose this may also need to be fixed for the llama model but I have not checked.

[vaibhavad]

Thanks @hatzel for spotting this! Can you also add changes for Llama, so that everything is in the same PR?

_no_split_modules = ["ModifiedLlamaDecoderLayer"]

This will need to be added here

[hatzel]

Done!

[vaibhavad]

Thanks a lot @hatzel!