nemo/lightning/pytorch/callbacks/peft.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import re
from abc import ABC, abstractmethod
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple

import lightning.pytorch as pl
import torch
import torch.nn as nn
from lightning.fabric.utilities.types import _PATH
from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
from lightning.pytorch.trainer.states import TrainerFn
from typing_extensions import override

from nemo.lightning.ckpt_utils import ADAPTER_META_FILENAME, HF_ADAPTER_CONFIG_FILENAME, HF_ADAPTER_PATH
from nemo.lightning.io.mixin import IOMixin
from nemo.lightning.io.pl import ckpt_to_dir, ckpt_to_weights_subdir
from nemo.lightning.megatron_parallel import MegatronParallel
from nemo.lightning.pytorch.callbacks.model_transform import ModelTransform
from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
from nemo.lightning.pytorch.utils import get_automodel_from_trainer, is_trainer_attached
from nemo.utils import logging
from nemo.utils.callbacks.dist_ckpt_io import AsyncCompatibleCheckpointIO

if TYPE_CHECKING:
    from megatron.core.dist_checkpointing.mapping import ShardedStateDict


class PEFT(IOMixin, ABC, ModelTransform):
    """Abstract base class for Parameter-Efficient Fine-Tuning (PEFT) methods.

    This class defines the interface for PEFT methods, which are used to fine-tune
    large language models efficiently by modifying only a small subset of the model's
    parameters.

    Example:
        class MyPEFT(PEFT):
            def transform(self, module, name=None, prefix=None):
                # Implement the transform logic
                pass


        peft = MyPEFT()
        peft_model = LargeLanguageModel(model_transform=peft)
    """

    @abstractmethod
    def transform(self, module, name=None, prefix=None):
        """Transform a single module according to the PEFT method.

        This method is called for each module in the model during the PEFT application process.
        It should be implemented by subclasses to define how individual modules are transformed
        for the specific PEFT technique.

        Args:
            module (nn.Module): The individual module to be transformed.
            name (Optional[str]): The name of the module within the model structure. Defaults to None.
            prefix (Optional[str]): A prefix to be added to the module name, typically used for
                                    nested modules. Defaults to None.

        Returns:
            nn.Module: The transformed module. This can be the original module with modifications,
                       a new module replacing the original, or the original module if no
                       transformation is needed for this specific module.

        Note:
            This method is automatically called for each module in the model when the PEFT
            instance is applied to the model using the __call__ method.
        """
        raise NotImplementedError("The transform method should be implemented by subclasses.")

    def __call__(self, model: nn.Module) -> nn.Module:
        """Apply the PEFT method to the entire model.

        This method freezes the model parameters and walks through the model
        structure, applying the transform method to each module.

        Args:
            model (nn.Module): The model to be fine-tuned.

        Returns:
            nn.Module: The transformed model with PEFT applied.
        """
        self.freeze_model(model)

        # walk model chunks
        if isinstance(model, MegatronParallel) and len(model) > 1:
            for model_chunk in model:
                model_chunk.walk(self.transform)
        elif isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel):
            model.module.walk(self.transform)
        else:
            model.walk(self.transform)

        if is_trainer_attached(model) and model.trainer.state.fn != TrainerFn.FITTING:
            self.freeze_model(model)
        return model

    def freeze_model(self, model: nn.Module) -> None:
        """Apply a default freeze method to the model.

        This method freezes all the model parameters. This method can be overridden by subclasses to
        implement custom freeze strategies (e.g. freeze only parts of the model)

        Args:
            model (nn.Module): The model to be fine-tuned.

        Returns:
            nn.Module: The transformed model with PEFT applied.
        """
        if isinstance(model, MegatronParallel) and len(model) > 1:
            for model_chunk in model:
                model_chunk.freeze()
        if isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel):
            model.module.freeze()
        else:
            model.freeze()
        if is_trainer_attached(model) and model.trainer.state.fn == TrainerFn.FITTING:
            model.train(mode=True)

    def get_wrappped_io(self):
        """
        This is a helper function to return a partial function that wraps the checkpoint I/O with the PEFT adapter.
        Can be overridden in each PEFT method class.
        """
        return partial(WrappedAdapterIO, peft=self)

    def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) -> None:
        """PTL callback setup function."""
        from nemo.lightning.pytorch.strategies.utils import create_checkpoint_io
        from nemo.lightning.pytorch.utils import get_automodel_from_trainer

        super().setup(trainer, pl_module, stage=stage)

        self._add_via_setattr = 'HFAutoModel' in type(pl_module).__name__
        trainer.strategy.trainer = trainer
        wrapped_io = self.get_wrappped_io()

        # automodel_setup_optimizers is either None or holds a reference to trainer.strategy.setup_optimizers
        self.automodel_setup_optimizers = None
        # automodel adds adapters in configure_model
        self.transform_already_applied = False
        if get_automodel_from_trainer(trainer) is not None:
            ckpt_io_kwargs = {"model_library": "huggingface", "lora": True}
            # Due to the workaround used in peft restoration, it makes restoration non-PTL conforming,
            # therefore need to short-circuit these two functions.
            trainer._checkpoint_connector.restore_training_state = lambda: True
            trainer._checkpoint_connector.restore_model = lambda: True
            self.automodel_setup_optimizers = trainer.strategy.setup_optimizers
            self.transform_already_applied = True
            trainer.strategy.setup_optimizers = lambda x: True
        else:
            ckpt_io_kwarg_names = [
                "save_ckpt_format",
                "async_save",
                "torch_dist_multiproc",
                "assume_constant_structure",
                "parallel_save",
                "parallel_save_within_dp",
                "parallel_load",
                "load_directly_on_device",
            ]
            ckpt_io_kwargs = {
                arg: getattr(trainer.strategy, arg)
                for arg in filter(lambda x: hasattr(trainer.strategy, x), ckpt_io_kwarg_names)
            }
        trainer.strategy._checkpoint_io = create_checkpoint_io(wrapping_ckpt_io=wrapped_io, **ckpt_io_kwargs)
        self.wrapped_io = (
            trainer.strategy._checkpoint_io._checkpoint_io
            if getattr(trainer.strategy, 'async_save', False)
            else trainer.strategy._checkpoint_io
        )
        trainer.strategy._init_model_parallel = False
        # it will enter the following if statement if the model is on the automodel workflow
        # where the PEFT is applied in the configure_model
        if self.transform_already_applied:
            trainer.strategy._init_model_parallel = True
        trainer.strategy._setup_optimizers = False

    def set_params_to_save(self, trainer: pl.Trainer) -> None:
        """
        Set params to be saved for PEFT. This function is called in apply_transform.
        Can be overridden in each PEFT method class.
        """
        self.params_to_save = set(
            name for name, param in trainer.lightning_module.named_parameters() if param.requires_grad
        )
        for module_name, module in trainer.lightning_module.named_modules():
            if hasattr(module, "track_running_stats"):
                for buffer_name, buffer in module.named_buffers():
                    if buffer is not None:
                        self.params_to_save.add(module_name + "." + buffer_name)

    def apply_transform(self, trainer):
        """
        This function does the following:
        1. Apply PEFT model transform.
        2. Set up model parallel and optimizer, which were skipped in setup
        3. Load weights and optimizer state dict
        4. Set up `finalize_model_grads` from mcore.
        """
        # automodel adds adapters in configure_model
        if not getattr(self, 'transform_already_applied', False):
            super().apply_transform(trainer)
        self.set_params_to_save(trainer)

        # Handle automodel and return early.
        if (
            self.wrapped_io.adapter_ckpt_path is not None
            and Path(self.wrapped_io.adapter_ckpt_path).parts[-1] == HF_ADAPTER_PATH
        ):
            # Automodel adapter restoration is handled in restore_automodel.
            return self.restore_automodel(trainer, self.wrapped_io.adapter_ckpt_path.parent)
        elif getattr(self, 'transform_already_applied', False) == True or self.automodel_setup_optimizers is not None:
            if self.automodel_setup_optimizers is not None:
                logging.info("Setting up optimizers")
                self.automodel_setup_optimizers(trainer)
                self.automodel_setup_optimizers = None
            return

        adapter_sharded_state_dict = {}
        if self.wrapped_io.adapter_ckpt_path is not None:
            logging.info(f"Loading adapters from {self.wrapped_io.adapter_ckpt_path}")
            # create sharded state dict for adapter weights only to enable PEFT resume
            adapter_sharded_state_dict['state_dict'] = {
                k: v for k, v in trainer.model.sharded_state_dict().items() if self.adapter_key_filter(k)
            }

        if hasattr(trainer.strategy, "init_model_parallel"):
            logging.info("Initializing model parallel")
            trainer.strategy.init_model_parallel()

        if trainer.state.fn == TrainerFn.FITTING:
            logging.info("Setting up optimizers")
            trainer.strategy.setup_optimizers(trainer)
            if self.wrapped_io.adapter_ckpt_path is not None and trainer.strategy.should_restore_optimizer_states():
                # PEFT resume, load optimizer state
                adapter_sharded_state_dict['optimizer'] = [
                    trainer.strategy.optimizer_sharded_state_dict(is_loading=True)
                ]

        if adapter_sharded_state_dict:
            adapter_state = self.wrapped_io.load_checkpoint(
                self.wrapped_io.adapter_ckpt_path, sharded_state_dict=adapter_sharded_state_dict
            )
            trainer.strategy.load_model_state_dict(adapter_state, strict=False)
            if trainer.state.fn == TrainerFn.FITTING:
                # Load optimizer
                trainer.strategy.load_optimizer_state_dict(adapter_state, selective_restore=False)
                # Load lr scheduler
                if (lr_schedulers := adapter_state.get('lr_schedulers', None)) is not None:
                    for config, lrs_state in zip(trainer.lr_scheduler_configs, lr_schedulers):
                        config.scheduler.load_state_dict(lrs_state)

        for cb in trainer.callbacks[::-1]:
            if isinstance(cb, MegatronOptimizerModule):
                cb.on_fit_start(trainer, trainer.lightning_module)
                break
        else:
            # i.e., this is an mcore model; elif not supported here.
            if get_automodel_from_trainer(trainer) is None:
                logging.warning(
                    "MegatronOptimizerModule not found in trainer callbacks. finalize_model_grads is not "
                    "properly set up for PEFT."
                )

    def restore_automodel(self, trainer, path):
        """restores automodel's adapter and optimizer state dict"""

        def pop_fqn_prefix(fqn, prefix='model'):
            """helper function to remove first "model" from fqn"""
            parts = fqn.split('.')
            assert parts[0] == prefix
            return '.'.join(parts[1:])

        adapter_state = self.wrapped_io.load_checkpoint(path)
        # Ensure all keys from adapter_state are contained in model
        state_dict = trainer.lightning_module.state_dict()
        for key in adapter_state['state_dict'].keys():
            assert key in state_dict, (key, state_dict.keys())

        # Move to cpu and load state
        from nemo.lightning.pytorch.strategies.utils import to_cpu

        trainer.strategy.load_model_state_dict(
            {'state_dict': {pop_fqn_prefix(k): to_cpu(v) for k, v in adapter_state['state_dict'].items()}},
            strict=False,
        )

        # Ensure adapters have grad enabled
        for key, param in trainer.lightning_module.named_parameters():
            param.requires_grad_(key in adapter_state['state_dict'])

        if trainer.state.fn == TrainerFn.FITTING:
            # Restore optim and LR Scheduler
            assert self.automodel_setup_optimizers is not None, "Expected automodel_setup_optimizers to be valid"
            self.automodel_setup_optimizers(trainer)
            # Load optimizer
            trainer.strategy.load_optimizer_state_dict(adapter_state)
            # Load lr scheduler
            if (lr_schedulers := adapter_state.get('lr_schedulers', None)) is not None:
                for config, lrs_state in zip(trainer.lr_scheduler_configs, lr_schedulers):
                    config.scheduler.load_state_dict(lrs_state)

    def adapter_key_filter(self, key: str) -> bool:
        """
        Given a key in the state dict, return whether the key is an adapter (or base model).
        This function can be subclassed in each PEFT method class.
        """
        if isinstance(key, tuple):
            return key[1].requires_grad
        return key in self.params_to_save or ".adapter." in key or key.endswith(".adapters")


