turn off pth saving which is useless now

Author: lchu6

fms_fsdp/utils/checkpointing_utils.py

@@ -314,34 +314,3 @@ def save(
         )
 
         return self._cleanup()
-
-    def save_single_file(
-        self,
-        step,
-        model,
-        is_compiled=False,
-        **kwargs,
-    ):
-        # Note: metadata kwargs cannot contain any of:
-        # (step, model)
-        pth_path = os.path.join(self.ckp_path[:-12], "pth", "step_" + str(step))
-        os.makedirs(pth_path, exist_ok=True)
-        save_name = os.path.join(pth_path, "consolidated.00.pth")
-        save_time = time.time()
-        with FSDP.state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-        ):
-            if is_compiled:
-                model_state = model._orig_mod.state_dict()
-            else:
-                model_state = model.state_dict()
-        if self.rank == 0:
-            metadata = kwargs
-            metadata["step"] = step
-            metadata["model_state"] = model_state
-            torch.save(metadata, save_name)
-        self.report("Checkpoint written", model_save_time=time.time() - save_time)
-
-        return self._cleanup()

main_training_llama.py

@@ -167,8 +167,6 @@ def main(**kwargs):
         tokens_seen,
     )
 
-    checkpointer.save_single_file(cfg.num_steps, model)
-
     dist.barrier()
     dist.destroy_process_group()
 

main_training_mamba.py

@@ -167,8 +167,6 @@ def main(**kwargs):
         tokens_seen,
     )
 
-    checkpointer.save_single_file(cfg.num_steps, model)
-
     dist.barrier()
     dist.destroy_process_group()
 

speculator/train_speculator_utils.py

@@ -412,6 +412,7 @@ def train_speculator(
 
         if (
             batch_idx % cfg.checkpoint_interval == 0
+            or batch_idx == cfg.num_steps
             or do_ckpt(cfg.ckpt_save_path) is True
         ):
             torch.cuda.empty_cache()
@@ -425,13 +426,6 @@ def train_speculator(
             torch.cuda.empty_cache()
             do_ckpt(cfg.ckpt_save_path, reset=True)
 
-    checkpointer.save_single_file(
-        batch_idx,
-        speculator,
-        tokens_seen=elapsed_tokens + n_tok,
-        is_compiled=cfg.use_torch_compile,
-    )
-
 
 class EmbedGPTBigCode(GPTBigCode):
     # Overrides the forward function of GPTBigCode to allow returning embedding vectors