zero_to_fp32: restore persistent buffers

deepspeed/runtime/engine.py

@@ -1915,6 +1915,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()
@@ -1934,6 +1935,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():

deepspeed/utils/zero_to_fp32.py

@@ -20,13 +20,27 @@
 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 = 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}'")
+
+    return file
+
+
 def get_optim_files(checkpoint_dir):
 
     if not os.path.isdir(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(
@@ -35,13 +49,34 @@ def get_optim_files(checkpoint_dir):
     return optim_files
 
 
+def parse_model_state(file):
+
+    # 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")
+    buffer_names = state_dict["buffer_names"]
+    if debug:
+        print(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
+
+
 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 +121,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 +145,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():