From 201dd2fdae8eaefc3febd860d36734d977299916 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 12 Jul 2021 14:55:41 -0700 Subject: [PATCH] revert part of #1220 https://github.com/microsoft/DeepSpeed/pull/1220 fixed the leak, but lead to another problem. reverting that part so that we could do release and will work on it after the release. @jeffra --- deepspeed/runtime/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 93af1dda458b..d24e99af488e 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -2045,7 +2045,7 @@ def get_layer_state_dict(module, prefix=""): # gather one layer at a time to be memory-efficient with deepspeed.zero.GatheredParameters(list( module.parameters(recurse=False)), - modifier_rank=0): + modifier_rank=None): if torch.distributed.get_rank() == 0: for name, param in module.named_parameters(recurse=False): if param is None: