Distillflow 0.2.0__py3-none-any.whl

distillflow/model/checkpoint.py

@@ -0,0 +1,146 @@
+import inspect
+from functools import wraps, partial
+from types import MethodType
+from typing import Callable, Union, Any, Optional, Dict, Tuple
+
+import torch
+from transformers import PreTrainedModel
+
+from .args import ModelArgs
+from ..common import get_logger
+
+logger = get_logger(__name__)
+
+LAYERNORM_NAMES = {"norm", "ln"}
+
+def get_unsloth_gradient_checkpointing_func() -> Callable:
+    class UnslothGradientCheckpointing(torch.autograd.Function):
+        r"""
+        Saves VRAM by smartly offloading to RAM.
+        """
+
+        @staticmethod
+        @torch.cuda.amp.custom_fwd
+        def forward(
+                ctx: torch.autograd.Function,
+                forward_function: torch.Module,
+                hidden_states: torch.Tensor,
+                *args: Union[torch.Tensor, Any],
+        ) -> torch.Tensor:
+            saved_hidden_states = hidden_states.to("cpu", non_blocking=True)
+            with torch.no_grad():
+                output = forward_function(hidden_states, *args)
+
+            ctx.save_for_backward(saved_hidden_states)
+            ctx.forward_function = forward_function
+            ctx.args = args
+            return output
+
+        @staticmethod
+        @torch.cuda.amp.custom_bwd
+        def backward(ctx: "torch.autograd.Function", grad_output: "torch.Tensor") -> "torch.Tensor":
+            (hidden_states,) = ctx.saved_tensors
+            hidden_states = hidden_states.to("cuda", non_blocking=True).detach()
+            hidden_states.requires_grad_(True)
+            with torch.enable_grad():
+                (output,) = ctx.forward_function(hidden_states, *ctx.args)
+
+            torch.autograd.backward(output, grad_output)
+            return (None, hidden_states.grad) + (None,) * len(ctx.args)
+
+    return UnslothGradientCheckpointing.apply
+
+
+def get_custom_gradient_checkpointing_func(gradient_checkpointing_func: Callable) -> Callable:
+    r"""
+    Only applies gradient checkpointing to trainable layers.
+    """
+
+    @wraps(gradient_checkpointing_func)
+    def custom_gradient_checkpointing_func(func: Callable, *args: Union["torch.Tensor", Any], **kwargs):
+        underlying = func.func if isinstance(func, partial) else func
+        module: "torch.nn.Module" = underlying.__self__
+
+        if any(param.requires_grad for param in module.parameters()):
+            for arg in args:
+                if torch.is_tensor(arg) and torch.is_floating_point(arg):
+                    arg.requires_grad_(True)
+
+        return gradient_checkpointing_func(func, *args, **kwargs)
+
+    if hasattr(gradient_checkpointing_func, "__self__"):  # fix unsloth gc test case
+        custom_gradient_checkpointing_func.__self__ = gradient_checkpointing_func.__self__
+
+    return custom_gradient_checkpointing_func
+
+
+def _gradient_checkpointing_enable(
+        self: PreTrainedModel,
+        gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
+        use_unsloth_gc: bool = False,
+) -> None:
+    r"""
+    Activates gradient checkpointing for the current model.
+
+    Modification of the original method to enable gradient checkpointing for block-wise optimizer.
+    """
+    from torch.utils.checkpoint import checkpoint
+
+    if not self.supports_gradient_checkpointing:
+        raise ValueError("{} does not support gradient checkpointing.".format(self.__class__.__name__))
+
+    if gradient_checkpointing_kwargs is None:
+        gradient_checkpointing_kwargs = {"use_reentrant": True}
+
+    if use_unsloth_gc:
+        gradient_checkpointing_func = get_unsloth_gradient_checkpointing_func()
+    else:
+        gradient_checkpointing_func = partial(checkpoint, **gradient_checkpointing_kwargs)
+
+    gradient_checkpointing_func = get_custom_gradient_checkpointing_func(gradient_checkpointing_func)
+    if "value" in inspect.signature(self._set_gradient_checkpointing).parameters:  # old GC format
+        self.apply(partial(self._set_gradient_checkpointing, value=True))
+        self.enable_input_require_grads()
+        logger.warning("You are using the old GC format, some features (e.g. BAdam) will be invalid.")
+    else:  # have already enabled input require gradients
+        self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func)
+
+
+def _fp32_forward_post_hook(
+        module: torch.nn.Module, args: Tuple[torch.Tensor], output: torch.Tensor
+) -> torch.Tensor:
+    return output.to(torch.float32)
+
+
+def prepare_model_for_training(model: "PreTrainedModel", model_args: ModelArgs) -> None:
+    r"""
+    Includes:
+        (1) cast the layernorm in fp32
+        (2) make output embedding layer require grads
+        (3) add the upcasting of the lm_head in fp32
+    """
+    if model_args.upcast_layernorm:
+        logger.info("Upcasting layernorm weights in float32.")
+        for name, param in model.named_parameters():
+            if param.ndim == 1 and any(ln_name in name for ln_name in LAYERNORM_NAMES):
+                param.data = param.data.to(torch.float32)
+
+    if not model_args.disable_gradient_checkpointing:
+        if not getattr(model, "supports_gradient_checkpointing", False):
+            logger.warning("Current model does not support gradient checkpointing.")
+        else:
+            # use_reentrant=False might increase VRAM usage (have not been empirically verified yet)
+            # According to: https://github.com/huggingface/transformers/issues/28339
+            gradient_checkpointing_enable = partial(
+                _gradient_checkpointing_enable, use_unsloth_gc=model_args.use_unsloth_gc
+            )
+            model.gradient_checkpointing_enable = MethodType(gradient_checkpointing_enable, model)
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True})
+            setattr(model.config, "use_cache", False)  # turn off when gradient checkpointing is enabled
+            logger.info("Gradient checkpointing enabled.")
+
+    if model_args.upcast_lmhead_output:
+        output_layer = model.get_output_embeddings()
+        if isinstance(output_layer, torch.nn.Linear) and output_layer.weight.dtype != torch.float32:
+            logger.info("Upcasting lm_head outputs in float32.")
+            output_layer.register_forward_hook(_fp32_forward_post_hook)

