Unverified Commit 857eb87c authored by Manuel R. Ciosici's avatar Manuel R. Ciosici Committed by GitHub
Browse files

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

parent 81ac45f8
......@@ -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
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)
if self.is_zero3():
# automatically assign the optimal config values based on model config
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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment