Final singlefile checkpoint saves one folder up (#127)

* Final singlefile checkpoint saves one folder up

Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>

* save file under new pth subfolder

Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>

* Repath for easier consumption/conversion

Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>

---------

Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>

Author: daviswer

fms_fsdp/utils/checkpointing_utils.py

@@ -324,7 +324,9 @@ def save_single_file(
     ):
         # Note: metadata kwargs cannot contain any of:
         # (step, model)
-        save_name = os.path.join(self.ckp_path, "step_" + str(step) + "_ckp.pth")
+        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,