Skip to content

Commit c5e29d4

Browse files
authored
Fix typing annotations for FSDP and DeepSpeed in TrainingArguments (huggingface#24549)
* Fix typing annotations for FSDP and DeepSpeed in TrainingArguments * Change dict to Dict
1 parent daccde1 commit c5e29d4

File tree

1 file changed

+5
-5
lines changed

1 file changed

+5
-5
lines changed

src/transformers/training_args.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -976,12 +976,12 @@ class TrainingArguments:
976976
)
977977
},
978978
)
979-
fsdp_config: Optional[str] = field(
979+
fsdp_config: Optional[Union[str, Dict]] = field(
980980
default=None,
981981
metadata={
982982
"help": (
983-
"Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a"
984-
"fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`."
983+
"Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a"
984+
"fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`."
985985
)
986986
},
987987
)
@@ -994,11 +994,11 @@ class TrainingArguments:
994994
)
995995
},
996996
)
997-
deepspeed: Optional[str] = field(
997+
deepspeed: Optional[Union[str, Dict]] = field(
998998
default=None,
999999
metadata={
10001000
"help": (
1001-
"Enable deepspeed and pass the path to deepspeed json config file (e.g. ds_config.json) or an already"
1001+
"Enable deepspeed and pass the path to deepspeed json config file (e.g. `ds_config.json`) or an already"
10021002
" loaded json file as a dict"
10031003
)
10041004
},

0 commit comments

Comments
 (0)