wehub-resource-sync
283 lines
9.5 KiB
Python
283 lines
9.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Dict, Iterable, Tuple, Optional
|
|
from os import path
|
|
|
|
import torch
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.ops.op_builder import RaggedUtilsBuilder
|
|
from deepspeed.runtime.config_utils import DeepSpeedConfigModel
|
|
from .layer_container_base import LayerContainer
|
|
from ..inference_parameter import InferenceParameter, STR_TO_DTYPE
|
|
from ..inference_utils import elem_size
|
|
|
|
|
|
def pad_to_aligned_offset(offset: int, alignment: int = 256) -> int:
|
|
"""
|
|
Pad the provided offset to a well-aligned value.
|
|
"""
|
|
return ((offset + alignment - 1) // alignment) * alignment
|
|
|
|
|
|
class TensorMetadata(DeepSpeedConfigModel):
|
|
"""
|
|
A class to represent a tensor specification.
|
|
"""
|
|
dtype: Optional[str] = None
|
|
shape: Optional[Tuple[int, ...]] = None
|
|
strides: Optional[Tuple[int, ...]] = None
|
|
offset: int
|
|
|
|
|
|
class ParameterMetadata(DeepSpeedConfigModel):
|
|
"""
|
|
A class to represent a parameter specification.
|
|
"""
|
|
core_param: Optional[TensorMetadata] = None
|
|
aux_params: Dict[str, TensorMetadata] = {}
|
|
|
|
|
|
class LayerMetadata(DeepSpeedConfigModel):
|
|
"""
|
|
A class to represent a layer specification.
|
|
"""
|
|
params: Dict[str, ParameterMetadata] = {}
|
|
|
|
|
|
class ModelMetadata(DeepSpeedConfigModel):
|
|
"""
|
|
A class to represent a model specification.
|
|
"""
|
|
policy: str = ""
|
|
layers: Dict[str, LayerMetadata] = {}
|
|
|
|
|
|
def make_param_filename(base: str, rank: int, n_ranks: int) -> str:
|
|
"""
|
|
Make a filename for a parameter file.
|
|
|
|
Arguments:
|
|
rank: Rank of the file.
|
|
n_ranks: Total number of ranks.
|
|
|
|
Returns:
|
|
str: Filename.
|
|
"""
|
|
return path.join(base, f"params_rank_{rank}_of_{n_ranks}.pt")
|
|
|
|
|
|
def make_metadata_filename(base: str, rank: int, n_ranks: int) -> str:
|
|
"""
|
|
Make a filename for a metadata file.
|
|
|
|
Arguments:
|
|
rank: Rank of the file.
|
|
n_ranks: Total number of ranks.
|
|
|
|
Returns:
|
|
str: Filename.
|
|
"""
|
|
return path.join(base, f"metadata_rank_{rank}_of_{n_ranks}.json")
|
|
|
|
|
|
def make_model_config_filename(base: str) -> str:
|
|
"""
|
|
Make a filename for a model config file.
|
|
|
|
Arguments:
|
|
base: Base directory.
|
|
|
|
Returns:
|
|
str: Filename.
|
|
"""
|
|
return path.join(base, "ds_model_config.json")
|
|
|
|
|
|
def flatten_inference_model(
|
|
transformer_containers: Iterable[LayerContainer],
|
|
non_transformer_container: LayerContainer,
|
|
policy_name: str,
|
|
) -> Tuple[torch.Tensor, ModelMetadata]:
|
|
"""
|
|
Flatten the underlying parameters into
|
|
|
|
Arguments:
|
|
transformer_containers: Iterable of layer containers corresponding to the transformer
|
|
parameters.
|
|
non_transformer_container: Layer container corresponding to the non-transformer parameters.
|
|
policy_name: The name of the policy class (typically accessed with `type(policy).__name__`).
|
|
|
|
Returns:
|
|
Iterable[Any]: Flattened list of parameters.
|
|
"""
|
|
alloc_fn = RaggedUtilsBuilder().load().allocate_view_on
|
|
|
|
total_size = 0
|
|
metadata = ModelMetadata(policy=policy_name)
|
|
|
|
def process_layer(layer_container: LayerContainer, l_name: str, cur_offset: int) -> int:
|
|
"""
|
|
Iterate over the parameters of a single container and collect metadata for the final
|
|
flattened buffer.
|
|
|
|
Arguments:
|
|
layer_container: The layer container to process.
|
|
l_name: The name of the layer container to key the metadata.
|
|
cur_offset: The current offset into the flattened buffer.
|
|
|
|
Captured Variables:
|
|
metadata: The metadata object to populate.
|
|
|
|
Returns:
|
|
int: The updated offset into the flattened buffer.
|
|
"""
|
|
try:
|
|
_ = layer_container.is_populated
|
|
except ValueError as e:
|
|
raise ValueError(f"Layer container {l_name} is not populated.") from e
|
|
|
|
layer_metadata = LayerMetadata()
|
|
|
|
for p_name in layer_container.annotation_attrs:
|
|
param = getattr(layer_container, p_name)
|
|
param_metadata = ParameterMetadata()
|
|
|
|
if param is None:
|
|
param_metadata.core_param = TensorMetadata(offset=-1)
|
|
layer_metadata.params[p_name] = param_metadata
|
|
continue
|
|
|
|
param_metadata.core_param = TensorMetadata(dtype=str(param.dtype),
|
|
shape=param.shape,
|
|
strides=param.stride(),
|
|
offset=cur_offset)
|
|
|
|
cur_offset += pad_to_aligned_offset(elem_size(param.dtype) * param.numel())
|
|
|
|
for t_name, tensor in param.aux_attrs.items():
|
|
param_metadata.aux_params[t_name] = TensorMetadata(dtype=str(tensor.dtype),
|
|
shape=tensor.shape,
|
|
strides=tensor.stride(),
|
|
offset=cur_offset)
|
|
|
|
cur_offset += pad_to_aligned_offset(elem_size(tensor.dtype) * tensor.numel())
|
|
|
|
layer_metadata.params[p_name] = param_metadata
|
|
|
|
metadata.layers[l_name] = layer_metadata
|
|
return cur_offset
|
|
|
|
for i, layer in enumerate(transformer_containers):
|
|
l_name = f"transformer_layer_{i}"
|
|
total_size = process_layer(layer, l_name, total_size)
|
|
|
|
l_name = "non_transformer"
|
|
total_size = process_layer(non_transformer_container, l_name, total_size)
|
|
|
|
buffer = torch.empty(total_size, dtype=torch.uint8, device=get_accelerator().current_device())
|
|
|
|
def copy_layer(layer_container: LayerContainer, l_name: str) -> None:
|
|
"""
|
|
Local method for copying from the layer container to the flattened buffer.
|
|
|
|
Arguments:
|
|
layer_container: The layer container to copy from.
|
|
l_name: The name of the layer container to key the metadata.
|
|
|
|
Captured Variables:
|
|
buffer: The flattened buffer to copy into.
|
|
metadata: The metadata object to populate.
|
|
"""
|
|
l_metadata = metadata.layers[l_name]
|
|
for p_name in layer_container.annotation_attrs:
|
|
p_metadata = l_metadata.params[p_name]
|
|
param = getattr(layer_container, p_name)
|
|
|
|
if param is None:
|
|
continue
|
|
|
|
core_param = alloc_fn(param, buffer, p_metadata.core_param.offset)
|
|
core_param.copy_(param)
|
|
|
|
aux_params = {}
|
|
|
|
for t_name, tensor in param.aux_attrs.items():
|
|
t_view = alloc_fn(tensor, buffer, p_metadata.aux_params[t_name].offset)
|
|
aux_params[t_name] = t_view
|
|
t_view.copy_(tensor)
|
|
|
|
setattr(layer_container, p_name, InferenceParameter.initialize(core_param, **aux_params))
|
|
|
|
for i, layer in enumerate(transformer_containers):
|
|
l_name = f"transformer_layer_{i}"
|
|
copy_layer(layer, l_name)
|
|
|
|
l_name = "non_transformer"
|
|
copy_layer(non_transformer_container, l_name)
|
|
|
|
return buffer, metadata
|
|
|
|
|
|
def restore_inference_model(buffer: torch.Tensor, metadata: ModelMetadata,
|
|
transformer_containers: Iterable[LayerContainer],
|
|
non_transformer_container: LayerContainer) -> None:
|
|
"""
|
|
Restore the model from the buffer and metadata.
|
|
|
|
Arguments:
|
|
buffer: Buffer containing the model parameters.
|
|
metadata: Metadata for the model.
|
|
transformer_containers: Iterable of transformer layer containers.
|
|
non_transformer_container: Non-transformer layer container.
|
|
"""
|
|
alloc_fn = RaggedUtilsBuilder().load().allocate_view_like
|
|
|
|
def restore_layer(layer_container: LayerContainer, l_name: str) -> None:
|
|
"""
|
|
Local method for restoring a layer container from a flattened buffer. This
|
|
only constructs views for the parameters onto the buffer. No data movement
|
|
is performed.
|
|
|
|
Arguments:
|
|
layer_container: The layer container to restore.
|
|
l_name: The name of the layer container to key the metadata.
|
|
|
|
Captured Variables:
|
|
buffer: The flattened buffer to reconstruct views on top of.
|
|
metadata: The metadata object describing the each parameter in the model.
|
|
"""
|
|
l_metadata = metadata.layers[l_name]
|
|
|
|
for p_name in layer_container.annotation_attrs:
|
|
p_metadata = l_metadata.params[p_name]
|
|
|
|
if p_metadata.core_param.offset == -1:
|
|
layer_container.direct_injection(p_name, None)
|
|
continue
|
|
|
|
dummy_tensor = torch.empty([], dtype=STR_TO_DTYPE[p_metadata.core_param.dtype])
|
|
core_param = alloc_fn(p_metadata.core_param.shape, p_metadata.core_param.strides, dummy_tensor, buffer,
|
|
p_metadata.core_param.offset)
|
|
|
|
aux_params = {}
|
|
|
|
for t_name, t_metadata in p_metadata.aux_params.items():
|
|
dummy_tensor = torch.empty([], dtype=STR_TO_DTYPE[t_metadata.dtype])
|
|
t_view = alloc_fn(t_metadata.shape, t_metadata.strides, dummy_tensor, buffer, t_metadata.offset)
|
|
|
|
aux_params[t_name] = t_view
|
|
|
|
restored_param = InferenceParameter.initialize(core_param, **aux_params)
|
|
layer_container.direct_injection(p_name, restored_param)
|
|
|
|
for i, layer in enumerate(transformer_containers):
|
|
l_name = f"transformer_layer_{i}"
|
|
restore_layer(layer, l_name)
|
|
|
|
l_name = "non_transformer"
|
|
restore_layer(non_transformer_container, l_name)
|