Add support for `device_map="auto"` to OPT (#17382)

13541b4a · Sylvain Gugger · GitHub · 71cced8a · 13541b4a
Unverified Commit 13541b4a authored May 23, 2022 by Sylvain Gugger Committed by GitHub May 23, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/models/opt/modeling_opt.py src/transformers/models/opt/modeling_opt.py +1 -0

No files found.
--- a/src/transformers/models/opt/modeling_opt.py
+++ b/src/transformers/models/opt/modeling_opt.py
@@ -381,6 +381,7 @@ class OPTPreTrainedModel(PreTrainedModel):
    config_class = OPTConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
+    _no_split_modules = ["OPTDecoderLayer"]
    _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]

    def _init_weights(self, module):