feat: move callbacks to fms-hf-tuning

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

tuning/config/acceleration_configs/__init__.py

@@ -15,6 +15,7 @@
 # Local
 from .acceleration_framework_config import AccelerationFrameworkConfig
 from .attention_and_distributed_packing import AttentionAndDistributedPackingConfig
+from .callbacks import get_additional_accel_framework_callbacks
 from .fast_moe import FastMoeConfig
 from .fused_ops_and_kernels import FusedOpsAndKernelsConfig
 from .quantized_lora_config import QuantizedLoraConfig

tuning/config/acceleration_configs/callbacks.py

@@ -0,0 +1,10 @@
+# Local
+from .fast_moe import get_callbacks
+
+
+def get_additional_accel_framework_callbacks(active_plugins, **kwargs):
+    callbacks = []
+    for active_plugin in active_plugins:
+        if "ScatterMoEAccelerationPlugin" == active_plugin[0]:
+            callbacks.extend(get_callbacks(**kwargs))
+    return callbacks

tuning/config/acceleration_configs/fast_moe.py

@@ -14,10 +14,30 @@
 
 # Standard
 from dataclasses import dataclass
+import os
+
+# Third Party
+from transformers import (
+    Trainer,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+import torch
 
 # Local
 from .utils import ensure_nested_dataclasses_initialized, parsable_dataclass
 
+is_recover_safetensors_from_dcp_available = True
+try:
+    # Third Party
+    from fms_acceleration_moe.utils import recover_safetensors_from_dcp
+except ImportError:
+    is_recover_safetensors_from_dcp_available = False
+
 
 @parsable_dataclass
 @dataclass
@@ -34,3 +54,77 @@ class FastMoeConfig:
     def __post_init__(self):
         # ensure nested dataclasses initialized
         ensure_nested_dataclasses_initialized(self)
+
+
+def get_callbacks(**kwargs):
+    pretrained_model_name_or_path = kwargs.pop("pretrained_model_name_or_path")
+    trainer = kwargs.pop("trainer")
+    callbacks = []
+    if is_recover_safetensors_from_dcp_available:
+
+        class ConvertAndSaveHFCheckpointAtEverySave(TrainerCallback):
+            def __init__(self, pretrained_model_name_or_path: str, trainer: Trainer):
+                self.pretrained_model_name_or_path = pretrained_model_name_or_path
+                self.trainer = trainer
+
+            def on_save(
+                self,
+                args: TrainingArguments,
+                state: TrainerState,
+                control: TrainerControl,
+                **kwargs,
+            ):
+                """
+                Save all HF files and convert dcp checkpoint to safetensors at every save operation.
+                """
+
+                def checkpoint():
+                    checkpoint_dir = os.path.join(
+                        args.output_dir,
+                        f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
+                    )
+                    hf_converted_output_dir = os.path.join(
+                        checkpoint_dir, "hf_converted_checkpoint"
+                    )
+                    if os.path.exists(hf_converted_output_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 tokenizer
+                        if self.trainer.processing_class:
+                            self.trainer.processing_class.save_pretrained(
+                                hf_converted_output_dir
+                            )
+                        # save training args
+                        torch.save(
+                            args,
+                            os.path.join(hf_converted_output_dir, TRAINING_ARGS_NAME),
+                        )
+                        # save model config files
+                        self.trainer.model.config.save_pretrained(
+                            hf_converted_output_dir
+                        )
+
+                    except Exception as e:
+                        raise ValueError(
+                            f"Failed to convert the checkpoint {checkpoint_dir}\
+                                to a HF compatible checkpoint"
+                        ) from e
+
+                if state.is_world_process_zero:
+                    checkpoint()
+
+        callbacks.append(
+            ConvertAndSaveHFCheckpointAtEverySave(
+                pretrained_model_name_or_path, trainer
+            )
+        )
+    return callbacks

tuning/sft_trainer.py

@@ -48,6 +48,7 @@
     FastMoeConfig,
     FusedOpsAndKernelsConfig,
     QuantizedLoraConfig,
+    get_additional_accel_framework_callbacks,
 )
 from tuning.config.tracker_configs import (
     AimConfig,
@@ -408,10 +409,14 @@ def train(
         accelerator = None if not is_accelerate_available() else trainer.accelerator
 
         # ready for train may produce additional callbacks for the trainer
-        for x in framework.get_callbacks_and_ready_for_train(
-            model, accelerator, trainer, model_args.model_name_or_path
-        ):
+        for x in framework.get_callbacks_and_ready_for_train(model, accelerator):
             trainer.add_callback(x)
+        for clb in get_additional_accel_framework_callbacks(
+            active_plugins=framework.active_plugins,
+            trainer=trainer,
+            pretrained_model_name_or_path=model_args.model_name_or_path,
+        ):
+            trainer.add_callback(clb)
 
     resume_from_checkpoint = None
     # Check if resume flag is not passed (None), or if flag is true and