deepspeedai--deepspeed/deepspeed/inference/v2/model_implementations/inference_policy_base.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import json
from abc import ABC, ABCMeta, abstractmethod
from typing import Any, Iterable, List, Optional, Union

import torch

from ..config_v2 import RaggedInferenceEngineConfig
from ..checkpoint import CheckpointEngineBase
from ..logging import inference_logger
from .layer_container_base import LayerContainer
from .inference_model_base import DSInferenceModelBase
from .flat_model_helpers import (
    flatten_inference_model,
    make_param_filename,
    make_metadata_filename,
    ModelMetadata,
    restore_inference_model,
)

POLICIES = {}


class ContainerMap:

    def __init__(self) -> None:
        self._prefix_map = {}
        self._transformer_params = None
        self._non_transformer_params = None

    @property
    def transformer_params(self) -> Iterable[LayerContainer]:
        return self._transformer_params

    @property
    def non_transformer_params(self) -> LayerContainer:
        return self._non_transformer_params

    def set_transformer_params(self, prefixes: Union[str, Iterable[str]], containers: List[LayerContainer]) -> None:
        if not isinstance(containers, list):
            raise ValueError(
                f"The transformer containers should be a list, of one container per layer, but got {type(containers)} instead."
            )

        self._transformer_prefixes = prefixes if isinstance(prefixes, list) else [prefixes]
        self._transformer_params = containers

    def set_non_transformer_params(self, container: LayerContainer) -> None:
        self._non_transformer_params = container

    def set_unmapped_params(self, prefixes: Union[str, Iterable[str]]) -> None:
        self._unmapped_prefixes = prefixes

    def map_param(self, name, parameter) -> None:
        for unmapped_prefix in self._unmapped_prefixes:
            if name.startswith(unmapped_prefix):
                inference_logger().debug(f"Ignoring: {name} for {unmapped_prefix}")
                return

        for transformer_prefix in self._transformer_prefixes:
            if name.startswith(transformer_prefix):
                popped_name = name[len(transformer_prefix) + 1:]
                layer_idx = popped_name.split(".")[0]
                assert layer_idx.isdigit(
                ), f"expected name to start w. list index but got {layer_idx} instead, name={name}"
                layer_idx = int(layer_idx)
                inference_logger().debug(
                    f"Setting: {'.'.join(popped_name.split('.')[1:])} for layer-idx={layer_idx} to {parameter.shape}")
                self._transformer_params[layer_idx].set_dependency(".".join(popped_name.split(".")[1:]), parameter)
                return

        try:
            inference_logger().debug(f"Setting: {name} to {parameter.shape}")
            self._non_transformer_params.set_dependency(name, parameter)
        except ValueError:
            # Catch the ValueError here from the non_transformer_params because we are knowingly
            # calling it with something that may not match. This should allow us to raise a slightly more
            # informative error message.
            raise ValueError(f"Cannot find container for {name}, please double check the Containers/ContainerMap")

    def validate(self) -> None:
        if not self._non_transformer_params.is_initialized:
            raise RuntimeError("Non-transformer parameters not fully initialized after checkpoint load.")

        for layer_idx, container in enumerate(self._transformer_params):
            if not container.is_initialized:
                raise RuntimeError(
                    f"Transformer container at index {layer_idx} not fully initialized after checkpoint load.")


class PolicyMeta(ABCMeta):

    def __new__(cls, name, bases, dct):
        new_obj = super().__new__(cls, name, bases, dct)
        if name != "InferenceV2Policy":
            POLICIES[name] = new_obj
        return new_obj


