From 061a73d16f186e147b12cc75c3feab0bcd5bf80d Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Fri, 24 Jun 2022 18:06:30 +0200 Subject: [PATCH] [CodeGen] support device_map="auto" for sharded checkpoints (#17871) --- src/transformers/models/codegen/modeling_codegen.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/codegen/modeling_codegen.py b/src/transformers/models/codegen/modeling_codegen.py index ea9a7ef6a2..06581e732c 100644 --- a/src/transformers/models/codegen/modeling_codegen.py +++ b/src/transformers/models/codegen/modeling_codegen.py @@ -332,6 +332,7 @@ class CodeGenPreTrainedModel(PreTrainedModel): config_class = CodeGenConfig base_model_prefix = "transformer" supports_gradient_checkpointing = True + _no_split_modules = ["CodeGenBlock"] def __init__(self, *inputs, **kwargs): super().__init__(*inputs, **kwargs)