distillflow/model/finetuning_args.py

@@ -0,0 +1,132 @@
+from typing import Literal, List, Optional
+
+from pydantic import Field, BaseModel
+
+
+class FreezeArgs(BaseModel):
+    r"""
+    Arguments pertaining to the freeze (partial-parameter) training.
+    """
+
+    freeze_trainable_layers: int = Field(
+        default=2,
+        description=(
+                "The number of trainable layers for freeze (partial-parameter) fine-tuning. "
+                "Positive numbers mean the last n layers are set as trainable, "
+                "negative numbers mean the first n layers are set as trainable."
+        )
+    )
+    freeze_trainable_modules: str = Field(
+        default="all",
+        description=(
+                "Name(s) of trainable modules for freeze (partial-parameter) fine-tuning. "
+                "Use commas to separate multiple modules. "
+                "Use `all` to specify all the available modules."
+        )
+    )
+    freeze_extra_modules: Optional[str] = Field(
+        default=None,
+        description=(
+                "Name(s) of modules apart from hidden layers to be set as trainable "
+                "for freeze (partial-parameter) fine-tuning. "
+                "Use commas to separate multiple modules."
+        )
+    )
+
+
+class LoraArgs(BaseModel):
+    r"""
+    Arguments pertaining to the LoRA training.
+    """
+
+    additional_target: Optional[str] = Field(
+        default=None,
+        description=(
+                "Name(s) of modules apart from LoRA layers to be set as trainable "
+                "and saved in the final checkpoint. "
+                "Use commas to separate multiple modules."
+        )
+    )
+    lora_alpha: Optional[int] = Field(
+        default=None,
+        description="The scale factor for LoRA fine-tuning (default: lora_rank * 2)."
+    )
+    lora_dropout: float = Field(
+        default=0.0,
+        description="Dropout rate for the LoRA fine-tuning."
+    )
+    lora_rank: int = Field(
+        default=8,
+        description="The intrinsic dimension for LoRA fine-tuning."
+    )
+    lora_target: str = Field(
+        default="all",
+        description=(
+                "Name(s) of target modules to apply LoRA. "
+                "Use commas to separate multiple modules. "
+                "Use `all` to specify all the linear modules."
+        )
+    )
+    loraplus_lr_ratio: Optional[float] = Field(
+        default=None,
+        description="LoRA plus learning rate ratio (lr_B / lr_A)."
+    )
+    loraplus_lr_embedding: float = Field(
+        default=1e-6,
+        description="LoRA plus learning rate for lora embedding layers."
+    )
+    use_rslora: bool = Field(
+        default=False,
+        description="Whether or not to use the rank stabilization scaling factor for LoRA layer."
+    )
+    use_dora: bool = Field(
+        default=False,
+        description="Whether or not to use the weight-decomposed lora method (DoRA)."
+    )
+    pissa_init: bool = Field(
+        default=False,
+        description="Whether or not to initialize a PiSSA adapter."
+    )
+    pissa_iter: int = Field(
+        default=16,
+        description="The number of iteration steps performed by FSVD in PiSSA. Use -1 to disable it."
+    )
+    pissa_convert: bool = Field(
+        default=False,
+        description="Whether or not to convert the PiSSA adapter to a normal LoRA adapter."
+    )
+    create_new_adapter: bool = Field(
+        default=False,
+        description="Whether or not to create a new adapter with randomly initialized weight."
+    )
+
+
+class FinetuningArgs(FreezeArgs, LoraArgs, BaseModel):
+    r"""
+    Arguments pertaining to which techniques we are going to fine-tuning with.
+    """
+
+    finetuning_type: Literal["lora", "freeze", "full"] = Field(
+        default="full",
+        description="Which fine-tuning method to use."
+    )
+    use_llama_pro: bool = Field(
+        default=False,
+        description="Whether or not to make only the parameters in the expanded blocks trainable."
+    )
+
+    def __post_init__(self):
+        def split_arg(arg):
+            if isinstance(arg, str):
+                return [item.strip() for item in arg.split(",")]
+            return arg
+
+        self.freeze_trainable_modules: List[str] = split_arg(self.freeze_trainable_modules)
+        self.freeze_extra_modules: Optional[List[str]] = split_arg(self.freeze_extra_modules)
+        self.lora_alpha: int = self.lora_alpha or self.lora_rank * 2
+        self.lora_target: List[str] = split_arg(self.lora_target)
+        self.additional_target: Optional[List[str]] = split_arg(self.additional_target)
+
+        assert self.finetuning_type in ["lora", "freeze", "full"], "Invalid fine-tuning method."
+        if self.use_llama_pro and self.finetuning_type == "full":
+            raise ValueError("`use_llama_pro` is only valid for Freeze or LoRA training.")

