rapidfireai 0.0.1__py3-none-any.whl → 0.9.9__py3-none-any.whl

rapidfireai/ml/trainer.py

@@ -0,0 +1,309 @@
+import os
+from typing import Union
+
+import torch
+from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
+from transformers.utils.logging import set_verbosity_error
+from trl import DPOConfig, DPOTrainer, GRPOConfig, GRPOTrainer, SFTConfig, SFTTrainer
+
+from rapidfireai.ml.callbacks import GenerationMetricsCallback, MLflowLoggingCallback
+from rapidfireai.ml.checkpoint_utils import (
+    ensure_gradient_compatibility,
+    load_checkpoint_from_disk,
+    load_checkpoint_from_shared_memory,
+    load_or_create_ref_model,
+    move_tensors_to_cpu,
+    move_tensors_to_device,
+    restore_trainer_from_disk,
+    restore_trainer_from_shared_memory,
+)
+from rapidfireai.utils.constants import SHMObjectType
+from rapidfireai.utils.datapaths import DataPath
+from rapidfireai.utils.shm_manager import SharedMemoryManager
+from rapidfireai.utils.trainer_config import TrainerConfig
+
+set_verbosity_error()
+
+
+def create_trainer_instance(
+    trainer_config: TrainerConfig,
+    shm_manager: SharedMemoryManager,
+    use_shared_memory: bool = False,
+    mlflow_manager=None,
+    chunk_id: int = 0,
+) -> tuple[Union[SFTTrainer, DPOTrainer, GRPOTrainer, None], str]:
+    """
+    Create a trainer instance with proper state restoration.
+    """
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(trainer_config.worker_id)
+    device = "cuda:0"
+
+    trainer = None
+    config_leaf = trainer_config.config_leaf
+    trainer_type = config_leaf.get("trainer_type", "SFT")
+    training_args = config_leaf.get("training_args", {})
+    additional_trainer_kwargs = config_leaf.get("additional_kwargs", {})
+    compute_metrics = additional_trainer_kwargs.get("compute_metrics", None)
+
+    # Configure training arguments
+    training_args = _configure_training_args(training_args)
+    trainer_config_obj = _create_trainer_config_object(trainer_type, training_args)
+    # check if peft params is empty dict
+    if config_leaf.get("peft_params") == {}:
+        is_peft = False
+    else:
+        is_peft = True
+    # Load model and tokenizer
+    if use_shared_memory:
+        model_instance, tokenizer = load_checkpoint_from_shared_memory(trainer_config, shm_manager, is_peft=is_peft)
+    else:
+        model_instance, tokenizer = load_checkpoint_from_disk(trainer_config, is_peft=is_peft)
+    # add model name to model config
+    config_leaf["model_name"] = model_instance.config._name_or_path
+
+    # Handle reference model for DPO
+    ref_model_instance = None
+    if config_leaf.get("trainer_type") == "DPO":
+        model_instance, ref_model_instance = _setup_reference_model(
+            model_instance, trainer_config, config_leaf, use_shared_memory, shm_manager, device, is_peft
+        )
+
+    model_instance = model_instance.to(device)
+
+    trainer_kwargs, formatting_func, additional_trainer_kwargs = _prepare_trainer_kwargs(
+        model_instance,
+        trainer_config_obj,
+        tokenizer,
+        trainer_config,
+        additional_trainer_kwargs,
+        ref_model_instance,
+        config_leaf,
+    )
+
+    callbacks, additional_trainer_kwargs = _setup_callbacks(  # FIXME: avoid returning additional_trainer_kwargs
+        mlflow_manager,
+        trainer_config,
+        chunk_id,
+        compute_metrics,
+        additional_trainer_kwargs,
+        tokenizer,
+        training_args,
+        formatting_func,
+    )
+
+    if callbacks:
+        trainer_kwargs["callbacks"] = callbacks
+
+    trainer_kwargs.update(additional_trainer_kwargs)
+    trainer_kwargs = {k: v for k, v in trainer_kwargs.items() if v is not None}
+
+    trainer = _create_trainer_by_type(trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager)
+    return trainer, config_leaf["model_name"]
+
+
+def _configure_training_args(training_args: dict) -> dict:
+    """Configure training arguments with default values."""
+    training_args.pop("max_steps", None)
+    training_args["num_train_epochs"] = 1
+    training_args["save_strategy"] = "no"
+    training_args["do_train"] = True
+    training_args["do_eval"] = True
+    training_args["dataloader_pin_memory"] = False
+    training_args["no_cuda"] = False
+    training_args["local_rank"] = -1
+    training_args["disable_tqdm"] = True
+
+    if "save_steps" in training_args:
+        training_args.pop("save_steps")
+
+    return training_args
+
+
+def _create_trainer_config_object(trainer_type: str, training_args: dict):
+    """Create the appropriate trainer config object based on trainer type."""
+    if trainer_type == "SFT":
+        return SFTConfig(**training_args)
+    elif trainer_type == "DPO":
+        return DPOConfig(**training_args)
+    elif trainer_type == "GRPO":
+        return GRPOConfig(**training_args)
+    else:
+        raise ValueError(f"Unsupported trainer type: {trainer_type}")
+
+
+def _setup_reference_model(
+    model_instance, trainer_config, config_leaf, use_shared_memory, shm_manager, device, is_peft
+):
+    """Setup reference model for DPO training."""
+    ref_model_instance = None
+    training_args = config_leaf.get("training_args", {})
+    if is_peft and not training_args.get("force_use_ref_model", False):
+        model_adapter_name = training_args.get("model_adapter_name", "default")
+        ref_adapter_name = training_args.get("ref_adapter_name", "reference")
+
+        if model_adapter_name is not None and ref_adapter_name is not None:
+            if use_shared_memory:
+                peft_config = LoraConfig(**config_leaf["peft_params"])
+                if trainer_config.completed_steps == 0 and trainer_config.warm_started_from is None:
+                    reference_state_dict = get_peft_model_state_dict(model_instance)
+                    reference_state_dict = move_tensors_to_cpu(reference_state_dict)
+                    shm_manager.save_model_object(
+                        trainer_config.run_id, SHMObjectType.REF_STATE_DICT, reference_state_dict
+                    )
+                else:
+                    reference_state_dict = shm_manager.load_model_object(
+                        trainer_config.run_id, SHMObjectType.REF_STATE_DICT
+                    )
+                    reference_state_dict = move_tensors_to_device(reference_state_dict, device)
+                model_instance.add_adapter(ref_adapter_name, peft_config)
+                model_instance.set_adapter(ref_adapter_name)
+                set_peft_model_state_dict(model_instance, reference_state_dict, adapter_name=ref_adapter_name)
+                model_instance.set_adapter(model_adapter_name)
+            else:
+                base_run_path = DataPath.base_run_path(trainer_config.run_id)
+                ref_model_path = DataPath.ref_model_path(base_run_path)
+                reference_adapter_path = ref_model_path / "reference"
+
+                if not reference_adapter_path.exists():
+                    os.makedirs(reference_adapter_path, exist_ok=True)
+                    model_instance.save_pretrained(reference_adapter_path)
+                torch.cuda.empty_cache()
+                model_instance.load_adapter(
+                    reference_adapter_path, adapter_name=ref_adapter_name, device_map={"": device}
+                )
+                model_instance.set_adapter(model_adapter_name)
+            model_instance = model_instance.to(device)
+    else:
+        ref_model_instance = load_or_create_ref_model(
+            model_instance, trainer_config, device, use_shared_memory, shm_manager
+        )
+        ref_model_instance = ref_model_instance.to(device)
+    return model_instance, ref_model_instance
+
+
+def _prepare_trainer_kwargs(
+    model_instance,
+    trainer_config_obj,
+    tokenizer,
+    trainer_config,
+    additional_trainer_kwargs,
+    ref_model_instance,
+    config_leaf,
+):
+    """Prepare keyword arguments for trainer creation."""
+    if config_leaf.get("trainer_type") == "DPO  ":
+        model_instance = ensure_gradient_compatibility(
+            model_instance, hasattr(model_instance, "peft_config")
+        )  # FIXME: change function for DPO
+    trainer_kwargs = {
+        "model": model_instance,
+        "args": trainer_config_obj,
+        "processing_class": tokenizer,
+    }
+
+    train_dataset = trainer_config.train_dataset
+    eval_dataset = trainer_config.eval_dataset
+    formatting_func = None
+
+    if additional_trainer_kwargs.get("formatting_func") is not None:
+        formatting_func = additional_trainer_kwargs.get("formatting_func")
+        train_dataset = train_dataset.map(formatting_func)  # FIXME: add try exception with batched/unbatched
+        if eval_dataset is not None:
+            eval_dataset = eval_dataset.map(formatting_func)
+        additional_trainer_kwargs_copy = additional_trainer_kwargs.copy()
+        additional_trainer_kwargs_copy.pop("formatting_func")
+        additional_trainer_kwargs = additional_trainer_kwargs_copy
+
+    trainer_kwargs["train_dataset"] = train_dataset
+    if eval_dataset is not None:
+        trainer_kwargs["eval_dataset"] = eval_dataset
+
+    if config_leaf.get("trainer_type") == "DPO" and ref_model_instance is not None:
+        trainer_kwargs["ref_model"] = ref_model_instance
+
+    if config_leaf.get("trainer_type") == "GRPO":
+        reward_funcs = config_leaf.get("reward_funcs")
+        if reward_funcs is not None:
+            trainer_kwargs["reward_funcs"] = reward_funcs
+
+    return trainer_kwargs, formatting_func, additional_trainer_kwargs
+
+
+def _setup_callbacks(
+    mlflow_manager,
+    trainer_config,
+    chunk_id,
+    compute_metrics,
+    additional_trainer_kwargs,
+    tokenizer,
+    training_args,
+    formatting_func,
+):
+    """Setup callbacks for the trainer."""
+    callbacks = []
+
+    if mlflow_manager is not None and trainer_config.mlflow_run_id is not None:
+        mlflow_callback = MLflowLoggingCallback(
+            mlflow_manager=mlflow_manager,
+            mlflow_run_id=trainer_config.mlflow_run_id,
+            completed_steps=trainer_config.completed_steps,
+            chunk_id=chunk_id,
+            num_epochs_completed=trainer_config.num_epochs_completed,
+        )
+        callbacks.append(mlflow_callback)
+
+    if compute_metrics is not None and additional_trainer_kwargs.get("generation_config") is not None:
+        compute_metrics_function = compute_metrics
+        if formatting_func is not None:
+            formatted_eval_dataset = trainer_config.eval_dataset.map(formatting_func)
+        else:
+            formatted_eval_dataset = trainer_config.eval_dataset
+
+        generation_callback = GenerationMetricsCallback(
+            tokenizer=tokenizer,
+            eval_dataset=formatted_eval_dataset,
+            generation_config=additional_trainer_kwargs.get("generation_config"),
+            compute_metrics=compute_metrics_function,
+            batch_size=training_args.get("per_device_eval_batch_size"),
+            mlflow_manager=mlflow_manager,
+            mlflow_run_id=trainer_config.mlflow_run_id,
+            completed_steps=trainer_config.completed_steps,
+        )
+        callbacks.append(generation_callback)
+        additional_trainer_kwargs.pop("generation_config")
+        additional_trainer_kwargs.pop("compute_metrics")
+
+    return callbacks, additional_trainer_kwargs
+
+
+def _create_trainer_by_type(trainer_type, trainer_kwargs, trainer_config, use_shared_memory, shm_manager):
+    """Create trainer instance based on type with proper state restoration."""
+    if trainer_type == "SFT":
+        dummy_trainer = SFTTrainer(**trainer_kwargs)
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
+        trainer = SFTTrainer(**trainer_kwargs, optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler))
+        del dummy_trainer
+
+    elif trainer_type == "DPO":
+        dummy_trainer = DPOTrainer(**trainer_kwargs)
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
+        trainer = DPOTrainer(**trainer_kwargs, optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler))
+        num_trainable = sum(p.requires_grad for _, p in trainer.model.named_parameters())
+        num_opt_params = sum(len(g["params"]) for g in trainer.optimizer.state_dict()["param_groups"])
+        del dummy_trainer
+
+    elif trainer_type == "GRPO":
+        dummy_trainer = GRPOTrainer(**trainer_kwargs)
+        dummy_trainer.create_optimizer_and_scheduler(num_training_steps=trainer_config.total_steps)
+        trainer = GRPOTrainer(**trainer_kwargs, optimizers=(dummy_trainer.optimizer, dummy_trainer.lr_scheduler))
+        del dummy_trainer
+    else:
+        raise ValueError(f"Unsupported trainer type: {trainer_type}")
+
+    if trainer_config.completed_steps > 0:
+        if use_shared_memory:
+            trainer = restore_trainer_from_shared_memory(trainer, trainer_config, shm_manager)
+        else:
+            trainer = restore_trainer_from_disk(trainer, trainer_config)
+
+    return trainer