chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,220 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user