distillflow/model/generating_args.py

@@ -0,0 +1,72 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, Optional
+
+from transformers import GenerationConfig
+
+@dataclass
+class GeneratingArguments:
+    r"""
+    Arguments pertaining to specify the decoding parameters.
+    """
+
+    do_sample: bool = field(
+        default=True,
+        metadata={"help": "Whether or not to use sampling, use greedy decoding otherwise."},
+    )
+    temperature: float = field(
+        default=0.95,
+        metadata={"help": "The value used to modulate the next token probabilities."},
+    )
+    top_p: float = field(
+        default=0.7,
+        metadata={
+            "help": "The smallest set of most probable tokens with probabilities that add up to top_p or higher are kept."
+        },
+    )
+    top_k: int = field(
+        default=50,
+        metadata={"help": "The number of highest probability vocabulary tokens to keep for top-k filtering."},
+    )
+    num_beams: int = field(
+        default=1,
+        metadata={"help": "Number of beams for beam search. 1 means no beam search."},
+    )
+    max_length: int = field(
+        default=1024,
+        metadata={"help": "The maximum length the generated tokens can have. It can be overridden by max_new_tokens."},
+    )
+    max_new_tokens: int = field(
+        default=1024,
+        metadata={"help": "The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt."},
+    )
+    repetition_penalty: float = field(
+        default=1.0,
+        metadata={"help": "The parameter for repetition penalty. 1.0 means no penalty."},
+    )
+    length_penalty: float = field(
+        default=1.0,
+        metadata={"help": "Exponential penalty to the length that is used with beam-based generation."},
+    )
+    default_system: Optional[str] = field(
+        default=None,
+        metadata={"help": "Default system message to use in chat completion."},
+    )
+    skip_special_tokens: bool = field(
+        default=True,
+        metadata={"help": "Whether or not to remove special tokens in the decoding."},
+    )
+
+    def to_dict(self, obey_generation_config: bool = False) -> Dict[str, Any]:
+        args = asdict(self)
+        if args.get("max_new_tokens", -1) > 0:
+            args.pop("max_length", None)
+        else:
+            args.pop("max_new_tokens", None)
+
+        if obey_generation_config:
+            generation_config = GenerationConfig()
+            for key in list(args.keys()):
+                if not hasattr(generation_config, key):
+                    args.pop(key)
+
+        return args

