fix: directly save final ckpt in save_model_dir

Signed-off-by: yashasvi <yashasvi@ibm.com>

Author: YashasviChaurasia

tuning/config/acceleration_configs/fast_moe.py

@@ -97,30 +97,25 @@ def on_save(
                 """
 
                 def checkpoint(checkpoint_dir, save_dir):
-                    hf_converted_output_dir = os.path.join(
-                        save_dir, "hf_converted_checkpoint"
-                    )
-                    if os.path.exists(hf_converted_output_dir):
+                    if os.path.exists(save_dir):
                         # If the folder already exists
                         # we return, since this is possible to happen
                         # saving the checkpointing at the end of the training
                         return
-                    os.mkdir(hf_converted_output_dir)
+
                     try:
                         recover_safetensors_from_dcp(
                             checkpoint_dir,
                             self.pretrained_model_name_or_path,
-                            hf_converted_output_dir,
+                            save_dir,
                         )
                         # Save tokenizer
                         if self.trainer.processing_class:
-                            self.trainer.processing_class.save_pretrained(
-                                hf_converted_output_dir
-                            )
+                            self.trainer.processing_class.save_pretrained(save_dir)
                         # Save training args
                         torch.save(
                             args,
-                            os.path.join(hf_converted_output_dir, TRAINING_ARGS_NAME),
+                            os.path.join(save_dir, TRAINING_ARGS_NAME),
                         )
 
                         # Unwrap FSDP module
@@ -135,16 +130,14 @@ def checkpoint(checkpoint_dir, save_dir):
                                 list(config_dict["target_modules"])
                             )
                             with open(
-                                os.path.join(
-                                    hf_converted_output_dir, "adapter_config.json"
-                                ),
+                                os.path.join(save_dir, "adapter_config.json"),
                                 "w",
                                 encoding="utf-8",
                             ) as f:
                                 json.dump(config_dict, f, indent=2)
 
                         else:
-                            model.config.save_pretrained(hf_converted_output_dir)
+                            model.config.save_pretrained(save_dir)
 
                     except Exception as e:
                         raise ValueError(
@@ -157,15 +150,19 @@ def checkpoint(checkpoint_dir, save_dir):
                     checkpoint_dir = os.path.join(
                         args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
                     )
-                    checkpoint(checkpoint_dir, checkpoint_dir)
+                    hf_converted_path = os.path.join(
+                        checkpoint_dir, "hf_converted_checkpoint"
+                    )
+                    if not os.path.exists(hf_converted_path):
+                        os.makedirs(hf_converted_path)
+                    checkpoint(checkpoint_dir, hf_converted_path)
 
                     # If final save directory is provided, save the model there
                     if (
                         getattr(self, "save_model_dir", None)
                         and state.global_step == state.max_steps
                     ):
-                        if not os.path.exists(self.save_model_dir):
-                            os.mkdir(self.save_model_dir)
+                        os.makedirs(self.save_model_dir, exist_ok=True)
                         checkpoint(checkpoint_dir, self.save_model_dir)
 
         callbacks.append(