class InferenceV2Policy(ABC, metaclass=PolicyMeta):
    """
    The InferenceV2Policy is the base class for all inference policies. An inference policy
    is responsible for instantiating the inference model and mapping the parameters from the
    checkpoint engine to the model itself.
    """

    def __init__(
        self,
        model_config: Any,
        checkpoint_engine: Optional[CheckpointEngineBase] = None,
        inf_checkpoint_path: Optional[str] = None,
    ) -> None:
        """
        Create the Policy with sufficient context to build the model. There are two supported
        model creation mechanisms.

        The first is the generalized ``checkpoint_engine`` which
        will iterate over the parameters of the model and provide them to the policy. These in
        turn will be sharded/transformed by the model implementation.

        The second is used to re-create a previously serialized DeepSpeed inference model. These
        checkpoints should not be used across different model backend configurations.

        TODO(cmikeh2): Enforce this in code
        """
        if checkpoint_engine is None and inf_checkpoint_path is None:
            raise ValueError("Either checkpoint_engine or ds_checkpoint_path must be provided.")

        if checkpoint_engine is not None and inf_checkpoint_path is not None:
            raise ValueError("Only one of checkpoint_engine or ds_checkpoint_path can be provided.")

        self._checkpoint_engine = checkpoint_engine
        self._inf_checkpoint_path = inf_checkpoint_path
        self._model_config = model_config

    def build_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> DSInferenceModelBase:
        """
        Completely instantiate the inference model. This will both create the ops needed to run the
        model, as well as load the model parameters via the checkpoint engine. For more context
        on each of these components please see ``instantiate_model`` and ``populate_model_parameters``.

        Arguments:
            engine_config: The config that has been used to instantiate the engine. This is used
                to communicate to the model implementation the limits on batches (sequences/tokens)
                and bound the size of intermediate buffers.
            mp_group: Object to enable communication between tensor parallel ranks.

        Returns:
            DSInferenceModelBase: An implementation of the inference model abstraction that will be
                run by the engine.
        """
        self.model = self.instantiate_model(engine_config, mp_group)
        self.populate_model_parameters()
        return self.model

    @abstractmethod
    def instantiate_model(self, engine_config: RaggedInferenceEngineConfig) -> DSInferenceModelBase:
        """
        Instantiate the inference model. Depending on the engine/model config, this could be where
        different model implementations could be selected.

        Arguments:
            engine_config: The config that has been used to instantiate the engine. This is used
                to communicate to the model implementation the limits on batches (sequences/tokens)
                and bound the size of intermediate buffers.

        Returns:
            DSInferenceModelBase: An implementation of the inference model abstraction that will be
                run by the engine.
        """
        ...

    @abstractmethod
    def build_container_map(self) -> ContainerMap:
        """
        Build a dictionary representing the structure of the string prefixes leading
        to the parameters to be mapped to the container.

        Returns:
            ContainerMap: An instantiated mapping describing how checkpoint prefixes map
                to ``LayerContainer`` instances.
        """
        raise NotImplementedError()

    def populate_model_parameters(self) -> None:
        """
        This model will iterate over the parameters (as provided by the checkpoint engine) and
        use the container map built by ``build_container_map`` to populate the model
        """

        container_map = self.build_container_map()

        if self._checkpoint_engine is not None:
            for name, parameter in self._checkpoint_engine.parameters():
                container_map.map_param(name, parameter)

            buffer, metadata = flatten_inference_model(container_map.transformer_params,
                                                       container_map.non_transformer_params, self.__class__.__name__)
        else:

            buffer_path = make_param_filename(self._inf_checkpoint_path, self.model.tp_rank, self.model.tp_size)
            metadata_path = make_metadata_filename(self._inf_checkpoint_path, self.model.tp_rank, self.model.tp_size)

            buffer = torch.load(buffer_path, weights_only=False)
            metadata = json.load(open(metadata_path, "r"))
            metadata = ModelMetadata.parse_raw(metadata)

            restore_inference_model(buffer, metadata, container_map.transformer_params,
                                    container_map.non_transformer_params)

        container_map.validate()

        self.model.set_parameters(transformer=container_map.transformer_params,
                                  non_transformer=container_map.non_transformer_params,
                                  flattened_param_buffer=buffer,
                                  flattened_param_metadata=metadata)