chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,282 @@
+# 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)