From c27232091f1b44feab884d5b8cb27d512213fd38 Mon Sep 17 00:00:00 2001 From: Joe Mayer Date: Mon, 16 Oct 2023 11:14:24 -0700 Subject: [PATCH] updating autotuner optimizer_mem factor --- deepspeed/autotuning/autotuner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/autotuning/autotuner.py b/deepspeed/autotuning/autotuner.py index c77415beb358..30d3d836c59f 100755 --- a/deepspeed/autotuning/autotuner.py +++ b/deepspeed/autotuning/autotuner.py @@ -286,7 +286,7 @@ def get_instantiation_memory_required_per_gpu(self, zero_stage): # ZeroStageEnum.disabled: params_mem = num_params * (2 if fp16_enabled else 4) gradients_mem = num_params * (2 if fp16_enabled else 4) - optimizer_mem = num_params * (16 if fp16_enabled else 8) + optimizer_mem = num_params * (4 if fp16_enabled else 8) if zero_stage >= ZeroStageEnum.optimizer_states: optimizer_mem = optimizer_mem / total_gpus