dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support reduce_bucket_size=auto for deepspeed stages <3 (#16496)

This commit is contained in:
Manuel R. Ciosici 2022-03-30 14:12:29 -07:00 committed by GitHub
parent 81ac45f85c
commit 857eb87cc4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/transformers/deepspeed.py
View File

@ -261,13 +261,13 @@ class HfTrainerDeepSpeedConfig(HfDeepSpeedConfig):
"""
This stage is run after we have the model and know num_training_steps.
Now we we can complete the configuration process.
Now we can complete the configuration process.
"""
# zero
hidden_size = model.config.hidden_size
self.fill_only("zero_optimization.reduce_bucket_size", hidden_size * hidden_size)
if self.is_zero3():
# automatically assign the optimal config values based on model config
hidden_size = model.config.hidden_size
self.fill_only("zero_optimization.reduce_bucket_size", hidden_size * hidden_size)
self.fill_only("zero_optimization.stage3_prefetch_bucket_size", 0.9 * hidden_size * hidden_size)
self.fill_only("zero_optimization.stage3_param_persistence_threshold", 10 * hidden_size)