class AdapterWrapper(nn.Module):
    """Abstract base class for wrapping modules with adapters in Parameter-Efficient Fine-Tuning (PEFT).

    This class wraps a module and its associated adapter, providing methods for
    managing the state dictionaries of both the main module and the adapter. It does not
    implement the forward method, which must be implemented by concrete subclasses.

    Attributes:
        to_wrap (nn.Module): The main module to be wrapped.
        adapter (nn.Module): The adapter module to be applied.

    Note:
        This class is abstract and cannot be instantiated directly. Subclasses must
        implement the forward method.

    Example:
        class LoRALinear(AdapterWrapper):
            def __init__(self, to_wrap, adapter):
                super().__init__(to_wrap, adapter)

            def forward(self, x):
                return self.to_wrap(x) + self.adapter(x)

        main_module = nn.Linear(100, 100)
        adapter = nn.Linear(100, 100)
        parallel_adapter = LoRALinear(main_module, adapter)
    """

    def __init__(self, to_wrap: nn.Module, adapter: nn.Module):
        super(AdapterWrapper, self).__init__()
        self.to_wrap = to_wrap
        self.adapter = adapter

    def base_linear_forward(self, x, *args, **kwargs):
        """
        Run the forward method of the linear module `to_wrap`.
        Return a tuple of three elements: linear_output, bias, layernorm_output

        x -> [layernorm/identity] -> layernorm_output -> [linear] -> linear_output, bias

        layernorm_output is different from input x only when linear layer is LayerNormColumnParallelLinear.
        """
        linear_output = self.to_wrap(x, *args, **kwargs)
        assert isinstance(
            linear_output, tuple
        ), f"{self.to_wrap} should return a tuple but instead returns {linear_output}"
        """ Four cases for the wrapped module's return values
        1. nothing: (out, None)
        2. return_bias: (out, bias)
        2. return_layernorm_output: ((out, ln_out), None)
        3. both: (out, bias, ln_out)
        """
        bias = None
        layernorm_output = x
        if len(linear_output) == 2:
            linear_output, bias = linear_output
            if isinstance(linear_output, tuple) and len(linear_output) == 2:
                linear_output, layernorm_output = linear_output
        elif len(linear_output) == 3:
            linear_output, bias, layernorm_output = linear_output

        return linear_output, bias, layernorm_output

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        """Retrieve the state dictionary of the wrapped module and adapter.

        This method overrides the default state_dict behavior to include both
        the main module's state and the adapter's state under a special 'adapters' key.

        Args:
            destination (Optional[dict]): A dictionary to store the state. If None, a new
                                          dictionary is created. Defaults to None.
            prefix (str): A prefix added to parameter and buffer names. Defaults to ''.
            keep_vars (bool): If True, returns variables instead of tensor values.
                              Defaults to False.

        Returns:
            dict: The state dictionary containing both the main module and adapter states.
        """

        if destination is None:
            destination = {}

        # Get state dict of the main module
        self.to_wrap.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)

        # Store adapter state dict under the "adapter" prefix in the destination dict
        self.adapter.state_dict(destination=destination, prefix=f'{prefix}adapter.', keep_vars=keep_vars)
        return destination

    def sharded_state_dict(
        self,
        prefix: str = '',
        sharded_offsets: Tuple[Tuple[int, int, int]] = (),
        metadata: Optional[dict] = None,
    ) -> "ShardedStateDict":
        """Retrieve the sharded state dictionary of the wrapped module and adapter.

        This method is used for distributed checkpointing, combining the sharded states
        of both the main module and the adapter.

        Args:
            prefix (str): A prefix added to parameter and buffer names. Defaults to ''.
            sharded_offsets (Tuple[Tuple[int, int, int]]): Offsets for sharded parameters.
                                                           Defaults to an empty tuple.
            metadata (Optional[dict]): Additional metadata for the sharded state.
                                       Defaults to None.

        Returns:
            ShardedStateDict: The combined sharded state dictionary.
        """
        sharded_state_dict = {}
        sharded_state_dict.update(self.to_wrap.sharded_state_dict(prefix, sharded_offsets, metadata))
        sharded_state_dict.update(self.adapter.sharded_state_dict(f"{prefix}adapter.", sharded_offsets, metadata))
        return sharded_state_dict


