From 48949ae99cfc3983fd6053c6af0c048e5b4ad86a Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 7 Jun 2021 22:23:24 -0700 Subject: [PATCH 1/3] zero_to_fp32: restore buffers --- deepspeed/runtime/engine.py | 22 +++++++++++++++++ deepspeed/utils/zero_to_fp32.py | 42 +++++++++++++++++++++++++++++++-- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 28ad2e977f17..1ea1ee181491 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -1903,6 +1903,7 @@ def _save_checkpoint(self, save_dir, tag, client_state={}): self._curr_ckpt_path = os.path.join(save_dir, tag) state = dict(module=self.module_state_dict(), + buffer_names=self._get_buffer_names(), optimizer=self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None, lr_scheduler=self.lr_scheduler.state_dict() @@ -1922,6 +1923,27 @@ def _save_checkpoint(self, save_dir, tag, client_state={}): torch.save(state, save_path) self._curr_save_path = None + def _get_buffer_names(self): + buffer_names = [] + + # we save buffer names so that we could extract later the real buffers from the saved + # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they + # are intermixed with param placeholders + + # have to traverse the tree to be able to skip non-persistent buffers + def get_layer_named_buffers(module, prefix=""): + for name, buf in module.named_buffers(recurse=False): + if buf is not None and name not in module._non_persistent_buffers_set: + buffer_names.append(prefix + name) + + for name, child in module.named_children(): + if child is not None: + get_layer_named_buffers(child, prefix + name + ".") + + get_layer_named_buffers(self.module, prefix="") + + return buffer_names + def _get_param_shapes(self): param_shapes = OrderedDict() for name, param in self.module.named_parameters(): diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 2d98a39e3fc7..551e010436e3 100644 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -20,6 +20,20 @@ debug = 0 +def get_model_state_file(checkpoint_dir): + + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + file = f"{checkpoint_dir}/zero_pp_rank_0_mp_rank_00_model_states.pt" + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find '{file}' in directory '{checkpoint_dir}'") + + return file + + def get_optim_files(checkpoint_dir): if not os.path.isdir(checkpoint_dir): @@ -35,13 +49,28 @@ def get_optim_files(checkpoint_dir): return optim_files +def parse_model_state(file): + state_dict = torch.load(file) + + buffers = {} + + if "buffer_names" not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict["buffer_names"] + if debug: + print(buffer_names) + + buffers = {k: v.cpu() for k, v in state_dict["module"].items() if k in buffer_names} + return buffers + + def parse_optim_states(files): state_dicts = [] for f in files: state_dicts.append(torch.load(f)) if not "zero_stage" in state_dicts[0]['optimizer_state_dict']: - raise ValueError(f"non zero checkpoint") + raise ValueError(f"{files[0]} is not a zero checkpoint") zero_stage = state_dicts[0]['optimizer_state_dict']["zero_stage"] world_size = state_dicts[0]['optimizer_state_dict']["partition_count"] param_shapes = state_dicts[0]["param_shapes"] @@ -86,7 +115,9 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file): """ print(f"Processing zero checkpoint '{checkpoint_dir}'") + model_file = get_model_state_file(checkpoint_dir) optim_files = get_optim_files(checkpoint_dir) + buffers = parse_model_state(model_file) zero_stage, world_size, param_shapes, fp32_flat_groups = parse_optim_states(optim_files) print( f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") @@ -108,9 +139,16 @@ def convert_zero_chkpt_to_fp32_consolid_state_dict(checkpoint_dir, output_file): # XXX: memory usage doubles here (zero2) full_single_fp32_vector = torch.cat(fp32_flat_groups, 0) + state_dict = OrderedDict() + + # buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + # params # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support # out-of-core computing solution - state_dict = OrderedDict() offset = 0 total_numel = 0 for name, shape in param_shapes.items(): From cb13a3ae2748ce1a3f23fd44b954ab97bee4ade9 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 7 Jun 2021 22:49:57 -0700 Subject: [PATCH 2/3] load to cpu, recover fp32 --- deepspeed/utils/zero_to_fp32.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 551e010436e3..5b4f7bba3cad 100644 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -50,9 +50,10 @@ def get_optim_files(checkpoint_dir): def parse_model_state(file): - state_dict = torch.load(file) - buffers = {} + # load to cpu + device = torch.device('cpu') + state_dict = torch.load(file, map_location=device) if "buffer_names" not in state_dict: raise ValueError(f"{file} is not a model state checkpoint") @@ -60,7 +61,12 @@ def parse_model_state(file): if debug: print(buffer_names) - buffers = {k: v.cpu() for k, v in state_dict["module"].items() if k in buffer_names} + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = { + k: v.float() + for k, + v in state_dict["module"].items() if k in buffer_names + } return buffers From 322c82edc8639b11f41c03540229c0c1ac2c2361 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Tue, 8 Jun 2021 09:43:14 -0700 Subject: [PATCH 3/3] cross-platform path sep --- deepspeed/utils/zero_to_fp32.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 5b4f7bba3cad..1b58c403618c 100644 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -26,7 +26,7 @@ def get_model_state_file(checkpoint_dir): raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") # there should be only one file - file = f"{checkpoint_dir}/zero_pp_rank_0_mp_rank_00_model_states.pt" + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") if not os.path.exists(file): raise FileNotFoundError(f"can't find '{file}' in directory '{checkpoint_dir}'") @@ -40,7 +40,7 @@ def get_optim_files(checkpoint_dir): raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") # XXX: need to test that this simple glob rule works for multi-node setup too - optim_files = sorted(glob.glob(f"{checkpoint_dir}/*_optim_states.pt")) + optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt"))) if len(optim_files) == 0: raise FileNotFoundError(