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

rapidfireai/ml/checkpoint_utils.py

@@ -0,0 +1,540 @@
+import copy
+import json
+import os
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+from peft import LoraConfig, PeftModel, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
+from transformers import AutoTokenizer
+from trl import DPOTrainer, GRPOTrainer, SFTTrainer
+
+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
+
+
+def move_tensors_to_device(obj, device: torch.device):
+    """Recursively move all tensors in a nested structure to device"""
+    if isinstance(obj, torch.Tensor):
+        return obj.to(device, non_blocking=True)
+    elif isinstance(obj, dict):
+        return {key: move_tensors_to_device(value, device) for key, value in obj.items()}
+    elif isinstance(obj, list):
+        return [move_tensors_to_device(item, device) for item in obj]
+    elif isinstance(obj, tuple):
+        return tuple(move_tensors_to_device(item, device) for item in obj)
+    else:
+        return obj
+
+
+def move_tensors_to_cpu(obj):
+    """Recursively move all tensors in a nested structure to CPU"""
+    if isinstance(obj, torch.Tensor):
+        return obj.cpu().clone()
+    elif isinstance(obj, dict):
+        return {key: move_tensors_to_cpu(value) for key, value in obj.items()}
+    elif isinstance(obj, list):
+        return [move_tensors_to_cpu(item) for item in obj]
+    elif isinstance(obj, tuple):
+        return tuple(move_tensors_to_cpu(item) for item in obj)
+    else:
+        return obj
+
+
+def ensure_gradient_compatibility(model, use_peft: bool = False):
+    """Ensure model parameters have proper gradient settings"""
+    model.train()
+    if use_peft:
+        model.base_model.eval()
+        for name, param in model.named_parameters():
+            if any(adapter_key in name.lower() for adapter_key in ["lora", "adapter", "modules_to_save"]):
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+    else:
+        for param in model.parameters():
+            param.requires_grad = True
+    for n, p in model.named_parameters():
+        if "reference" in n:
+            p.requires_grad = False
+    model.train()
+    torch.set_grad_enabled(True)
+
+    return model
+
+
+def _configure_tokenizer(tokenizer: AutoTokenizer) -> None:
+    """Configure tokenizer with proper padding token."""
+    if tokenizer.pad_token is None:
+        if tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+        else:
+            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+
+def create_model_instance(
+    model_config: dict,
+    create_model_fn: Callable,
+    checkpoint_path: Optional[str] = None,
+    is_peft: bool = False,
+    device: Optional[str] = None,
+) -> tuple[nn.Module, AutoTokenizer]:
+    """Create a model instance from a model configuration"""
+    if device is not None:
+        model_config["model_kwargs"]["device_map"] = {"": device}
+    if checkpoint_path and not is_peft:
+        model_config["model_name"] = checkpoint_path
+
+    model_instance, tokenizer = create_model_fn(model_config)
+
+    if is_peft and checkpoint_path:
+        model_instance = PeftModel.from_pretrained(model_instance, checkpoint_path)
+    _configure_tokenizer(tokenizer)
+
+    return model_instance, tokenizer
+
+
+def save_checkpoint_to_shared_memory(
+    trainer: Union[SFTTrainer, DPOTrainer, GRPOTrainer], trainer_config: TrainerConfig, shm_manager: SharedMemoryManager
+) -> None:
+    """Save checkpoint to shared memory"""
+    checkpoint = {}
+
+    if hasattr(trainer.model, "peft_config"):
+        peft_state_dict = get_peft_model_state_dict(
+            trainer.model,
+            adapter_name=trainer_config.config_leaf.get("training_args").get("model_adapter_name", "default"),
+        )
+        checkpoint["state_dict"] = {k: v.cpu().clone() for k, v in peft_state_dict.items()}
+        checkpoint["adapter_config"] = trainer.model.peft_config
+
+    if trainer.optimizer is not None:
+        checkpoint["optimizer_state"] = move_tensors_to_cpu(trainer.optimizer.state_dict())
+
+    if trainer.lr_scheduler is not None:
+        checkpoint["scheduler_state"] = move_tensors_to_cpu(trainer.lr_scheduler.state_dict())
+
+    if hasattr(trainer, "state"):
+        checkpoint["trainer_state"] = move_tensors_to_cpu(trainer.state.__dict__.copy())
+
+    if hasattr(trainer, "rng_state"):
+        checkpoint["rng_state"] = move_tensors_to_cpu(trainer.rng_state)
+
+    if hasattr(trainer.args, "__dict__"):
+        checkpoint["training_args"] = trainer.args.__dict__.copy()
+
+    if hasattr(trainer.model, "generation_config"):
+        checkpoint["generation_config"] = trainer.model.generation_config.to_dict()
+
+    if hasattr(trainer, "scaler"):
+        checkpoint["scaler"] = move_tensors_to_cpu(trainer.scaler.state_dict())
+
+    if hasattr(trainer.model, "config"):
+        config = trainer.model.config
+        if hasattr(config, "special_tokens_map"):
+            checkpoint["special_tokens_map"] = config.special_tokens_map
+        if hasattr(config, "tokenizer_config"):
+            checkpoint["tokenizer_config"] = config.tokenizer_config
+
+    shm_manager.save_model_object(trainer_config.run_id, SHMObjectType.CHECKPOINTS, checkpoint)
+
+
+def load_checkpoint_from_shared_memory(
+    trainer_config: TrainerConfig, shm_manager: SharedMemoryManager, ref: bool = False, is_peft: bool = False
+) -> tuple[nn.Module, AutoTokenizer, dict]:
+    """Load checkpoint from shared memory"""
+    run_id = trainer_config.run_id
+    device = "cuda:0"
+    base_model = None
+    model_id = trainer_config.config_leaf.get("model_name")
+
+    if trainer_config.warm_started_from is not None and not shm_manager.model_exists(run_id):
+        shm_manager.create_warm_start_checkpoint(run_id, trainer_config.warm_started_from)
+
+    if is_peft:
+        if not shm_manager.model_exists(model_id):
+            base_model, tokenizer = create_model_instance(
+                trainer_config.config_leaf,
+                trainer_config.create_model_fn,
+                checkpoint_path=None,
+                is_peft=is_peft,
+                device=device,
+            )
+            if model_id is None:
+                model_id = base_model.config._name_or_path
+            save_model_to_shared_memory(
+                base_model, tokenizer, trainer_config, shm_manager, SHMObjectType.BASE_MODEL, model_id
+            )
+        else:
+            base_model, tokenizer = load_model_from_shared_memory(
+                trainer_config, shm_manager, SHMObjectType.BASE_MODEL, model_id
+            )
+
+    if base_model == "" or (not is_peft and not shm_manager.model_exists(run_id)):
+        base_model, tokenizer = create_model_instance(
+            trainer_config.config_leaf,
+            trainer_config.create_model_fn,
+            checkpoint_path=None,
+            is_peft=is_peft,
+            device=device,
+        )
+
+    model = base_model
+    peft_config = LoraConfig(**trainer_config.config_leaf.get("peft_params", {}))
+    if is_peft:
+        model = get_peft_model(model, peft_config)
+
+    # Load weights from shared memory
+    if trainer_config.completed_steps > 0 or trainer_config.warm_started_from is not None:
+        checkpoint = shm_manager.load_model_object(run_id, SHMObjectType.CHECKPOINTS)
+
+        if "adapter_config" in checkpoint:
+            if trainer_config.config_leaf.get("trainer_type") == "DPO" and is_peft:
+                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.add_adapter(
+                    trainer_config.config_leaf.get("training_args", {}).get("ref_adapter_name"), peft_config
+                )
+                model.set_adapter(
+                    trainer_config.config_leaf.get("training_args", {}).get("model_adapter_name", "default")
+                )
+                set_peft_model_state_dict(
+                    model,
+                    reference_state_dict,
+                    adapter_name=trainer_config.config_leaf.get("training_args").get("ref_adapter_name"),
+                )
+                model.set_adapter(trainer_config.config_leaf.get("training_args").get("model_adapter_name", "default"))
+
+        if checkpoint.get("state_dict"):
+            state_dict = {k: v.to(device) for k, v in checkpoint["state_dict"].items()}
+            if is_peft:
+                set_peft_model_state_dict(
+                    model,
+                    state_dict,
+                    adapter_name=trainer_config.config_leaf.get("training_args", {}).get(
+                        "model_adapter_name", "default"
+                    ),
+                )
+                if trainer_config.config_leaf.get("trainer_type") == "DPO" and is_peft:
+                    model.set_adapter(
+                        trainer_config.config_leaf.get("training_args", {}).get("model_adapter_name", "default")
+                    )
+            else:
+                model.load_state_dict(state_dict)
+
+        elif not is_peft:
+            model, tokenizer = load_model_from_shared_memory(
+                trainer_config, shm_manager, SHMObjectType.FULL_MODEL, trainer_config.run_id
+            )
+
+    return model, tokenizer
+
+
+def load_model_from_shared_memory(
+    trainer_config: TrainerConfig, shm_manager: SharedMemoryManager, model_object_type: SHMObjectType, model_id: str
+) -> tuple[nn.Module, AutoTokenizer]:
+    """Load model from shared memory"""
+    model_data = shm_manager.load_model_object(model_id, model_object_type)
+    model = copy.deepcopy(model_data[model_object_type])
+    tokenizer = model_data["tokenizer"]
+    bnb_modules = move_tensors_to_device(model_data["bnb_modules"], device="cuda:0")
+    model = get_model_to_device(model, bnb_modules, device="cuda:0")
+    return model, tokenizer
+
+
+def save_model_to_shared_memory(
+    model: Union[nn.Module, str],
+    tokenizer: AutoTokenizer,
+    trainer_config: TrainerConfig,
+    shm_manager: SharedMemoryManager,
+    model_type: str,
+    model_id: str,
+) -> None:
+    """Save model to shared memory"""
+    if model_type != SHMObjectType.FULL_MODEL and shm_manager.model_exists(model_id):
+        return
+    model_cpu = model.cpu()
+    model_data = {model_type: model_cpu, "tokenizer": tokenizer}
+    shm_manager.save_model_object(model_id, model_type, model_data)
+
+
+def load_or_create_ref_model(
+    model_instance,
+    trainer_config: TrainerConfig,
+    device: str,
+    use_shared_memory: bool,
+    shm_manager: SharedMemoryManager,
+) -> Optional[nn.Module]:
+    """Load or create reference model for DPO training based on configuration"""
+    config_leaf = trainer_config.config_leaf
+    device = "cuda:0"
+    ref_model_name = trainer_config.config_leaf.get("ref_model_config", {}).get("model_name", None)
+    model_id = trainer_config.config_leaf.get("model_name")
+    ref_model_id = "ref_" + (trainer_config.config_leaf.get("ref_model_config", {}).get("model_name") or model_id)
+    if use_shared_memory and shm_manager.model_exists(ref_model_id):
+        ref_model_instance, _ = load_model_from_shared_memory(
+            trainer_config, shm_manager, SHMObjectType.REF_FULL_MODEL, ref_model_id
+        )
+    else:
+        if ref_model_name is not None:
+            ref_model_instance, _ = create_model_instance(
+                config_leaf.get("ref_model_config"), trainer_config.create_model_fn, device=device
+            )
+        elif trainer_config.completed_steps == 0:
+            ref_model_instance = copy.deepcopy(model_instance)
+        save_model_to_shared_memory(
+            ref_model_instance, None, trainer_config, shm_manager, SHMObjectType.REF_FULL_MODEL, ref_model_id
+        )
+
+    return ref_model_instance
+
+
+def get_model_to_device(model, bnb_modules, device="cuda:0"):
+    """Move model from shared memory to specified device with proper BitsAndBytes restoration"""
+    for name, param in model.named_parameters():
+        if param.data is not None:
+            param.data = move_tensors_to_device(param.data, device)
+
+    for name, buffer in model.named_buffers():
+        if isinstance(buffer, torch.Tensor) and buffer is not None:
+            parent_module = model
+            attr_path = name.split(".")
+
+            for attr in attr_path[:-1]:
+                parent_module = getattr(parent_module, attr)
+
+            device_buffer = move_tensors_to_device(buffer, device)
+            setattr(parent_module, attr_path[-1], device_buffer)
+
+    for name, module in model.named_modules():
+        if not hasattr(module, "weight"):
+            continue
+
+        try:
+            import bitsandbytes as bnb
+
+            bnb_layer_types = [bnb.nn.Linear4bit, bnb.nn.LinearFP4, bnb.nn.LinearNF4, bnb.nn.Params4bit]
+        except ImportError:
+            continue
+
+        is_bnb_layer = any(isinstance(module, layer_type) for layer_type in bnb_layer_types)
+
+        if is_bnb_layer and name in bnb_modules:
+            bnb_attrs = bnb_modules[name]
+            weight = module.weight
+
+            if hasattr(weight, "data") and weight.data is not None:
+                weight.data = move_tensors_to_device(weight.data, device)
+
+            if "quant_state_data" in bnb_attrs:
+                if not hasattr(weight, "quant_state") or weight.quant_state is None:
+                    from bitsandbytes.functional import QuantState
+
+                    quant_data = bnb_attrs["quant_state_data"]
+
+                    weight.quant_state = QuantState(absmax=quant_data["absmax"], code=quant_data["code"])
+
+                quant_data = bnb_attrs["quant_state_data"]
+
+                for attr, value in quant_data.items():
+                    if isinstance(value, torch.Tensor):
+                        value = move_tensors_to_device(value, device)
+                    setattr(weight.quant_state, attr, value)
+                if "state2_data" in bnb_attrs:
+                    for attr, value in bnb_attrs["state2_data"].items():
+                        if isinstance(value, torch.Tensor):
+                            value = move_tensors_to_device(value, device)
+                        setattr(weight.quant_state.state2, attr, value)
+
+            for attr, value in bnb_attrs.items():
+                if attr not in ["quant_state_data", "quant_state_class", "weight_class", "state2_data"]:
+                    if isinstance(value, torch.Tensor):
+                        value = move_tensors_to_device(value, device)
+                    setattr(weight, attr, value)
+
+        elif hasattr(module, "weight") and hasattr(module.weight, "data") and module.weight.data is not None:
+            if name not in bnb_modules:
+                module.weight.data = move_tensors_to_device(module.weight.data, device)
+
+    model = model.to(device)
+    return model
+
+
+def restore_trainer_from_shared_memory(
+    trainer: Union[SFTTrainer, DPOTrainer, GRPOTrainer],
+    trainer_config: TrainerConfig,
+    shm_manager: SharedMemoryManager,
+) -> Union[SFTTrainer, DPOTrainer, GRPOTrainer]:
+    """Restore complete training state to trainer"""
+    try:
+        device = next(trainer.model.parameters()).device
+
+        if shm_manager.model_exists(trainer_config.run_id):
+            training_state = shm_manager.load_model_object(trainer_config.run_id, SHMObjectType.CHECKPOINTS)
+        else:
+            raise ValueError(f"Training state for run {trainer_config.run_id} not found in shared memory")
+
+        if training_state.get("trainer_state") is not None and hasattr(trainer, "state"):
+            trainer_state_dict = training_state["trainer_state"]
+            device_trainer_state = move_tensors_to_device(trainer_state_dict, device)
+            for key, value in device_trainer_state.items():
+                if hasattr(trainer.state, key):
+                    setattr(trainer.state, key, value)
+
+        if training_state.get("optimizer_state") is not None and trainer.optimizer is not None:
+            optimizer_state = training_state["optimizer_state"]
+            device_optimizer_state = move_tensors_to_device(optimizer_state, device)
+            trainer.optimizer.load_state_dict(device_optimizer_state)
+
+        if training_state.get("scheduler_state") is not None and trainer.lr_scheduler is not None:
+            scheduler_state = training_state["scheduler_state"]
+            device_scheduler_state = move_tensors_to_device(scheduler_state, device)
+            trainer.lr_scheduler.load_state_dict(device_scheduler_state)
+
+            if hasattr(trainer.state, "global_step"):
+                trainer.lr_scheduler._step_count = trainer.state.global_step + 1
+
+        if training_state.get("rng_state") is not None:
+            rng_state = training_state["rng_state"]
+            device_rng_state = move_tensors_to_device(rng_state, device)
+            trainer.rng_state = device_rng_state
+
+        if training_state.get("generation_config") is not None and hasattr(trainer.model, "generation_config"):
+            trainer.model.generation_config = type(trainer.model.generation_config)(
+                **training_state["generation_config"]
+            )
+        if training_state.get("scaler") is not None:
+            trainer.scaler.load_state_dict(training_state["scaler"])
+
+    except Exception as e:
+        print(f"Warning: Error restoring training state: {e}")
+
+    return trainer
+
+
+def save_checkpoint_to_disk(
+    trainer: Union[SFTTrainer, DPOTrainer, GRPOTrainer],
+    trainer_config: TrainerConfig,
+    first: bool = False,
+    last: bool = False,
+    completed_steps: int = 0,
+) -> None:
+    base_run_path = DataPath.base_run_path(trainer_config.run_id)
+    if first:
+        checkpoint_path = DataPath.initial_checkpoint_path(base_run_path)
+    elif last:
+        checkpoint_path = DataPath.final_checkpoint_path(base_run_path)
+    else:
+        checkpoint_path = DataPath.intermediate_checkpoint_path(base_run_path) / "checkpoint"
+    os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
+    os.makedirs(checkpoint_path, exist_ok=True)
+
+    trainer.model.save_pretrained(checkpoint_path)
+
+    trainer_state_path = checkpoint_path / "trainer_state.json"
+    trainer_state_dict = trainer.state.__dict__.copy()
+    with open(trainer_state_path, "w") as f:
+        json.dump(trainer_state_dict, f, indent=2)
+
+    if trainer.optimizer is not None:
+        optimizer_path = checkpoint_path / "optimizer.pt"
+        optimizer_state = move_tensors_to_cpu(trainer.optimizer.state_dict())
+        torch.save(optimizer_state, optimizer_path)
+
+    if trainer.lr_scheduler is not None:
+        scheduler_path = checkpoint_path / "scheduler.pt"
+        scheduler_state = move_tensors_to_cpu(trainer.lr_scheduler.state_dict())
+        torch.save(scheduler_state, scheduler_path)
+
+    if hasattr(trainer, "rng_state"):
+        rng_state_path = checkpoint_path / "rng_state.pth"
+        torch.save(trainer.rng_state, rng_state_path)
+
+
+def load_checkpoint_from_disk(
+    trainer_config: TrainerConfig, ref: bool = False, is_peft: bool = False
+) -> tuple[nn.Module, AutoTokenizer, dict]:
+    """Load checkpoint from disk"""
+    device = "cuda:0"
+    checkpoint_path = None
+    if trainer_config.warm_started_from is not None and trainer_config.completed_steps == 0:
+        base_run_path = DataPath.base_run_path(trainer_config.warm_started_from)
+        checkpoint_path = DataPath.intermediate_checkpoint_path(base_run_path) / "checkpoint"
+    elif trainer_config.completed_steps > 0:
+        base_run_path = DataPath.base_run_path(trainer_config.run_id)
+        checkpoint_path = DataPath.intermediate_checkpoint_path(base_run_path) / "checkpoint"
+
+    model_instance, tokenizer = create_model_instance(
+        trainer_config.config_leaf, trainer_config.create_model_fn, checkpoint_path, is_peft=is_peft, device=device
+    )
+    if is_peft and checkpoint_path is None:
+        model_instance = get_peft_model(model_instance, LoraConfig(**trainer_config.config_leaf.get("peft_params", {})))
+
+    if ref:
+        model_instance, tokenizer = create_model_instance(
+            trainer_config.config_leaf.get("ref_model_config", {}), trainer_config.create_model_fn, device=device
+        )
+
+    return model_instance, tokenizer
+
+
+def restore_trainer_from_disk(
+    trainer: Union[SFTTrainer, DPOTrainer, GRPOTrainer], trainer_config: TrainerConfig
+) -> Union[SFTTrainer, DPOTrainer, GRPOTrainer]:
+    """Restore trainer from disk with proper state accumulation"""
+    base_run_path = DataPath.base_run_path(trainer_config.run_id)
+    checkpoint_path = DataPath.intermediate_checkpoint_path(base_run_path) / "checkpoint"
+    device = "cuda:0"
+
+    trainer_state_path = checkpoint_path / "trainer_state.json"
+    if trainer_state_path.exists():
+        with open(trainer_state_path) as f:
+            trainer_state_dict = json.load(f)
+
+        for key, value in trainer_state_dict.items():
+            if hasattr(trainer.state, key):
+                setattr(trainer.state, key, value)
+
+    optimizer_path = checkpoint_path / "optimizer.pt"
+    if optimizer_path.exists() and trainer.optimizer is not None:
+        optimizer_state = torch.load(optimizer_path, map_location=device)
+        model_device = next(trainer.model.parameters()).device
+        optimizer_state = move_tensors_to_device(optimizer_state, model_device)
+        trainer.optimizer.load_state_dict(optimizer_state)
+
+    lr_scheduler_path = checkpoint_path / "scheduler.pt"
+    if lr_scheduler_path.exists() and trainer.lr_scheduler is not None:
+        lr_scheduler_state = torch.load(lr_scheduler_path, map_location=device)
+        model_device = next(trainer.model.parameters()).device
+        lr_scheduler_state = move_tensors_to_device(lr_scheduler_state, model_device)
+        trainer.lr_scheduler.load_state_dict(lr_scheduler_state)
+
+        if hasattr(trainer.state, "global_step") and trainer.lr_scheduler is not None:
+            trainer.lr_scheduler._step_count = trainer.state.global_step + 1
+
+    rng_state_path = checkpoint_path / "rng_state.pth"
+    if rng_state_path.exists():
+        rng_state = torch.load(rng_state_path, map_location=device, weights_only=False)
+        model_device = next(trainer.model.parameters()).device
+        rng_state = move_tensors_to_device(rng_state, model_device)
+        trainer.rng_state = rng_state
+
+    return trainer
+
+
+def save_ref_model_to_disk(model_instance: nn.Module, trainer_config: TrainerConfig, ref: bool = False) -> None:
+    """Save reference model to disk"""
+    base_run_path = DataPath.base_run_path(trainer_config.run_id)
+    ref_model_path = DataPath.ref_model_path(base_run_path)
+    os.makedirs(ref_model_path, exist_ok=True)
+
+    if hasattr(model_instance, "peft_config"):
+        peft_state_dict = get_peft_model_state_dict(model_instance)
+        torch.save(peft_state_dict, ref_model_path / "pytorch_model.bin")
+    else:
+        torch.save(model_instance.state_dict(), ref_model_path / "pytorch_model.bin")