class WrappedAdapterIO(_WrappingCheckpointIO, AsyncCompatibleCheckpointIO):  # noqa: F821
    """
    A wrapper class for checkpoint I/O operations, specifically designed for PEFT (Parameter-Efficient Fine-Tuning).

    This class handles the complexities of saving and loading checkpoints for both initial PEFT training and resuming
    PEFT training. It ensures that only the necessary adapter weights are saved and loaded, while also preserving the
    base model weights.

    **Usage:**

    1. **Initial PEFT Training:**
       - The class handles the saving of only adapter weights.
       - Metadata about the base model checkpoint is stored for future reference.

    2. **PEFT Resume:**
       - The class loads both base model and adapter weights.
       - The previously stored metadata is used to locate the correct base model checkpoint.

    **Attributes:**

    - `peft`: The PEFT instance associated with the wrapped checkpoint I/O.
    - `model_ckpt_path`: The path to the base model checkpoint.
    - `adapter_ckpt_path`: The path to the adapter checkpoint.
    Note that the paths are set by save/load functions and users do not need to set them.

    **Methods:**

    - `save_checkpoint`: Saves the adapter weights and metadata to the specified path.
    - `load_checkpoint`: Loads the base model and adapter weights based on the specified path and metadata.
    """

    peft: Optional[PEFT] = None
    model_ckpt_path: Optional[Path] = None
    adapter_ckpt_path: Optional[Path] = None

    def __init__(
        self, checkpoint_io: Optional["CheckpointIO"] = None, peft: Optional[PEFT] = None  # noqa: F821
    ) -> None:
        self.peft = peft
        super().__init__(checkpoint_io)

    @override
    def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
        assert self.checkpoint_io is not None
        state_key = None
        for k in ['sharded_state_dict', 'state_dict']:
            if k in checkpoint:
                state_key = k
                break
        assert state_key is not None, "Expected checkpoint to contain `sharded_state_dict` or `state_dict`"
        assert state_key in checkpoint, "Expected state_key to be in checkpoint"

        state_dict = checkpoint.pop(state_key)
        checkpoint[state_key] = dict(filter(lambda item: self.peft.adapter_key_filter(item[0]), state_dict.items()))
        ckpt_keys = list(checkpoint[state_key].keys())
        request = self.checkpoint_io.save_checkpoint(checkpoint, path, storage_options=storage_options)

        from nemo.utils.get_rank import is_global_rank_zero

        if is_global_rank_zero():
            base_dir = ckpt_to_weights_subdir(path, is_saving=True)

            from nemo.lightning.io.hf import HFCheckpointIO

            if isinstance(self.checkpoint_io, HFCheckpointIO):
                metadata = self._create_lora_hf_config(ckpt_keys)
                hf_adapter_base = base_dir.parent / HF_ADAPTER_PATH
                hf_adapter_base.mkdir(parents=True, exist_ok=True)
                adapter_meta_path = hf_adapter_base / HF_ADAPTER_CONFIG_FILENAME
            else:
                base_dir.mkdir(parents=True, exist_ok=True)
                metadata = {"model_ckpt_path": str(self.model_ckpt_path)}
                adapter_meta_path = base_dir / ADAPTER_META_FILENAME

            with open(adapter_meta_path, "w") as f:
                json.dump(metadata, f)
        return request

    def _create_lora_hf_config(self, ckpt_keys):
        """Creates a HF lora config from a NeMo Lora config"""

        def extract_matched_module_names(ckpt_keys, target_modules):
            """
            Extracts module names from a list of checkpoint keys that match the target modules.

            This function processes a list of target module patterns, where each pattern may or may
            not contain a wildcard (`'*'`). The function matches these patterns against the
            checkpoint keys, with the following behavior:
            - Patterns containing '*' will be expanded to match any sequence of characters
              except a dot (`.`).
            - Patterns without '*' are matched literally.

            Args:
                ckpt_keys (list of str): A list of strings representing checkpoint keys to be
                    searched.
                target_modules (list of str): A list of target module patterns. Some patterns may
                    contain wildcards (`'*'`), which match any characters except a dot.

            Returns:
                list of str: A list of module names from `target_modules` that match any of the
                `ckpt_keys`. The result is returned as a list of unique module names.

            Example:
                ckpt_keys = [
                    "model.model.layers.27.self_attn.k_proj",
                    "model.model.layers.27.self_attn.v_proj",
                    "model.model.layers.27.self_attn.mlp"
                ]
                target_modules = ["*proj"]

                extract_matched_module_names(ckpt_keys, target_modules)
                # Output: ['k_proj', 'v_proj']

            Notes:
                - This function uses regular expressions to match the target patterns in the
                  checkpoint keys.
                - Wildcards are expanded as `[^.]+` to ensure that the match doesn't cross dot
                  (`.`) boundaries.
            """
            re_target_modules = list(filter(lambda x: '*' in x, target_modules))
            if len(re_target_modules) == 0:
                return target_modules
            non_re_target_modules = list(filter(lambda x: not '*' in x, target_modules))
            combined_pattern = '|'.join(
                map(lambda x: x.replace('*', '[^.]+'), re_target_modules),
            )
            ans = set(non_re_target_modules)
            for key in ckpt_keys:
                ans.update(re.findall(combined_pattern, key))
            return list(ans)

        from peft import LoraConfig

        from nemo.collections.llm.peft import DoRA

        # Contains all target module names, without any regular expression
        materialized_module_names = extract_matched_module_names(ckpt_keys, self.peft.target_modules)
        lora_config = LoraConfig(
            r=self.peft.dim,
            target_modules=materialized_module_names,
            lora_alpha=self.peft.alpha,
            lora_dropout=self.peft.dropout,
            use_dora=isinstance(self.peft, DoRA),
        )
        lora_config = lora_config.to_dict()
        lora_config["peft_type"] = "LORA"
        lora_config["megatron_core"] = None
        lora_config["target_modules"] = materialized_module_names
        return lora_config

    @override
    def load_checkpoint(
        self,
        path: _PATH,
        sharded_state_dict=None,
        map_location: Optional[Callable] = None,
        strict: Optional['StrictHandling'] | bool = None,  # noqa: F821
    ) -> Dict[str, Any]:
        """
        =====================
        Initial PEFT Training
        =====================
        Initial PEFT training requires loading the base model weights. In this case, this function is called by
        trainer.strategy.setup() -> megatron_strategy.restore_model() -> megatron_strategy.load_checkpoint().
        `path = PosixPath(<base_path>)`, and sharded_state_dict contains only base model weights

        ===========
        PEFT Resume
        ===========
        PEFT resume requires loading two set of model weights, 1) base model weights and 2) adapter weights
        Base model weights could be imported from e.g. HF, and is frozen during PEFT training.
        Adapter weights contains the training metadata that will need to be loaded.
        As such, this function will be entered twice during PEFT training resume.

        For the FIRST TIME this function is called by trainer._checkpoint_connector._restore_modules_and_callbacks.
        `path = AdapterPath(<adapter_path>, base_model_path=<base_path>)`, and sharded_state_dict contains only base
        model weights

        For the SECOND TIME this function is called by PEFT.apply_transform (above, in the same file).
        `path = PosixPath(<adapter_path>)`, and sharded_state_dict contains only adapter weights.
        """

        assert self.checkpoint_io is not None

        adapter_ckpt = None
        base = ckpt_to_dir(path)
        if getattr(path, "base_model_path", None):
            # PEFT Resume, FIRST TIME
            self.adapter_ckpt_path = Path(str(path))
            adapter_ckpt = self.checkpoint_io.load_checkpoint(path, sharded_state_dict={})  # Loads only metadata
            # path is adapter path to restore the training metadata, but switch to loading base model here.
            path = self.model_ckpt_path = path.base_model_path
        elif (adapter_meta_path := base / ADAPTER_META_FILENAME).exists():
            # PEFT Resume, SECOND TIME
            with open(adapter_meta_path, "r") as f:
                metadata = json.load(f)
            self.model_ckpt_path = Path(metadata['model_ckpt_path'])
            self.adapter_ckpt_path = path
        elif (base / HF_ADAPTER_PATH / HF_ADAPTER_CONFIG_FILENAME).exists():
            self.adapter_ckpt_path = path / HF_ADAPTER_PATH
        else:
            # Initial PEFT Training
            self.model_ckpt_path = path

        # Note: this will include the Trainer-state of the model-checkpoint
        model_ckpt = self.checkpoint_io.load_checkpoint(path, sharded_state_dict, map_location, strict)
        if adapter_ckpt is not None:
            # PEFT Resume, FIRST TIME
            adapter_ckpt['state_dict'].update(model_ckpt['state_dict'])
            return adapter_ckpt
        return model_ckpt