distillflow/model/liger_kernel.py

@@ -0,0 +1,46 @@
+import inspect
+from transformers import PretrainedConfig
+from ..common import logger
+
+from distillflow.model.args import ModelArgs
+
+logger = logger.get_logger(__name__)
+
+def apply_liger_kernel(
+    config: PretrainedConfig,
+    model_args: ModelArgs,
+    is_trainable: bool,
+    require_logits: bool,
+) -> None:
+    if not is_trainable or not model_args.enable_liger_kernel:
+        return
+
+    model_type = getattr(config, "model_type", None)
+    if model_type == "gemma":
+        from liger_kernel.transformers import apply_liger_kernel_to_gemma as apply_liger_kernel
+    elif model_type == "gemma2":
+        from liger_kernel.transformers import apply_liger_kernel_to_gemma2 as apply_liger_kernel
+    elif model_type == "llama":
+        from liger_kernel.transformers import apply_liger_kernel_to_llama as apply_liger_kernel
+    elif model_type == "mistral":
+        from liger_kernel.transformers import apply_liger_kernel_to_mistral as apply_liger_kernel
+    elif model_type == "mixtral":
+        from liger_kernel.transformers import apply_liger_kernel_to_mixtral as apply_liger_kernel
+    elif model_type == "phi3":
+        from liger_kernel.transformers import apply_liger_kernel_to_phi3 as apply_liger_kernel
+    elif model_type == "qwen2":
+        from liger_kernel.transformers import apply_liger_kernel_to_qwen2 as apply_liger_kernel
+    elif model_type == "qwen2_vl":
+        from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl as apply_liger_kernel
+    else:
+        logger.warning("Current model does not support liger kernel.")
+        return
+
+    if require_logits and "fused_linear_cross_entropy" in inspect.signature(apply_liger_kernel).parameters:
+        logger.info("Current training stage does not support chunked cross entropy.")
+        kwargs = {"fused_linear_cross_entropy": False}
+    else:
+        kwargs = {}
+
+    apply_liger_kernel(**kwargs)
+    logger.info("Liger kernel has been applied to the model.")

distillflow/model/loader.py

