From 44976b64fd0282fc77aa0537a59536e3ae38e7de Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 24 Jul 2023 10:45:50 -0700 Subject: [PATCH 1/3] [zero_to_fp32] 3x less cpu memory requirements --- deepspeed/utils/zero_to_fp32.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 6b595ce40add..d1a5d5eae41e 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -143,7 +143,10 @@ def parse_optim_states(files, ds_checkpoint_dir): total_files = len(files) state_dicts = [] for f in files: - state_dicts.append(torch.load(f, map_location=device)) + state_dict = torch.load(f, map_location=device) + # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights + del state_dict["optimizer_state_dict"]["optimizer_state_dict"] + state_dicts.append(state_dict) if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: raise ValueError(f"{files[0]} is not a zero checkpoint") From b8e2a38baab6357fa70be7061e698b87f7fdec44 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 24 Jul 2023 10:52:00 -0700 Subject: [PATCH 2/3] safer version --- deepspeed/utils/zero_to_fp32.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index d1a5d5eae41e..f228db9f5469 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -145,7 +145,8 @@ def parse_optim_states(files, ds_checkpoint_dir): for f in files: state_dict = torch.load(f, map_location=device) # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights - del state_dict["optimizer_state_dict"]["optimizer_state_dict"] + # and also handle the case where it was already removed by another helper script + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) state_dicts.append(state_dict) if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: From 0b515ca71e6f1f74b2477ad66de9e37a902f1ecd Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 24 Jul 2023 10:53:05 -0700 Subject: [PATCH 3/3] style --- deepspeed/utils/zero_to_fp32.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index f228db9f5469..fd5a2eb8bfdd 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -146,7 +146,7 @@ def parse_optim_states(files, ds_checkpoint_dir): state_dict = torch.load(f, map_location=device) # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights # and also handle the case where it was already removed by another helper script - state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) state_dicts.append(state_dict) if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: