SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[CodeGen] support device_map="auto" for sharded checkpoints (#17871)

Browse Source
This commit is contained in:
Suraj Patil
2022-06-24 18:06:30 +02:00
committed by GitHub
parent d6b6fb9963
commit 061a73d16f
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/transformers/models/codegen/modeling_codegen.py
View File

@@ -332,6 +332,7 @@ class CodeGenPreTrainedModel(PreTrainedModel):
config_class = CodeGenConfig config_class = CodeGenConfig
base_model_prefix = "transformer" base_model_prefix = "transformer"
supports_gradient_checkpointing = True supports_gradient_checkpointing = True
_no_split_modules = ["CodeGenBlock"]
def __init__(self, *inputs, **kwargs): def __init__(self, *inputs, **kwargs):
super().__init__(*inputs, **kwargs) super().__init__(*inputs, **kwargs)