@@ -0,0 +1,286 @@
+import math
+import os
+from contextlib import nullcontext
+from types import MethodType
+from typing import Dict, Any
+
+from transformers import PreTrainedTokenizer, AutoConfig, PreTrainedModel, \
+    PretrainedConfig, AutoModelForCausalLM, is_torch_npu_available
+from transformers.integrations import is_deepspeed_zero3_enabled, is_deepspeed_available
+from transformers.modeling_utils import is_fsdp_enabled
+from transformers.utils import is_torch_sdpa_available, is_flash_attn_2_available
+from transformers.utils.versions import require_version
+
+from .adapter import init_adapter
+from .checkpoint import prepare_model_for_training
+from .liger_kernel import apply_liger_kernel
+from .quantization import _configure_quantization, QuantizationMethod
+from .tokenizer import load_tokenizer
+from .unsloth import load_unsloth_pretrained_model
+from ..common.logger import get_logger
+from .args import ModelArgs
+import torch
+
+from ..common import count_parameters, infer_optim_dtype, get_current_device
+
+logger = get_logger(__name__)
+
+def get_init_kwargs(model_args: ModelArgs) -> Dict[str, Any]:
+    r"""
+    Gets arguments to load config/tokenizer/model.
+
+    Note: including inplace operation of model_args.
+    """
+    return {
+        "trust_remote_code": True,
+        "cache_dir": model_args.cache_dir,
+        "revision": model_args.model_revision,
+        "token": model_args.hf_hub_token
+    }
+
+def _register_autoclass(config: PretrainedConfig, model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer"):
+    if "AutoConfig" in getattr(config, "auto_map", {}):
+        config.__class__.register_for_auto_class()
+    if "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
+        model.__class__.register_for_auto_class()
+    if "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
+        tokenizer.__class__.register_for_auto_class()
+
+def _configure_attn_implementation(
+        config: PretrainedConfig, model_args: ModelArgs, is_trainable: bool
+) -> None:
+    if getattr(config, "model_type", None) == "gemma2" and is_trainable:
+        if model_args.flash_attn == "auto" or model_args.flash_attn == "fa2":
+            if is_flash_attn_2_available():
+                require_version("flash_attn>=2.6.3", "To fix: pip install flash_attn>=2.6.3")
+                if model_args.flash_attn != "fa2":
+                    logger.warning("Gemma-2 should use flash attention 2, change `flash_attn` to fa2.")
+                    model_args.flash_attn = "fa2"
+            else:
+                logger.warning("FlashAttention-2 is not installed, use eager attention.")
+                model_args.flash_attn = "disabled"
+        elif model_args.flash_attn == "sdpa":
+            logger.warning("Gemma-2 should use soft-capping attention, while the SDPA attention does not support it.")
+
+    if model_args.flash_attn == "auto":
+        return
+
+    elif model_args.flash_attn == "disabled":
+        requested_attn_implementation = "eager"
+
+    elif model_args.flash_attn == "sdpa":
+        if not is_torch_sdpa_available():
+            logger.warning("torch>=2.1.1 is required for SDPA attention.")
+            return
+
+        requested_attn_implementation = "sdpa"
+    elif model_args.flash_attn == "fa2":
+        if not is_flash_attn_2_available():
+            logger.warning("FlashAttention-2 is not installed.")
+            return
+
+        requested_attn_implementation = "flash_attention_2"
+    else:
+        raise NotImplementedError("Unknown attention type: {}".format(model_args.flash_attn))
+
+    if getattr(config, "model_type", None) == "internlm2":  # special case for custom models
+        setattr(config, "attn_implementation", requested_attn_implementation)
+    else:
+        setattr(config, "_attn_implementation", requested_attn_implementation)
+
+def _patch_config(
+        config: PretrainedConfig,
+        tokenizer: PreTrainedTokenizer,
+        model_args: ModelArgs,
+        init_kwargs: Dict[str, Any],
+        is_trainable: bool,
+        torch_dtype: torch.dtype
+) -> None:
+    if is_torch_npu_available():
+        use_jit_compile = os.environ.get("JIT_COMPILE", "0").lower() in ["true", "1"]
+        torch.npu.set_compile_mode(jit_compile=use_jit_compile)
+
+    _configure_attn_implementation(config, model_args, is_trainable)
+    _configure_quantization(config, tokenizer, model_args.quantization_args, init_kwargs, torch_dtype)
+
+    if model_args.use_cache and not is_trainable:
+        setattr(config, "use_cache", True)
+        logger.info("Using KV cache for faster generation.")
+
+    if getattr(config, "model_type", None) == "qwen":
+        setattr(config, "use_flash_attn", model_args.flash_attn == "fa2")
+        for dtype_name, dtype in [("fp16", torch.float16), ("bf16", torch.bfloat16), ("fp32", torch.float32)]:
+            setattr(config, dtype_name, torch_dtype == dtype)
+
+    if getattr(config, "model_type", None) == "qwen2" and is_trainable and model_args.flash_attn == "fa2":
+        setattr(config, "use_cache", False)  # qwen2 does not support use_cache when using flash attn
+
+    # if "LlavaLlamaForCausalLM" in getattr(config, "architectures", []):
+    #     raise ValueError("Please download llava models with hf-compatible format: https://huggingface.co/llava-hf")
+
+    # deepspeed zero3 is not compatible with low_cpu_mem_usage
+
+    init_kwargs["low_cpu_mem_usage"] = model_args.low_cpu_mem_usage and (not is_deepspeed_available())
+
+    # cast data type of the model if:
+    # 1. not deepspeed zero3 and not fsdp (keep zero3 or fsdp in float32)
+    # 2. quantization_bit is not None (qlora)
+
+    quantization_args = model_args.quantization_args
+    if (not is_deepspeed_zero3_enabled() and not is_fsdp_enabled()) or quantization_args.quantization_bit is not None:
+        init_kwargs["torch_dtype"] = torch_dtype
+
+        if init_kwargs["low_cpu_mem_usage"]:  # device map requires low_cpu_mem_usage=True
+            if "device_map" not in init_kwargs:
+                init_kwargs["device_map"] = {"": get_current_device()}
+
+            if init_kwargs.get("device_map", None) == "auto":
+                init_kwargs["offload_folder"] = model_args.offload_folder
+
+
+def _noisy_mean_initialization(embed_weight: "torch.Tensor", num_new_tokens: int) -> None:
+    embedding_dim = embed_weight.size(1)
+    avg_weight = embed_weight[:-num_new_tokens].mean(dim=0, keepdim=True)
+    noise_weight = torch.empty_like(embed_weight[-num_new_tokens:])
+    noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
+    embed_weight[-num_new_tokens:] = avg_weight + noise_weight
+
+def _resize_embedding_layer(model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer") -> None:
+    r"""
+    Resize token embeddings.
+    """
+    # if is_deepspeed_zero3_enabled():
+    #     import deepspeed  # type: ignore
+    #
+    #     params = [model.get_input_embeddings().weight]
+    #     if model.get_output_embeddings() is not None and not model.config.tie_word_embeddings:
+    #         params.append(model.get_output_embeddings().weight)
+    #
+    #     context_maybe_zero3 = deepspeed.zero.GatheredParameters(params, modifier_rank=0)
+    # else:
+    context_maybe_zero3 = nullcontext()
+
+    with context_maybe_zero3:
+        current_embedding_size = model.get_input_embeddings().weight.size(0)
+
+    if len(tokenizer) > current_embedding_size:
+        if getattr(model, "quantization_method", None):
+            raise ValueError("Cannot resize embedding layers of a quantized model.")
+
+        if not isinstance(model.get_output_embeddings(), torch.nn.Linear):
+            raise ValueError("Current model does not support resizing embedding layers.")
+
+        model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
+        with context_maybe_zero3:
+            new_embedding_size = model.get_input_embeddings().weight.size(0)
+            num_new_tokens = new_embedding_size - current_embedding_size
+            _noisy_mean_initialization(model.get_input_embeddings().weight.data, num_new_tokens)
+            _noisy_mean_initialization(model.get_output_embeddings().weight.data, num_new_tokens)
+
+        logger.info("Resized token embeddings from {} to {}.".format(current_embedding_size, new_embedding_size))
+
+def _patch_model(
+        model: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        model_args: ModelArgs,
+        is_trainable: bool,
+) -> None:
+    gen_config = model.generation_config  # check and fix generation config
+    if not gen_config.do_sample and (
+            (gen_config.temperature is not None and gen_config.temperature != 1.0)
+            or (gen_config.top_p is not None and gen_config.top_p != 1.0)
+            or (gen_config.typical_p is not None and gen_config.typical_p != 1.0)
+    ):
+        gen_config.do_sample = True
+
+    if "GenerationMixin" not in str(model.generate.__func__):
+        model.generate = MethodType(PreTrainedModel.generate, model)
+
+    if model_args.resize_vocab:
+        _resize_embedding_layer(model, tokenizer)
+
+    if is_trainable:
+        prepare_model_for_training(model, model_args)
+        # autocast_projector_dtype(model, model_args)
+        # add_z3_leaf_module(model)
+
+    attn = getattr(model.config, "_attn_implementation", None)
+    if attn == "flash_attention_2":
+        logger.info("Using FlashAttention-2 for faster training and inference.")
+    elif attn == "sdpa":
+        logger.info("Using torch SDPA for faster training and inference.")
+
+def load_model(
+        model_args: ModelArgs,
+        is_trainable: bool = False,
+) -> (PreTrainedModel, PreTrainedTokenizer):
+    tokenizer = load_tokenizer(model_args)
+    init_kwargs = get_init_kwargs(model_args)
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **init_kwargs)
+    if model_args.infer_dtype != "auto" and not is_trainable:
+        torch_dtype = getattr(torch, model_args.infer_dtype)
+    else:
+        torch_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
+
+    _patch_config(config, tokenizer, model_args, init_kwargs, is_trainable, torch_dtype)
+
+    # More details here: https://github.com/linkedin/Liger-Kernel
+    apply_liger_kernel(config, model_args, is_trainable, require_logits=True)
+
+    model = None
+    lazy_load = False
+    if model_args.use_unsloth:
+        if model_args.adapter_name_or_path is not None:
+            lazy_load = True
+        elif is_trainable:
+            model = load_unsloth_pretrained_model(config, model_args)
+
+    quantization_args = model_args.quantization_args
+    if quantization_args.quantization_method == QuantizationMethod.GPTQ.value:
+        from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
+        quantize_config = BaseQuantizeConfig(
+            bits=quantization_args.quantization_bit,
+            group_size=128,       # Group size (optional, can be None)
+            desc_act=False        # Disable activation descriptor (optional)
+        )
+        model = AutoGPTQForCausalLM.from_pretrained(model_args.model_name_or_path, bits=quantization_args.quantization_bit, group_size=128, quantize_config=quantize_config)
+
+    if model is None and not lazy_load:
+        init_kwargs["config"] = config
+        init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path
+        model = AutoModelForCausalLM.from_pretrained(**init_kwargs)
+
+    if not lazy_load:
+        _patch_model(model, tokenizer, model_args, is_trainable)
+        _register_autoclass(config, model, tokenizer)
+
+    model = init_adapter(config, model, model_args, is_trainable)
+    if not is_trainable:
+        model.requires_grad_(False)
+        for param in model.parameters():
+            if param.data.dtype == torch.float32 and torch_dtype != torch.float32:
+                param.data = param.data.to(torch_dtype)
+
+        model.eval()
+    else:
+        model.train()
+
+    trainable_params, all_param = count_parameters(model)
+    if is_trainable:
+        param_stats = "trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
+            trainable_params, all_param, 100 * trainable_params / all_param
+        )
+    else:
+        param_stats = "all params: {:,}".format(all_param)
+
+    logger.info(param_stats)
+
+    if model_args.print_param_status:
+        for name, param in model.named_parameters():
+            print(
+                "name: {}, dtype: {}, device: {}, trainable: {}".format(
+                    name, param.dtype, param.device, param.requires_grad
+                )
+            )
+
+    return model, tokenizer