From d2bea2270a0918d1efaf668d1171f607d9684a88 Mon Sep 17 00:00:00 2001 From: "anthony.301" Date: Fri, 2 Sep 2022 11:40:26 -0500 Subject: [PATCH] Fix for Zero3 when MP>1 and at least one batch param undefined --- deepspeed/runtime/zero/stage3.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 2777fd728a81..2a78f8f92f57 100755 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -158,7 +158,8 @@ def __init__(self, max_live_parameters=max_live_parameters, param_persistence_threshold=param_persistence_threshold, model_persistence_threshold=model_persistence_threshold, - offload_param_config=offload_optimizer_config) + offload_param_config=offload_optimizer_config, + mpu=mpu) self.persistent_parameters = self.parameter_offload.persistent_parameters self._configure_offloading(offload_optimizer_config, offload_param_config)