chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,84 @@
|
||||
# Adding Support for a New Model in DeepSpeed Inference V2
|
||||
|
||||
Adding supoprt for a new model in DeepSpeed Inference requires developing three related components:
|
||||
- Containers: These describe the parameters contained in the model
|
||||
- Model implementation: How should the model be computed.
|
||||
- Policy: The map for adding parameters to your containers and creating the model implementation.
|
||||
|
||||
In this tutorial, we will assume that you'd like to use a relatively traditionally styled Transformer model and will be able to inherit from `DSTransformerModelBase` and can take advantage of the utilities that provides.
|
||||
|
||||
## Defining Your Containers
|
||||
|
||||
A container is the bridge between the original model's parameters and how to transform them to serve them for inference. For a model implementation, there are two primary kinds of containers: transformer containers and non-transformer containers. A transformer container consists of the parameters for a single Transformer layer in the model. So this includes your traditional parameters like the projections for the fully connected network, or query-key-value projections. The non-transformer container will contain basically everything else! However, before defining these containers, we need to understand how to define an individual parameter.
|
||||
|
||||
In DeepSpeed inference, the original model parameters are populated into the model and mapped as dependencies to a parameter. A `Parameter` has two primary components: its dependencies and its `finalize` method. Let's do an example. In Llama models, the native format is for the `query`, `key`, and `value` projections to be performed independently. However, we can achieve higher throughput by fusing them into a single larger projection. We can define this fusion with a parameter:
|
||||
|
||||
```python
|
||||
from deepspeed.inference.module_implementations.parameter_base import ParameterBase
|
||||
|
||||
class UnfusedQKVParameter(ParameterBase):
|
||||
query: torch.Tensor
|
||||
key: torch.Tensor
|
||||
value: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
fused_param = torch.cat([self.query, self.key, self.value], dim=0)
|
||||
return self.inference_model.transform_qkv_param(fused_param)
|
||||
```
|
||||
|
||||
Let's walk through each part of this implementation. First, parameters should inherit from `ParameterBase`. This will allow it to automatically determine when its dependencies are met and set the appropriate components of a parent `LayerContainer`. The second key component is the type annotations on the class itself. Each type annotation represents a dependency of the parameter. Since the original Llama mode has separate query, key, and value dependencies, our fused parameter will declare dependencies for each. Finally, we have the `finalize` method. This method is automatically called once all dependencies on the layer are met and should return the final parameter.
|
||||
|
||||
In this `finalize` method, we are doing two things: the first is the act of fusing the parameters together through the concatenate method. Note that each of the dependencies can be accessed via `self.{name}`. The second is calling `self.inference_model.transform_qkv_param`. A parameter's finalize method always has access to the inference model. In this case we are using that to use a feature provided by `DSTransformerBase`. This method will automatically shard the parameter for tensor parallelism and then pass it to the linear module implementation to perform additional optimizations or shape transformations, like quantization.
|
||||
|
||||
Since many patterns are very common in Transformer models, `model_implementations.common_parameters` provides implementations for many of the patterns (all compatible with `DSTransformerBase`) to help accelerate development.
|
||||
|
||||
Once all parameters are created, we need to compose them into a layer container. In our simplified Llama model, let's assume there's only QKV and attention output projection matrices. A layer container would appear as the following:
|
||||
|
||||
```python
|
||||
from deepspeed.inference.module_implementations.layer_container_base import LayerContainer
|
||||
|
||||
class ExampleContainer(LayerContainer):
|
||||
qkvw: UnfusedQKVParameter
|
||||
|
||||
attn_o: AttentionOutputParameter
|
||||
|
||||
PARAM_MAPPING: {
|
||||
"self_attn.q_proj.weight": "qkvw.query",
|
||||
"self_attn.k_proj.weight": "qkvw.key",
|
||||
"self_attn.v_proj.weight": "qkvw.value",
|
||||
"self_attn.o_proj.weight": "attn_o.params",
|
||||
}
|
||||
```
|
||||
|
||||
Once again, we have a couple of key components. The first are parameter type annotations. Each annotation corresponds to a parameter that can be used in the model implementation. In the model implementation, I can simply write `container.qkvw` to access my fused and transformed QKV parameter. The second key component is the `PARAM_MAPPING` dictionary. This is our explicit mapping of the names of parameters in the source model to a parameter dependency. This mapping dictionary will be used by the policy to automatically populate dependencies.
|
||||
|
||||
Once you have written `LayerContainer`s for both the transformer and non-transformer parameters, it's time to work on the model implementation!
|
||||
|
||||
## Building a Model Implementation that Inherits from `DSTransformerBase`
|
||||
|
||||
By inheriting from `DSTransformerBase`, most of the implementation work for sharding and transforming parameters will be automatically handled for you. However, there are four key tasks that still need to be completed.
|
||||
|
||||
1. Defining the abstract properties based on your model configuration.
|
||||
2. Configuring embedding and unembedding modules and the forward implementations for them.
|
||||
3. Configuring the attention configuration and desired KV cache behaviors.
|
||||
4. Writing the forward implementation for your layer.
|
||||
|
||||
## Writing a Policy
|
||||
|
||||
The `InferenceV2Policy` is the level of composition. This is the object that will be passed directly to the inference engine and will compose the model implementation and your containers to create an end-to-end solution. There are two main components to be implemented: the first is to create the model that you defined earlier. This is done by implementing the `instantiate_model` method of the policy. In general, this can just be implemented by calling the constructor for your model and passing the engine config, tensor-parallel communication object, and your custom model config.
|
||||
|
||||
The second component is to define how the parameters from the checkpoint will map to each container. From the section on `LayerContainer`s above, you may remember that the `LayerContainer` can handle the internal routing of a checkpoint parameter to its dependency. In order to find the correct `LayerContainer` though, we need a second abstraction: the `ContainerMap`.
|
||||
|
||||
A `ContainerMap` performs this mapping by categorizing checkpoint prefix strings to the type of container they map to. Typically, the easiest way to do this is through iterating over a model checkpoint's state dict or by iterating over the `named_parameters` of a PyTorch model. There are three types of mappings to define: the transformer mappings, the non-transformer mappings, and the what we'll call the rest. Let's work through an example:
|
||||
|
||||
```python
|
||||
from deepspeed.inference.module_implementations.inference_policy_base import ContainerMap
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [MyTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
map.set_transformer_params("model.layers", transformer_containers)
|
||||
|
||||
non_transformer_container = MyNonTransformerContainer(self.model)
|
||||
```
|
||||
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .inference_model_base import DSInferenceModelBase
|
||||
from .inference_transformer_base import DSTransformerModelBase, DSMoETransformerModelBase
|
||||
from .inference_policy_base import InferenceV2Policy, ContainerMap
|
||||
from .sharding import *
|
||||
|
||||
# Model Implementations
|
||||
from .llama_v2 import *
|
||||
from .opt import *
|
||||
from .mistral import *
|
||||
from .mixtral import *
|
||||
from .falcon import *
|
||||
from .phi import *
|
||||
from .phi3 import *
|
||||
from .qwen import *
|
||||
from .qwen_v2 import *
|
||||
from .qwen_v2_moe import *
|
||||
from .exaone4 import *
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .attn_output_parameters import *
|
||||
from .embedding_parameters import *
|
||||
from .mlp_parameters import *
|
||||
from .moe_parameters import *
|
||||
from .norm_parameters import *
|
||||
from .qkv_parameters import *
|
||||
from .unembed_parameters import *
|
||||
from .invfreq_parameters import *
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Common Attention Output Parameter Patterns
|
||||
"""
|
||||
|
||||
|
||||
class AttentionOutputParameter(ParameterBase):
|
||||
"""
|
||||
Attention output parameter container.
|
||||
|
||||
Note: The differentiation for something like GQA for this matrix is primarily
|
||||
encompassed in the sharding logic, which is currently expected to be performed by
|
||||
the model implementation.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Unsharded attention output parameter of shape [model_dim, model_dim]
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_attn_out_param(self.params)
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Embedding containers.
|
||||
"""
|
||||
|
||||
|
||||
class EmbeddingParameter(ParameterBase):
|
||||
"""
|
||||
Embedding container. This should be safe to use for all types of embeddings (i.e. word, position,
|
||||
and token type).
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Vocabulary parameter of shape [vocab_size, model_dim].
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_embedding_param(self.params)
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Common InvFreq Parameter Patterns
|
||||
"""
|
||||
|
||||
|
||||
class InvFreqParameter(ParameterBase):
|
||||
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.params.to(self.inference_model.activation_dtype.value)
|
||||
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
MLP Parameter Containers
|
||||
"""
|
||||
|
||||
|
||||
class MLP1Parameter(ParameterBase):
|
||||
"""
|
||||
First MLP projection weight container. This performs a straight pass-through to the
|
||||
model implementation for transformation.
|
||||
"""
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
# NOTE(cmikeh2): If we are gated but not in the format specified below, we should trigger a permutation here.
|
||||
# I am not currently aware of any models that use this format (or how we should even detect it; probably should
|
||||
# just be a different param entirely, but until then we'll just assume the format is correct).
|
||||
return self.inference_model.transform_mlp_1_param(self.params)
|
||||
|
||||
|
||||
class GatedMLPParameter(ParameterBase):
|
||||
"""
|
||||
Gated MLP projection container.
|
||||
"""
|
||||
|
||||
gate_params: torch.Tensor
|
||||
"""
|
||||
Weight parameter for the gating matrix.
|
||||
"""
|
||||
|
||||
up_params: torch.Tensor
|
||||
"""
|
||||
For lack of a better name, the non-gating weight parameters.
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
"""
|
||||
Our gated format (this is different from InferenceV1!) is to have the gate and activated neurons
|
||||
interleaved. So if we have 4 output neurons (two effective neurons) with 4 input neurons, the finalized
|
||||
parameter will look like:
|
||||
[g0_0, g0_1, g0_2, g0_3]
|
||||
[a0_0, a0_1, a0_2, a0_3]
|
||||
[g1_0, g1_1, g1_2, g1_3]
|
||||
[a1_0, a1_1, a1_2, a1_3]
|
||||
|
||||
As a reference, in inference v1, the format is:
|
||||
[g0_0, g0_1, g0_2, g0_3]
|
||||
[g1_0, g1_1, g1_2, g1_3]
|
||||
[a0_0, a0_1, a0_2, a0_3]
|
||||
[a1_0, a1_1, a1_2, a1_3]
|
||||
"""
|
||||
assert self.gate_params.shape[0] == self.up_params.shape[
|
||||
0], "Gated MLP parameters must have the same number of neurons."
|
||||
total_neurons = self.gate_params.shape[0] + self.up_params.shape[0]
|
||||
|
||||
# flip the order if even with the correct tokenizer we get wrong output
|
||||
#fused_param = torch.cat([self.up_params, self.gate_params], dim=-1).reshape(total_neurons, -1)
|
||||
fused_param = torch.cat([self.gate_params, self.up_params], dim=-1).reshape(total_neurons, -1)
|
||||
return self.inference_model.transform_mlp_1_param(fused_param)
|
||||
|
||||
|
||||
class FusedGatedMLPParameter(ParameterBase):
|
||||
"""
|
||||
Gated MLP projection container.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Weight parameter for the fused gating and non-gating weight parameters.
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
gate_params = self.params[:self.params.shape[0] // 2]
|
||||
up_params = self.params[self.params.shape[0] // 2:]
|
||||
total_neurons = gate_params.shape[0] + up_params.shape[0]
|
||||
fused_param = torch.cat([gate_params, up_params], dim=-1).reshape(total_neurons, -1)
|
||||
return self.inference_model.transform_mlp_1_param(fused_param)
|
||||
|
||||
|
||||
class MLP2Parameter(ParameterBase):
|
||||
"""
|
||||
Second MLP projection weight container. This performs a straight pass-through to the
|
||||
model implementation for transformation.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Full weight parameter.
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_mlp_2_param(self.params)
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase, ParamList
|
||||
"""
|
||||
Moe Parameters
|
||||
|
||||
These parameters are compatible with any model inheriting from ``DSMoETransformerModelBase``.
|
||||
"""
|
||||
|
||||
|
||||
class MoEGatingWeightParameter(ParameterBase):
|
||||
"""
|
||||
Gating weight matrix.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Projection matrix from the input activations to the gate logits.
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_moe_gate_param(self.params)
|
||||
|
||||
|
||||
class UnfusedMoEMLP1Parameter(ParameterBase):
|
||||
"""
|
||||
This container should be used when the experts are held in separate parameters
|
||||
and need to be joined into a single group.
|
||||
"""
|
||||
|
||||
experts: ParamList("n_experts") # noqa: F821
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
stacked_experts = torch.stack([p for p in self.experts], dim=0)
|
||||
return self.inference_model.transform_moe_mlp_1_param(stacked_experts)
|
||||
|
||||
|
||||
class UnfusedMoEMLP2Parameter(ParameterBase):
|
||||
"""
|
||||
This container should be used when the experts are held in separate parameters
|
||||
and need to be joined into a single group.
|
||||
"""
|
||||
|
||||
experts: ParamList("n_experts") # noqa: F821
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
stacked_experts = torch.stack([p for p in self.experts], dim=0)
|
||||
return self.inference_model.transform_moe_mlp_2_param(stacked_experts)
|
||||
|
||||
|
||||
class UnfusedMoEGatedMLPParameter(ParameterBase):
|
||||
"""
|
||||
MoE Parameter for a gated activation function in which the gating matrix is not
|
||||
fused in the same parameter as the non-gating matrix.
|
||||
|
||||
This is a stacked version of the ``GatedMLPParameter``. Please see that class for more
|
||||
documentation on the layout of the parameters.
|
||||
"""
|
||||
|
||||
gating_experts: ParamList("n_experts") # noqa: F821
|
||||
|
||||
up_experts: ParamList("n_experts") # noqa: F821
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
transposed_experts = []
|
||||
for gate, up in zip(self.gating_experts, self.up_experts):
|
||||
assert gate.shape[0] == up.shape[0], "Gated MLP parameters must have the same number of neurons."
|
||||
total_neurons = gate.shape[0] + up.shape[0]
|
||||
fused_expert = torch.cat([gate, up], dim=-1).reshape(total_neurons, -1)
|
||||
transposed_experts.append(fused_expert)
|
||||
|
||||
stacked_experts = torch.stack(transposed_experts, dim=0)
|
||||
return self.inference_model.transform_moe_mlp_1_param(stacked_experts)
|
||||
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Common Attention Output Parameter Patterns
|
||||
"""
|
||||
|
||||
|
||||
class NormParameter(ParameterBase):
|
||||
"""
|
||||
Simple normalization container.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_norm_param(self.params)
|
||||
@@ -0,0 +1,115 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Common QKV Parameter Patterns
|
||||
"""
|
||||
|
||||
|
||||
class FusedQKVParameter(ParameterBase):
|
||||
"""
|
||||
Traditional fused QKV parameters for QKV projection. This is functionally
|
||||
a direct copy.
|
||||
|
||||
src_qkv_w shape: [3 * out_features, in_features]
|
||||
qkv_w shape: [3 * out_features, in_features]
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_qkv_param(self.params)
|
||||
|
||||
|
||||
class UnfusedQKVParameter(ParameterBase):
|
||||
"""
|
||||
QKV parameter container for unfused QKV projection.
|
||||
|
||||
src_param shapes: 3 x [out_features, in_features]
|
||||
dst_param shape: [3 x out_features, in_features]
|
||||
"""
|
||||
|
||||
q_params: torch.Tensor
|
||||
|
||||
k_params: torch.Tensor
|
||||
|
||||
v_params: torch.Tensor
|
||||
|
||||
def finalize(self):
|
||||
fused_param = torch.cat([self.q_params, self.k_params, self.v_params], dim=0)
|
||||
return self.inference_model.transform_qkv_param(fused_param)
|
||||
|
||||
|
||||
def megatron_qkv_reshape(param: torch.Tensor, head_size: int, n_heads: int) -> torch.Tensor:
|
||||
assert param.shape[0] == 3 * n_heads * head_size
|
||||
|
||||
all_heads = torch.chunk(param, chunks=3 * n_heads, dim=0)
|
||||
q_heads = all_heads[::3]
|
||||
k_heads = all_heads[1::3]
|
||||
v_heads = all_heads[2::3]
|
||||
return torch.cat([q_heads, k_heads, v_heads], dim=0)
|
||||
|
||||
|
||||
class MegatronQKVParameter(ParameterBase):
|
||||
"""
|
||||
QKV parameter container for Megatron-style QKV projection. Megatron stores the parameter
|
||||
as [n_heads, 3, head_size, in_features] whereas our inference system is built around
|
||||
[3, n_heads, head_size, in_features]. This container handles the conversion.
|
||||
|
||||
Note: this container expects the model implementation to implement properties for
|
||||
`head_size` and `n_heads`.
|
||||
|
||||
src_qkv_w shape: [3 * out_features, in_features]
|
||||
qkv_w shape: [3 * out_features, in_features]
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
head_size = self.inference_model.head_size
|
||||
n_heads = self.inference_model.n_heads
|
||||
|
||||
transposed_param = megatron_qkv_reshape(self.params, head_size, n_heads)
|
||||
return self.inference_model.transform_qkv_param(transposed_param)
|
||||
|
||||
|
||||
def transform_gqa_megatron(src_param: torch.Tensor, head_size: int, n_q_heads: int, n_kv_heads: int) -> torch.Tensor:
|
||||
assert src_param.shape[0] == (2 * n_kv_heads + n_q_heads) * head_size
|
||||
|
||||
head_ratio = n_q_heads // n_kv_heads
|
||||
|
||||
# Reshape to get the groups as the leading dimension
|
||||
groups_leading_view = src_param.reshape(n_kv_heads, 2 + head_ratio, head_size, -1)
|
||||
q_heads = groups_leading_view[:, :head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1])
|
||||
k_heads = groups_leading_view[:, head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1])
|
||||
v_heads = groups_leading_view[:, head_ratio + 1, :, :].reshape(-1, groups_leading_view.shape[-1])
|
||||
# Squeeze will remove extra dimension for bias
|
||||
return torch.cat([q_heads, k_heads, v_heads], dim=0).squeeze()
|
||||
|
||||
|
||||
class GQAMegatronQKVParameter(ParameterBase):
|
||||
"""
|
||||
QKV parameter for Megatron-style QKV projection with GQA-style QKV projection. In this
|
||||
storage format each of the groups is stored consecutively, so there will be multiple q_heads,
|
||||
then one k head, and one v head.
|
||||
|
||||
Note: this container expects the model implementation to implement properties for
|
||||
`head_size`, `n_q_heads`, and `n_kv_heads`.
|
||||
|
||||
src_qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features]
|
||||
qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features]
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
head_size = self.inference_model.head_size
|
||||
n_q_heads = self.inference_model.n_heads_q
|
||||
n_kv_heads = self.inference_model.n_heads_kv
|
||||
transposed_param = transform_gqa_megatron(self.params, head_size, n_q_heads, n_kv_heads)
|
||||
return self.inference_model.transform_qkv_param(transposed_param)
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from ...model_implementations.parameter_base import ParameterBase
|
||||
"""
|
||||
Unembedding containers.
|
||||
"""
|
||||
|
||||
|
||||
class UnembedParameter(ParameterBase):
|
||||
"""
|
||||
Unembedding parameter. This will likely be mapped to the same original weight in the model as the
|
||||
embedding, but we have a different preferred sharding approach.
|
||||
"""
|
||||
|
||||
params: torch.Tensor
|
||||
"""
|
||||
Unembedding parameter of shape [vocab_size, model_dim].
|
||||
"""
|
||||
|
||||
def finalize(self) -> torch.Tensor:
|
||||
return self.inference_model.transform_unembed_param(self.params)
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import Exaone4Policy
|
||||
@@ -0,0 +1,49 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.inference.v2.model_implementations.common_parameters import *
|
||||
from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
|
||||
|
||||
|
||||
class Exaone4TransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the EXAONE 4.0 model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: GatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
q_norm_gamma: NormParameter
|
||||
k_norm_gamma: NormParameter
|
||||
post_attn_norm_gamma: NormParameter
|
||||
post_ff_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate_proj.weight": "mlp_1_w.gate_params",
|
||||
"mlp.up_proj.weight": "mlp_1_w.up_params",
|
||||
"mlp.down_proj.weight": "mlp_2_w.params",
|
||||
"self_attn.q_norm.weight": "q_norm_gamma.params",
|
||||
"self_attn.k_norm.weight": "k_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "post_attn_norm_gamma.params",
|
||||
"post_feedforward_layernorm.weight": "post_ff_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class Exaone4NonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the EXAONE 4.0 model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from ...model_implementations import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
from ...kernels.core_ops.cuda_rms_norm.rms_norm import CUDARMSNorm
|
||||
|
||||
from .container import Exaone4NonTransformerContainer, Exaone4TransformerContainer
|
||||
|
||||
|
||||
class Exaone4InferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for EXAONE 4.0 models.
|
||||
|
||||
Key differences from Mistral/Llama:
|
||||
- Post-norm architecture (norm after attn/mlp, not before)
|
||||
- QK-Norm (RMSNorm on Q and K projections per head)
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[Exaone4NonTransformerContainer]
|
||||
_transformer: Optional[Iterable[Exaone4TransformerContainer]]
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_position_embeddings
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return getattr(self._config, "head_dim", self.model_dim // self.n_heads)
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
activation = self._config.hidden_act.lower()
|
||||
if activation == "silu":
|
||||
return ActivationType.SiGLU
|
||||
elif activation == "gelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "relu":
|
||||
return ActivationType.ReGLU
|
||||
else:
|
||||
raise NotImplementedError(f"Activation {activation} not supported")
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
rope_theta = getattr(self._config, "rope_theta", 1000000.0)
|
||||
return RotateHalfConfig(theta_base=rope_theta)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._qk_norm = CUDARMSNorm(
|
||||
channels=self.head_size,
|
||||
fp_dtype=torch.float16 if self.activation_dtype == DtypeEnum.fp16 else torch.bfloat16,
|
||||
epsilon=getattr(self._config, "rms_norm_eps", 1e-5),
|
||||
)
|
||||
|
||||
def _apply_qk_norm(self, hidden_states: torch.Tensor, q_norm_gamma: torch.Tensor,
|
||||
k_norm_gamma: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply RMSNorm to Q and K projections independently per head.
|
||||
hidden_states shape: [tokens, (n_q + n_kv + n_kv) * head_size]
|
||||
"""
|
||||
tokens = hidden_states.shape[0]
|
||||
local_n_heads = self.n_heads_q_local
|
||||
local_n_heads_kv = self.n_heads_kv_local
|
||||
q_len = local_n_heads * self.head_size
|
||||
kv_len = local_n_heads_kv * self.head_size
|
||||
|
||||
q = hidden_states[:, :q_len].contiguous()
|
||||
k = hidden_states[:, q_len:q_len + kv_len].contiguous()
|
||||
v = hidden_states[:, q_len + kv_len:]
|
||||
|
||||
# Reshape to [tokens * n_heads, head_size] for per-head RMSNorm
|
||||
q = q.view(-1, self.head_size)
|
||||
self._qk_norm(q, q, q_norm_gamma)
|
||||
q = q.view(tokens, q_len)
|
||||
|
||||
k = k.view(-1, self.head_size)
|
||||
self._qk_norm(k, k, k_norm_gamma)
|
||||
k = k.view(tokens, kv_len)
|
||||
|
||||
hidden_states[:, :q_len] = q
|
||||
hidden_states[:, q_len:q_len + kv_len] = k
|
||||
|
||||
return hidden_states
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
return embed
|
||||
|
||||
def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
EXAONE 4.0 uses post-norm architecture:
|
||||
hidden = attn(hidden)
|
||||
hidden = post_attn_norm(hidden)
|
||||
residual = residual + hidden
|
||||
hidden = mlp(residual)
|
||||
hidden = post_ff_norm(hidden)
|
||||
residual = residual + hidden
|
||||
"""
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
# Attention block
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None)
|
||||
hidden_states = self._apply_qk_norm(hidden_states, cur_params.q_norm_gamma, cur_params.k_norm_gamma)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
# Post-attn norm + residual add
|
||||
_, hidden_states = self.norm(hidden_states, None, cur_params.post_attn_norm_gamma, beta=None)
|
||||
residual.add_(hidden_states)
|
||||
|
||||
# MLP block
|
||||
hidden_states = self.mlp_1(residual, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
# Post-ff norm + residual add
|
||||
_, hidden_states = self.norm(hidden_states, None, cur_params.post_ff_norm_gamma, beta=None)
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, residual
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer(layer_idx, residual, residual, wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import Exaone4NonTransformerContainer, Exaone4TransformerContainer
|
||||
from .model import Exaone4InferenceModel
|
||||
|
||||
|
||||
class Exaone4Policy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Exaone4InferenceModel:
|
||||
return Exaone4InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [Exaone4TransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(Exaone4NonTransformerContainer(self.model))
|
||||
map.set_unmapped_params([])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import FalconPolicy
|
||||
@@ -0,0 +1,129 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Falcon 7b model looks like this:
|
||||
|
||||
FalconForCausalLM(
|
||||
(transformer): FalconModel(
|
||||
(word_embeddings): Embedding(65024, 4544)
|
||||
(h): ModuleList(
|
||||
(0-31): 32 x FalconDecoderLayer(
|
||||
(self_attention): FalconAttention(
|
||||
(maybe_rotary): FalconRotaryEmbedding()
|
||||
(query_key_value): FalconLinear(in_features=4544, out_features=4672, bias=False)
|
||||
(dense): FalconLinear(in_features=4544, out_features=4544, bias=False)
|
||||
(attention_dropout): Dropout(p=0.0, inplace=False)
|
||||
)
|
||||
(mlp): FalconMLP(
|
||||
(dense_h_to_4h): FalconLinear(in_features=4544, out_features=18176, bias=False)
|
||||
(act): GELU(approximate='none')
|
||||
(dense_4h_to_h): FalconLinear(in_features=18176, out_features=4544, bias=False)
|
||||
)
|
||||
(input_layernorm): LayerNorm((4544,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
)
|
||||
(ln_f): LayerNorm((4544,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
(lm_head): Linear(in_features=4544, out_features=65024, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class FalconTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Falcon model.
|
||||
"""
|
||||
qkv_w: FusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: MLP1Parameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
ln_attn_gamma: NormParameter
|
||||
ln_attn_beta: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attention.query_key_value.weight": "qkv_w.params",
|
||||
"self_attention.dense.weight": "attn_out_w.params",
|
||||
"mlp.dense_h_to_4h.weight": "mlp_1_w.params",
|
||||
"mlp.dense_4h_to_h.weight": "mlp_2_w.params",
|
||||
"input_layernorm.weight": "ln_attn_gamma.params",
|
||||
"input_layernorm.bias": "ln_attn_beta.params",
|
||||
}
|
||||
|
||||
|
||||
class FalconNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Falcon model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm_gamma: NormParameter
|
||||
final_norm_beta: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"transformer.word_embeddings.weight": "word_emb.params",
|
||||
"transformer.ln_f.weight": "final_norm_gamma.params",
|
||||
"transformer.ln_f.bias": "final_norm_beta.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
|
||||
|
||||
'''
|
||||
# HF Falcon 40b model looks like this:
|
||||
|
||||
FalconForCausalLM(
|
||||
(transformer): FalconModel(
|
||||
(word_embeddings): Embedding(65024, 8192)
|
||||
(h): ModuleList(
|
||||
(0-59): 60 x FalconDecoderLayer(
|
||||
(self_attention): FalconAttention(
|
||||
(maybe_rotary): FalconRotaryEmbedding()
|
||||
(query_key_value): FalconLinear(in_features=8192, out_features=9216, bias=False)
|
||||
(dense): FalconLinear(in_features=8192, out_features=8192, bias=False)
|
||||
(attention_dropout): Dropout(p=0.0, inplace=False)
|
||||
)
|
||||
(mlp): FalconMLP(
|
||||
(dense_h_to_4h): FalconLinear(in_features=8192, out_features=32768, bias=False)
|
||||
(act): GELU(approximate='none')
|
||||
(dense_4h_to_h): FalconLinear(in_features=32768, out_features=8192, bias=False)
|
||||
)
|
||||
(ln_attn): LayerNorm((8192,), eps=1e-05, elementwise_affine=True)
|
||||
(ln_mlp): LayerNorm((8192,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
)
|
||||
(ln_f): LayerNorm((8192,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
(lm_head): Linear(in_features=8192, out_features=65024, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class FalconNewArchTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Falcon model.
|
||||
"""
|
||||
qkv_w: GQAMegatronQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: MLP1Parameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
ln_attn_gamma: NormParameter
|
||||
ln_attn_beta: NormParameter
|
||||
ln_mlp_gamma: NormParameter
|
||||
ln_mlp_beta: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attention.query_key_value.weight": "qkv_w.params",
|
||||
"self_attention.dense.weight": "attn_out_w.params",
|
||||
"mlp.dense_h_to_4h.weight": "mlp_1_w.params",
|
||||
"mlp.dense_4h_to_h.weight": "mlp_2_w.params",
|
||||
"ln_attn.weight": "ln_attn_gamma.params",
|
||||
"ln_attn.bias": "ln_attn_beta.params",
|
||||
"ln_mlp.weight": "ln_mlp_gamma.params",
|
||||
"ln_mlp.bias": "ln_mlp_beta.params",
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .container import FalconNonTransformerContainer, FalconTransformerContainer
|
||||
|
||||
|
||||
class FalconInferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[FalconNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[FalconTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties inherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties inherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return 4 * self._config.hidden_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_kv_heads if (self._config.new_decoder_architecture
|
||||
or not self._config.multi_query) else 1
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.GELU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.LayerNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> RotateHalfConfig:
|
||||
"""
|
||||
The positional embedding configuration for the model.
|
||||
"""
|
||||
return RotateHalfConfig()
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
assert self.config.parallel_attn, "Only parallel attention implementation is supported"
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
attn_ln_out = hidden_states
|
||||
attn_hidden_state = self.qkv(attn_ln_out, cur_params.qkv_w, b=None)
|
||||
attn_hidden_state = self.attn(attn_hidden_state, kv_cache, ragged_batch_info)
|
||||
attention_output = self.attn_out(attn_hidden_state, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.config.new_decoder_architecture:
|
||||
residual, mlp_ln_out = self.norm(residual,
|
||||
None,
|
||||
gamma=cur_params.ln_mlp_gamma,
|
||||
beta=cur_params.ln_mlp_beta)
|
||||
else:
|
||||
mlp_ln_out = hidden_states
|
||||
|
||||
mlp_hidden_state = self.mlp_1(mlp_ln_out, cur_params.mlp_1_w, b=None)
|
||||
mlp_output = self.mlp_2(mlp_hidden_state, cur_params.mlp_2_w, b=None)
|
||||
|
||||
mlp_output.add_(attention_output)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(mlp_output, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, mlp_output = self.norm(residual,
|
||||
mlp_output,
|
||||
next_params.ln_attn_gamma,
|
||||
beta=next_params.ln_attn_beta)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(mlp_output)
|
||||
|
||||
return residual, mlp_output
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm_gamma,
|
||||
beta=self._non_transformer.final_norm_beta)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual,
|
||||
None,
|
||||
gamma=self._transformer[0].ln_attn_gamma,
|
||||
beta=self._transformer[0].ln_attn_beta)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import FalconNonTransformerContainer, FalconTransformerContainer
|
||||
from .container import FalconNewArchTransformerContainer
|
||||
from .model import FalconInferenceModel
|
||||
|
||||
|
||||
class FalconPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> FalconInferenceModel:
|
||||
return FalconInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
trans_container_cls = FalconNewArchTransformerContainer if self._model_config.new_decoder_architecture else FalconTransformerContainer
|
||||
transformer_containers = [trans_container_cls(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['transformer.h'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(FalconNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(
|
||||
[f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)])
|
||||
|
||||
return map
|
||||
@@ -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)
|
||||
@@ -0,0 +1,272 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Iterable, Optional, Tuple, Type
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
from ..ragged import DSStateManager, RaggedBatchWrapper
|
||||
from ..ragged.manager_configs import KVCacheConfig
|
||||
from ..ragged import DSSequenceDescriptor
|
||||
from ..model_implementations.layer_container_base import LayerContainer
|
||||
from ..config_v2 import RaggedInferenceEngineConfig
|
||||
from .flat_model_helpers import ModelMetadata
|
||||
|
||||
try:
|
||||
from functools import cached_property
|
||||
except ImportError:
|
||||
|
||||
def cached_property(func):
|
||||
return property(func)
|
||||
|
||||
|
||||
"""
|
||||
This abstract class defines the interfaces that a model implementation should implement
|
||||
in order to include anything that may be called by the engine. Most models should be able
|
||||
to inherit from `DSInferenceTransformerModelBase` to reduce implementation work so it is recommended
|
||||
to begin there.
|
||||
"""
|
||||
"""
|
||||
Placeholder for typing the model config, which can vary based on model implementation/
|
||||
"""
|
||||
DSModelImplementationConfig = Type['DSModelImplementationConfig']
|
||||
"""
|
||||
Placeholder for typing the distributed comm object.
|
||||
|
||||
TODO(cmikeh2): Replace when we have a more defined API for the inference communication system.
|
||||
"""
|
||||
MPType = Type["MPType"]
|
||||
|
||||
|
||||
class DSInferenceModelBase(torch.nn.Module, ABC):
|
||||
"""
|
||||
Implementation of a model for inference composable with ragged batching.
|
||||
"""
|
||||
|
||||
_config: DSModelImplementationConfig
|
||||
"""
|
||||
Model-specific configuration. No abstraction surrounds this yet.
|
||||
"""
|
||||
|
||||
_engine_config: RaggedInferenceEngineConfig
|
||||
"""
|
||||
Engine configuration.
|
||||
"""
|
||||
|
||||
_base_mp_group: MPType
|
||||
"""
|
||||
Base communication group for Tensor-parallel inference.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[LayerContainer]
|
||||
"""
|
||||
Abstract container for storing both embedding (pre-transformer) and unembedding (post-transformer)
|
||||
parameters. This attribute should be None at model instantiation until the Policy sets
|
||||
the model parameters. These parameters are grouped together since many model implementations
|
||||
will tie the embedding and unembedding parameters together.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[LayerContainer]]
|
||||
"""
|
||||
List of abstract containers (1 per layer) for storing transformer (transformer)
|
||||
parameters. This attribute should be None at model instantiation until the Policy
|
||||
sets the model parameters.
|
||||
"""
|
||||
|
||||
state_manager: Optional[DSStateManager]
|
||||
"""
|
||||
Since the state manager is lazy initialized, by the engine, it is not guaranteed to be present
|
||||
until full initialization.
|
||||
"""
|
||||
|
||||
def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig,
|
||||
base_mp_group: MPType) -> None:
|
||||
"""
|
||||
Minimal initialization of the model.
|
||||
|
||||
Arguments:
|
||||
config (DSModelImplementationConfig): Model-specific configuration. No assumptions
|
||||
should be made about this config that are not closely tied to the specific
|
||||
model implementation.
|
||||
engine_config (RaggedInferenceEngineConfig): Engine configuration.
|
||||
base_mp_group (MPType): Base communication group for Tensor-parallel inference.
|
||||
"""
|
||||
super().__init__()
|
||||
self._config = config
|
||||
self._engine_config = engine_config
|
||||
self._base_mp_group = base_mp_group
|
||||
|
||||
# Set to None until the Policy sets the model parameters
|
||||
self._non_transformer = None
|
||||
self._transformer = None
|
||||
self._flattened_param_buffer = None
|
||||
self._flattened_param_metadata = None
|
||||
|
||||
@property
|
||||
def config(self) -> DSModelImplementationConfig:
|
||||
"""
|
||||
The model config.
|
||||
"""
|
||||
return self._config
|
||||
|
||||
def set_parameters(self, transformer: Iterable[LayerContainer], non_transformer: LayerContainer,
|
||||
flattened_param_buffer: torch.Tensor, flattened_param_metadata: ModelMetadata):
|
||||
"""
|
||||
Set the model parameters for the embedding, transformer, and unembedding containers.
|
||||
"""
|
||||
self._transformer = transformer
|
||||
self._non_transformer = non_transformer
|
||||
self._flattened_param_buffer = flattened_param_buffer
|
||||
self._flattened_param_metadata = flattened_param_metadata
|
||||
|
||||
def set_state_manager(self, state_manager: DSStateManager):
|
||||
"""
|
||||
Sets the state manager attribute. This is called by the inference engine after
|
||||
the model is fully initialized.
|
||||
"""
|
||||
self.state_manager = state_manager
|
||||
|
||||
@cached_property
|
||||
def tp_rank(self) -> int:
|
||||
"""
|
||||
The rank of the current process.
|
||||
|
||||
# TODO(cmikeh2): Kind of a hack right now, but this is too verbose to use at
|
||||
the frequency we need.
|
||||
"""
|
||||
return dist.get_rank(group=self._base_mp_group)
|
||||
|
||||
@cached_property
|
||||
def tp_size(self) -> int:
|
||||
"""
|
||||
The total number of processes.
|
||||
|
||||
# TODO(cmikeh2): Kind of a hack right now, but this is too verbose to use at
|
||||
the frequency we need.
|
||||
"""
|
||||
return dist.get_world_size(group=self._base_mp_group)
|
||||
|
||||
@property
|
||||
def model_config(self):
|
||||
"""
|
||||
The model config.
|
||||
"""
|
||||
return self._config
|
||||
|
||||
@property
|
||||
def engine_config(self):
|
||||
"""
|
||||
The engine config.
|
||||
"""
|
||||
return self._engine_config
|
||||
|
||||
@property
|
||||
def flattened_params(self) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
The flattened parameter buffer.
|
||||
"""
|
||||
return self._flattened_param_buffer
|
||||
|
||||
@property
|
||||
def flattened_param_metadata(self) -> Optional[ModelMetadata]:
|
||||
"""
|
||||
The flattened parameter metadata.
|
||||
"""
|
||||
return self._flattened_param_metadata
|
||||
|
||||
@abstractmethod
|
||||
def get_kv_requirements(self, sequence: DSSequenceDescriptor, max_new_tokens: int,
|
||||
max_new_blocks: Tuple[int, ...]) -> Tuple[int, torch.Tensor]:
|
||||
"""
|
||||
Given a sequence and the number of new tokens in the sequence, determine the
|
||||
number of new KV blocks needed to support the sequence. This method is
|
||||
used to help the engine provide schedulability APIs and can be used as a helper
|
||||
for ``maybe_allocate_kv``.
|
||||
|
||||
Args:
|
||||
sequence (DSSequenceDescriptor): The sequence for which to allocate KV-storage.
|
||||
max_new_tokens (int): Maximum number of tokens to hypothetically schedule.
|
||||
max_new_blocks (int): Maximum number of blocks to hypothetically allocate.
|
||||
|
||||
Returns:
|
||||
Tuple[int, torch.Tensor]: The tuple of number of tokens scheduled and number
|
||||
of blocks allocated (per KV cache). In general, only one of these numbers will
|
||||
match the corresponding input argument, but this is not guaranteed.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def get_remaining_block_capacity(self, sequence: DSSequenceDescriptor) -> int:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def maybe_allocate_kv(self, sequence: DSSequenceDescriptor, n_new_tokens: int) -> None:
|
||||
"""
|
||||
Given a sequence and the number of new tokens in the sequence, determine
|
||||
whether or not additional KV-storage is needed and allocate it if so.
|
||||
|
||||
Args:
|
||||
sequence (DSSequenceDescriptor): The sequence for which to allocate KV-storage.
|
||||
n_new_tokens (int): The number of new tokens in the sequence.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def kv_cache_config(self) -> Tuple[KVCacheConfig, ...]:
|
||||
"""
|
||||
Return the KV-cache configuration for this model. This should be a tuple of one or more
|
||||
KVCacheConfig objects (one for each distinct cache group).
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def max_sequence_length(self) -> int:
|
||||
"""
|
||||
The maximum sequence length supported by the model.
|
||||
"""
|
||||
...
|
||||
|
||||
def maybe_free_kv(self, sequence: DSSequenceDescriptor) -> None:
|
||||
"""
|
||||
After completing a forward pass, determine whether or not there are any KV blocks
|
||||
that maybe freed since they are no longer in use.
|
||||
|
||||
Consider the following example:
|
||||
|
||||
We have a block size of 4 and a local window size of 8. At the beginning of the forward
|
||||
pass there 10 tokens had been seen and the new forward has a size of 4. This would lend
|
||||
itself to the following cache structure prior to the forward:
|
||||
[[0, 1, 2*, 3*] [4*, 5*, 6*, 7*] [8*, 9*, x, x] [x x x x]]
|
||||
Where x's denote empty cache locations and * denote values that are needed for attention
|
||||
of the next open slot. After the forward, the cache would look like the following:
|
||||
[[0, 1, 2, 3] [4, 5, 6*, 7*] [8*, 9*, 10*, 11*] [12* 13* x x]]
|
||||
In this case, the first block is no longer needed since it is not needed for any future
|
||||
local attention windows. This function would be responsible for freeing that block.
|
||||
|
||||
Default behavior assumes no local patterns that require freeing and in general should
|
||||
be sufficient.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def prepare_batch(self, wrapped_batch: RaggedBatchWrapper) -> None:
|
||||
"""
|
||||
This will be called before each forward with the intent of building forward-specific metadata
|
||||
about a batch. The intent here is to build data structures like attention atoms without necessarily
|
||||
needing to implement graphable kernels to do so.
|
||||
|
||||
Abstract so as to force model implementations to opt out of doing anything here explicitly.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def forward(wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Complete a forward pass of the model. This interface should be graphable, so it
|
||||
should not rely on the ability to use python control flow.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -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)
|
||||
@@ -0,0 +1,617 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ..config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_utils import ActivationType, ceil_div, is_gated
|
||||
from ..model_implementations import *
|
||||
from ..model_implementations.sharding import *
|
||||
from ..modules.configs import (
|
||||
DSEmbeddingsConfig,
|
||||
DSLinearConfig,
|
||||
DSMoEConfig,
|
||||
DSNormConfig,
|
||||
DSSelfAttentionConfig,
|
||||
DSUnembedConfig,
|
||||
NormTypeEnum,
|
||||
PositionalEmbeddingType,
|
||||
RotateHalfConfig,
|
||||
)
|
||||
from ..modules import heuristics
|
||||
from ..ragged import (
|
||||
DSSequenceDescriptor,
|
||||
KVCacheConfig,
|
||||
RaggedBatchWrapper,
|
||||
)
|
||||
from .inference_model_base import (
|
||||
DSInferenceModelBase,
|
||||
DSModelImplementationConfig,
|
||||
MPType,
|
||||
)
|
||||
from ..inference_parameter import InferenceParameter
|
||||
|
||||
try:
|
||||
from functools import cached_property
|
||||
except ImportError:
|
||||
|
||||
def cached_property(func):
|
||||
return property(func)
|
||||
|
||||
|
||||
class DSTransformerModelBase(DSInferenceModelBase):
|
||||
"""
|
||||
Dimensioning properties
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def num_layers(self) -> int:
|
||||
"""
|
||||
Number of the layers in the model
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def model_dim(self) -> int:
|
||||
"""
|
||||
Size of embedding projection and residuals.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def vocab_size(self) -> int:
|
||||
"""
|
||||
Size of the vocabulary (including padding).
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def head_size(self) -> int:
|
||||
"""
|
||||
Size of each attention head.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def n_heads(self) -> int:
|
||||
"""
|
||||
The number of query heads on the model. This should not take into account
|
||||
any dimension reductions from model sharding.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
def n_heads_q(self) -> int:
|
||||
"""
|
||||
Alias to n_heads.
|
||||
"""
|
||||
return self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
"""
|
||||
The number of key and value heads on the model. For GQA or MQA, overload this attribute.
|
||||
Otherwise it adopts MHA formulations and uses n_heads. This should not take into account
|
||||
any dimension reductions from model sharding.
|
||||
"""
|
||||
return self.n_heads
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def intermediate_dim(self) -> int:
|
||||
"""
|
||||
The size of the (unsharded) intermediate projection dim. For a gated activation function
|
||||
this is the size of the input to the second MLP layer. This should not take into account
|
||||
any dimension reductions from model sharding.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
"""
|
||||
The type of positional embedding used by the model.
|
||||
"""
|
||||
...
|
||||
|
||||
"""
|
||||
Architectural properties
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def activation_dtype(self) -> torch.dtype:
|
||||
"""
|
||||
The activation dtype of the model.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
"""
|
||||
The activation function used in the MLP.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
"""
|
||||
The type of normalization used in the model.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
"""
|
||||
The positional embedding configuration for the model.
|
||||
"""
|
||||
...
|
||||
|
||||
"""
|
||||
Derived helpers
|
||||
"""
|
||||
|
||||
@cached_property
|
||||
def n_heads_q_local(self) -> int:
|
||||
"""
|
||||
Number of local heads post sharding.
|
||||
"""
|
||||
return get_local_heads(self.tp_rank, self.tp_size, self.n_heads_q, self.n_heads_kv)[0]
|
||||
|
||||
@cached_property
|
||||
def n_heads_kv_local(self) -> int:
|
||||
"""
|
||||
Number of local heads post sharding.
|
||||
"""
|
||||
return get_local_heads(self.tp_rank, self.tp_size, self.n_heads_q, self.n_heads_kv)[1]
|
||||
|
||||
@property
|
||||
def gated_mlp(self) -> bool:
|
||||
"""
|
||||
Return a boolean to determine whether the model uses a gated activation function.
|
||||
"""
|
||||
return is_gated(self.mlp_activation_fn)
|
||||
|
||||
"""
|
||||
Method implementations
|
||||
"""
|
||||
|
||||
def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig,
|
||||
base_mp_group: MPType) -> None:
|
||||
"""
|
||||
Base implementation for initialization. By default, this will initialize
|
||||
the traditional components of a transformer model:
|
||||
- Embedding
|
||||
- QKV projection
|
||||
- Self attention
|
||||
- Attention output projection
|
||||
- Feed forward network
|
||||
- Normalization
|
||||
- Unembedding
|
||||
|
||||
Arguments:
|
||||
config (DSModelImplementationConfig): Model-specific configuration. No assumptions
|
||||
should be made about this config that are not closely tied to the specific
|
||||
model implementation.
|
||||
engine_config (RaggedInferenceEngineConfig): Engine configuration.
|
||||
base_mp_group (MPType): Base communication group for Tensor-parallel inference.
|
||||
"""
|
||||
super().__init__(config, engine_config, base_mp_group)
|
||||
|
||||
self.make_norm_layer()
|
||||
self.make_qkv_layer()
|
||||
self.make_attn_layer()
|
||||
self.make_attn_out_layer()
|
||||
self.make_mlp_1_layer()
|
||||
self.make_mlp_2_layer()
|
||||
self.make_embedding_layer()
|
||||
self.make_unembedding_layer()
|
||||
self._kv_cache_config = None
|
||||
|
||||
######### Embedding #########
|
||||
def make_embedding_layer(self) -> None:
|
||||
"""
|
||||
Performs setup and creates embedding DSModule. This will set the `self.embed` attribute.
|
||||
"""
|
||||
|
||||
embed_config = DSEmbeddingsConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
residual_dtype=self.activation_dtype,
|
||||
embedding_dim=self.model_dim,
|
||||
)
|
||||
|
||||
self.embed = heuristics.instantiate_embed(embed_config, self._engine_config)
|
||||
|
||||
def transform_embedding_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Performs embedding sharding along the channels dimension.
|
||||
"""
|
||||
# Until we can do non-contiguous all-gather, we won't shard the embedding parameters.
|
||||
param = param.to(self.activation_dtype.value)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
######### Unembedding #########
|
||||
def make_unembedding_layer(self) -> None:
|
||||
"""
|
||||
Performs setup and creates an unembedding layer. This implementation assumes
|
||||
normalization prior to the LM head projection. If this does not match the model's
|
||||
implementation, override this method. This will set the ``self.unembed`` attribute.
|
||||
"""
|
||||
unembed_dim = sharded_unembed_dim(self.vocab_size, self.tp_rank, self.tp_size)
|
||||
|
||||
unembed_config = DSUnembedConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
max_sequences=self._engine_config.state_manager.max_ragged_sequence_count,
|
||||
dtype=self.activation_dtype,
|
||||
model_dim=self.model_dim,
|
||||
vocab_size=unembed_dim,
|
||||
norm_type=self.norm_type,
|
||||
)
|
||||
|
||||
self.unembed = heuristics.instantiate_unembed(unembed_config, self._engine_config)
|
||||
|
||||
if self.tp_size > 1:
|
||||
self._comm_logits = torch.empty(self.tp_size,
|
||||
self._engine_config.state_manager.max_ragged_sequence_count,
|
||||
unembed_dim,
|
||||
device=get_accelerator().current_device(),
|
||||
dtype=self.activation_dtype.value)
|
||||
self._return_logits = torch.empty(self._engine_config.state_manager.max_ragged_sequence_count,
|
||||
self.vocab_size,
|
||||
device=get_accelerator().current_device(),
|
||||
dtype=self.activation_dtype.value)
|
||||
|
||||
def transform_unembed_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Performs sharding along the vocab dimension.
|
||||
"""
|
||||
param = shard_unembed_param(param, self.tp_rank, self.tp_size).to(self.activation_dtype.value)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
######### QKV #########
|
||||
def make_qkv_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the QKV linear layer. This sets the
|
||||
`self.qkv` attribute.
|
||||
"""
|
||||
out_features = qkv_out_features(self.model_dim, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q,
|
||||
self.n_heads_kv)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=self.model_dim,
|
||||
out_channels=out_features,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.qkv = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
def transform_qkv_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Passes a QKV parameter to the underlying implementation for any necessary
|
||||
transformations.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have
|
||||
the shape (out_neurons, in_neurons)
|
||||
"""
|
||||
param = shard_qkv_param(param, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q, self.n_heads_kv)
|
||||
return self.qkv.transform_param(param)
|
||||
|
||||
######### Attention #########
|
||||
def make_attn_layer(self) -> None:
|
||||
"""
|
||||
Builds the attention layer for the model. This sets the `self.attn` attribute.
|
||||
"""
|
||||
softmax_scale = 1.0 / (self.head_size**0.5)
|
||||
|
||||
attn_config = DSSelfAttentionConfig(max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
n_heads_q=self.n_heads_q_local,
|
||||
n_heads_kv=self.n_heads_kv_local,
|
||||
head_size=self.head_size,
|
||||
max_sequences=self._engine_config.state_manager.max_ragged_sequence_count,
|
||||
scale_factor=softmax_scale,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
positional_embedding_type=self.positional_embedding_type,
|
||||
positional_embedding_config=self.positional_embedding_config)
|
||||
|
||||
self.attn = heuristics.instantiate_attention(attn_config, self._engine_config)
|
||||
|
||||
def get_kv_requirements(self, sequence: DSSequenceDescriptor, max_new_tokens: int,
|
||||
max_new_blocks: int) -> Tuple[int, int]:
|
||||
"""
|
||||
See ``DSInferenceModelBase.get_kv_requirements`` for documentation.
|
||||
|
||||
This method assumes an autoregressive dense attention pattern. Override this method
|
||||
if this does not match the model's attention pattern.
|
||||
"""
|
||||
total_tokens = sequence.seen_tokens + max_new_tokens
|
||||
req_blocks = ceil_div(total_tokens, self.attn.kv_block_size)
|
||||
block_lim = req_blocks - sequence.cur_allocated_blocks
|
||||
|
||||
if block_lim <= max_new_blocks:
|
||||
return max_new_tokens, block_lim
|
||||
|
||||
token_capacity = (max_new_blocks +
|
||||
sequence.cur_allocated_blocks) * self.attn.kv_block_size - sequence.seen_tokens
|
||||
|
||||
return token_capacity, max_new_blocks
|
||||
|
||||
def get_remaining_block_capacity(self, sequence: DSSequenceDescriptor) -> int:
|
||||
return sequence.seen_tokens % self.attn.kv_block_size
|
||||
|
||||
def maybe_allocate_kv(self, sequence: DSSequenceDescriptor, n_new_tokens: int) -> None:
|
||||
"""
|
||||
See ``DSInferenceModelBase.maybe_allocate_kv`` for documentation.
|
||||
|
||||
This method assumes an autoregressive dense attention pattern. Override this method
|
||||
if this does not match the model's attention pattern.
|
||||
"""
|
||||
free_block = self.state_manager.free_blocks[0]
|
||||
_, n_needed_blocks = self.get_kv_requirements(sequence, n_new_tokens, free_block)
|
||||
|
||||
if n_needed_blocks > 0:
|
||||
new_blocks = self.state_manager.allocate_blocks(n_needed_blocks)
|
||||
sequence.extend_kv_cache(new_blocks)
|
||||
|
||||
def kv_cache_config(self) -> Tuple[KVCacheConfig, ...]:
|
||||
"""
|
||||
See ``DSInferenceModelBase.kv_cache_config`` for documentation.
|
||||
|
||||
This method assumes an autoregressive dense attention pattern. Override this method
|
||||
if this does not match the model's attention pattern.
|
||||
"""
|
||||
if self._kv_cache_config is None:
|
||||
cache_shape = (self.num_layers, self.n_heads_kv_local, self.head_size)
|
||||
max_blocks = ceil_div(self.max_sequence_length, self.attn.kv_block_size)
|
||||
self._kv_cache_config = KVCacheConfig(block_size=self.attn.kv_block_size,
|
||||
cache_shape=cache_shape,
|
||||
cache_dtype=self.activation_dtype,
|
||||
max_blocks_per_allocation_group=max_blocks)
|
||||
return (self._kv_cache_config, )
|
||||
|
||||
def prepare_batch(self, wrapped_batch: RaggedBatchWrapper) -> None:
|
||||
"""
|
||||
See ``DSInferenceModelBase.prepare_batch`` for documentation.
|
||||
|
||||
This method assumes an autoregressive dense attention pattern. Override this method
|
||||
if this does not match the model's attention pattern.
|
||||
"""
|
||||
self.attn.build_atoms(wrapped_batch)
|
||||
|
||||
######### Attention output #########
|
||||
def make_attn_out_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the attention output linear layer. This sets the
|
||||
`self.attn_out` attribute.
|
||||
"""
|
||||
in_features = attn_out_in_features(self.model_dim, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q,
|
||||
self.n_heads_kv)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=in_features,
|
||||
out_channels=self.model_dim,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.attn_out = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
def transform_attn_out_param(self, param: torch.Tensor) -> Optional[InferenceParameter]:
|
||||
"""
|
||||
Shards an attention output projection parameter and passes it to the underlying
|
||||
implementation for any necessary transformations. This will return `None` for bias parameters
|
||||
if they are not on TP rank 0.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have
|
||||
the shape (out_neurons, in_neurons).
|
||||
"""
|
||||
param = shard_attn_out_param(param, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q,
|
||||
self.n_heads_kv)
|
||||
|
||||
if param is not None:
|
||||
param = self.attn_out.transform_param(param)
|
||||
|
||||
return param
|
||||
|
||||
######### MLP #########
|
||||
def make_mlp_1_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the first MLP in the feedforward network.
|
||||
This sets the `self.mlp_1` attribute.
|
||||
"""
|
||||
shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=self.model_dim,
|
||||
out_channels=shard_size,
|
||||
activation=self.mlp_activation_fn,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.mlp_1 = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
def transform_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Shards the first MLP parameter and passes it to the underlying implementation
|
||||
for any necessary transformations.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have
|
||||
the shape (out_neurons, in_neurons).
|
||||
"""
|
||||
param = shard_mlp_1_param(param, self.tp_rank, self.tp_size, gated=self.gated_mlp)
|
||||
|
||||
return self.mlp_1.transform_param(param)
|
||||
|
||||
def make_mlp_2_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the second MLP in the feedforward network.
|
||||
This sets the `self.mlp_2` attribute.
|
||||
"""
|
||||
shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=shard_size,
|
||||
out_channels=self.model_dim,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.mlp_2 = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
def transform_mlp_2_param(self, param: torch.Tensor) -> Optional[InferenceParameter]:
|
||||
"""
|
||||
Shards the second MLP parameter and passes it to the underlying implementation
|
||||
for any necessary transformations. This will return `None` for bias parameters
|
||||
if they are not on TP rank 0.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have
|
||||
the shape (out_neurons, in_neurons).
|
||||
"""
|
||||
param = shard_mlp_2_param(param, self.tp_rank, self.tp_size)
|
||||
|
||||
if param is not None:
|
||||
param = self.mlp_2.transform_param(param)
|
||||
|
||||
return param
|
||||
|
||||
######### Norm #########
|
||||
def make_norm_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the normalization layer for the model. This sets the `self.norm` attribute.
|
||||
|
||||
TODO(cmikeh2): In the future we'll distinguish between the different norm objects,
|
||||
but for now we'll just use the same one for all of them.
|
||||
"""
|
||||
norm_config = DSNormConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
type=self.norm_type,
|
||||
channels=self.model_dim,
|
||||
residual_dtype=self.activation_dtype,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config)
|
||||
|
||||
def transform_norm_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Passes a normalization parameter to the underlying implementation for any
|
||||
necessary transformations.
|
||||
|
||||
TODO(cmikeh2): In the future we'll distinguish between the different norm objects,
|
||||
but for now we'll just use the same one for all of them.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have
|
||||
shape (model_dim,)
|
||||
"""
|
||||
return self.norm.transform_param(param)
|
||||
|
||||
|
||||
class DSMoETransformerModelBase(DSTransformerModelBase):
|
||||
|
||||
@property
|
||||
def n_experts(self) -> int:
|
||||
"""
|
||||
Return the number of experts in the model.
|
||||
"""
|
||||
raise NotImplementedError("Attempted to access an unimplemented number of experts")
|
||||
|
||||
@property
|
||||
def n_top_k(self) -> int:
|
||||
"""
|
||||
Number of experts per token.
|
||||
"""
|
||||
raise NotImplementedError("Attempted to access an unimplemented number of experts per token")
|
||||
|
||||
@property
|
||||
def normalize_expert_scores(self) -> bool:
|
||||
"""
|
||||
Whether to normalize expert scores. If true, sum(expert_scores) = 1.
|
||||
"""
|
||||
raise NotImplementedError("Attempted to access an unimplemented normalization flag")
|
||||
|
||||
def make_moe_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the MoE layer for the model. This sets the `self.moe` attribute.
|
||||
"""
|
||||
sharded_dim = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
moe_config = DSMoEConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
model_dim=self.model_dim,
|
||||
intermediate_features=sharded_dim,
|
||||
activation=self.mlp_activation_fn,
|
||||
n_experts=self.n_experts,
|
||||
top_k=self.n_top_k,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
normalize_scores=self.normalize_expert_scores,
|
||||
)
|
||||
|
||||
self.moe = heuristics.instantiate_moe(moe_config, self._engine_config)
|
||||
|
||||
def transform_moe_gate_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Passes a MoE gate parameter to the underlying implementation for any necessary transformations.
|
||||
|
||||
TODO(cmikeh2): This will need to be updated/overridden for expert parallelism.
|
||||
"""
|
||||
return self.moe.transform_gate_param(param)
|
||||
|
||||
def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Shards the first MoE param and passes it to the underlying implementation. Since it's possible for an architecture
|
||||
to have both MoE and non-MoE layers, this can't be overloaded on the MLP1 transform. Furthermore, since both
|
||||
the MoE DSModule owns both MLP1 and MLP2, under certain sharding conditions it's not possible for the model implementation
|
||||
to infer from the shape whether to perform a different transformation based on MLP1 or MLP2. This (and the below)
|
||||
separations are intended to solve both these issues.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons).
|
||||
"""
|
||||
param = shard_mlp_1_param(param, self.tp_rank, self.tp_size, gated=self.gated_mlp, is_moe=True)
|
||||
|
||||
return self.moe.transform_moe_mlp_1_param(param)
|
||||
|
||||
def transform_moe_mlp_2_param(self, param: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Shards the second MoE param and passes it to the underlying implementation. See the above for context on why this API
|
||||
exists.
|
||||
|
||||
This will return `None` for expert bias params not on TP rank 0. NOTE(cmikeh2): Does it make sense to round-robin assign?
|
||||
My intuition is that this will make debugging much more difficult for minimal memory reduction.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons).
|
||||
"""
|
||||
param = shard_mlp_2_param(param, self.tp_rank, self.tp_size, is_moe=True)
|
||||
|
||||
if param is not None:
|
||||
param = self.moe.transform_moe_mlp_2_param(param)
|
||||
|
||||
return param
|
||||
@@ -0,0 +1,356 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import re
|
||||
from typing import Type
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.compat import get_annotations_from_namespace, get_annotations
|
||||
from .parameter_base import ParameterBase, ParametrizedList
|
||||
from ..inference_parameter import InferenceParameter
|
||||
|
||||
# Currently have dependency loops for the type hints.
|
||||
InferenceModel = Type["InferenceModel"]
|
||||
LayerContainer = Type["LayerContainer"] # noqa: F811
|
||||
|
||||
MAPPING_KEY = "PARAM_MAPPING"
|
||||
PLIST_HELPERS = "_ds_plist_strip_vals"
|
||||
|
||||
|
||||
def make_finalization_callback(all_names: str):
|
||||
"""
|
||||
Helper method for building the finalization callback for a LayerContainer. This
|
||||
is not client code and should not be used or called directly.
|
||||
"""
|
||||
|
||||
def finalization_callback(self, param: ParameterBase, finalized_param: torch.Tensor) -> None:
|
||||
"""
|
||||
Callback for when a parameter is finalized.
|
||||
"""
|
||||
self._finalized_params += 1
|
||||
|
||||
for name in all_names:
|
||||
if getattr(self, name) is param:
|
||||
setattr(self, name, finalized_param)
|
||||
|
||||
return finalization_callback
|
||||
|
||||
|
||||
class LayerMetaclass(type):
|
||||
"""
|
||||
MetaClass for the LayerContainer base class. This class will parse the annotations
|
||||
of the class that correspond to `ParameterBase` and create None initializers for each
|
||||
as well as a finalization callback that for when each `ParameterBase` is finalized
|
||||
and should be replaced with a Tensor.
|
||||
"""
|
||||
|
||||
def __new__(cls, clsname, bases, attrs):
|
||||
|
||||
annotations = get_annotations_from_namespace(attrs)
|
||||
|
||||
for base in bases:
|
||||
# We'll pick up all annotations on any base classes. This will allow us to
|
||||
# to use inheritance to share common parameter groups in base classes.
|
||||
annotations.update(get_annotations(base))
|
||||
|
||||
if hasattr(base, MAPPING_KEY):
|
||||
if MAPPING_KEY not in attrs:
|
||||
# This is likely a fail state. If a parent has MAPPING KEY but the child does
|
||||
# not, then we're guaranteed only a subset of the parameters will be mapped.
|
||||
attrs[MAPPING_KEY] = {}
|
||||
attrs[MAPPING_KEY].update(getattr(base, MAPPING_KEY))
|
||||
|
||||
all_names = [name for name, annotation in annotations.items() if issubclass(annotation, ParameterBase)]
|
||||
|
||||
if MAPPING_KEY in attrs:
|
||||
# If we have a mapping key at all, then we will enter the validation mode for building
|
||||
# helpers for mapping and ensuring we have complete mapping.
|
||||
|
||||
# First we'll build a flat list of every dependency for this layer.
|
||||
all_deps = set()
|
||||
for name in all_names:
|
||||
parameter_deps = [
|
||||
name for name, annotation in get_annotations(annotations[name]).items()
|
||||
if issubclass(annotation, (torch.Tensor, ParametrizedList))
|
||||
]
|
||||
|
||||
all_deps.update([f"{name}.{dep}" for dep in parameter_deps])
|
||||
|
||||
# Create static helper for doing the string processing only once.
|
||||
attrs[PLIST_HELPERS] = []
|
||||
|
||||
# Iterate over all the mappings
|
||||
for src_name, target_or_targets in attrs[MAPPING_KEY].items():
|
||||
if isinstance(target_or_targets, str):
|
||||
target_or_targets = [target_or_targets]
|
||||
|
||||
actual_targets = []
|
||||
for target_name in target_or_targets:
|
||||
base_dependency, dependency_attr = target_name.split(".")
|
||||
|
||||
# Check for invalid mappings
|
||||
if base_dependency not in all_names:
|
||||
raise ValueError(
|
||||
"Target parameter \"{}\" not found in this layer. Valid targets are {}".format(
|
||||
base_dependency, all_names))
|
||||
if dependency_attr not in get_annotations(annotations[base_dependency]):
|
||||
# This check is not universal (see below) if a single dependency is being
|
||||
# mapped to by a single row.
|
||||
raise ValueError(
|
||||
"Target dependency \"{}\" not found on parameter \"{}\". Valid targets are {}".format(
|
||||
dependency_attr, base_dependency,
|
||||
get_annotations(annotations[base_dependency]).keys()))
|
||||
if target_name not in all_deps:
|
||||
raise ValueError(
|
||||
"Target dependency \"{}\" was targeted with multiple mapping rules.".format(target_name))
|
||||
|
||||
# If we've made it this far, the dependency definitely exists.
|
||||
actual_targets.append(get_annotations(annotations[base_dependency])[dependency_attr])
|
||||
|
||||
all_deps.remove(target_name)
|
||||
|
||||
are_plists = [issubclass(target, ParametrizedList) for target in actual_targets]
|
||||
if all(are_plists):
|
||||
# We can do direct sets on everything but ParametrizedLists, so we'll only explicitly
|
||||
# handle these here.
|
||||
# TODO(cmikeh2): SPLIT, error if more than 1
|
||||
glob_count = src_name.count("*")
|
||||
if glob_count > 1:
|
||||
raise ValueError(
|
||||
"ParametrizedList index inference can only work with a single glob: {}".format(src_name))
|
||||
elif glob_count == 0:
|
||||
raise ValueError(
|
||||
"Must have wildcard (*) in source name for ParametrizedList mapping: {}".format(src_name))
|
||||
|
||||
wildcard_idx = src_name.find("*")
|
||||
prefix = src_name[:wildcard_idx]
|
||||
suffix = src_name[wildcard_idx + 1:]
|
||||
attrs[PLIST_HELPERS].append((prefix, suffix, target_or_targets))
|
||||
elif any(are_plists):
|
||||
raise ValueError("Cannot mix ParametrizedLists and Tensors in a single mapping rule.")
|
||||
|
||||
if len(all_deps) > 0:
|
||||
raise ValueError(
|
||||
"A parameter mapping was provided for {}, but the following dependencies were not mapped: {}".
|
||||
format(clsname, all_deps))
|
||||
|
||||
attrs["finalization_callback"] = make_finalization_callback(all_names)
|
||||
|
||||
new_obj = super().__new__(cls, clsname, bases, attrs)
|
||||
|
||||
setattr(new_obj, "_n_params", len(all_names))
|
||||
setattr(new_obj, "_annotation_attrs", all_names)
|
||||
|
||||
return new_obj
|
||||
|
||||
def __call__(cls, *args, **kwargs):
|
||||
instance = cls.__new__(cls, *args, **kwargs)
|
||||
instance.__init__(*args, **kwargs)
|
||||
|
||||
for name, annotation in get_annotations(instance).items():
|
||||
if issubclass(annotation, ParameterBase):
|
||||
# TODO(cmikeh2): Do we want to make this a property
|
||||
# It might also make sense to do this in the base class __init__
|
||||
# but since it is tied with the changes made in __new__ it feels
|
||||
# to me like it should be here.
|
||||
setattr(instance, name, annotation(instance.inference_model, instance))
|
||||
|
||||
return instance
|
||||
|
||||
|
||||
class LayerContainer(metaclass=LayerMetaclass): # noqa: F811
|
||||
"""
|
||||
Abstract base class for containing model parameters.
|
||||
|
||||
This is primarily a guidance abstraction since we do not put any restrictions
|
||||
on how the parameters are stored.
|
||||
|
||||
To use this class, annotate the class with `ParameterBase` subclasses and give them
|
||||
names. As a checkpoint is loaded into this container, the `ParameterBase` instances
|
||||
will be replaced with realized Tensors as soon as each of their dependencies are met.
|
||||
|
||||
To enable automatic mapping, add a static attribute `PARAM_MAPPING` to the class
|
||||
definition. This should be a dictionary mapping from a source string to one or
|
||||
more dependencies.
|
||||
|
||||
```python
|
||||
class MyLayer(LayerContainer):
|
||||
PARAM_MAPPING = {
|
||||
"path.to.param.dependency", "container_param_1.dependency",
|
||||
"path.to.param2.dependency", "container_param_2.dependency",
|
||||
"path.to.param3.*.dependency", "container_param_3.list_dependency"
|
||||
}
|
||||
|
||||
...
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, model: InferenceModel) -> None:
|
||||
"""
|
||||
Initialization of the LayerContainer. This method does not need to be overridden
|
||||
for any children classes.
|
||||
|
||||
Args:
|
||||
model (InferenceModel): Inference model that will be used to shard and transform
|
||||
parameters correctly, as well as provide specific information about the model
|
||||
for `ParameterizedList`s that may be part of one of the member `ParameterBase`s.
|
||||
"""
|
||||
self.inference_model = model
|
||||
self._finalized_params = 0
|
||||
|
||||
def _initialization_checker(self, check_device: bool = True) -> bool:
|
||||
"""
|
||||
Returns whether or not all parameters have been initialized and transformed by
|
||||
the model. Once this returns True, all the `ParameterBase` instances will be
|
||||
torch.Tensors.
|
||||
"""
|
||||
if self._finalized_params != self.n_params:
|
||||
return False
|
||||
|
||||
for name in self._annotation_attrs:
|
||||
tensor = getattr(self, name)
|
||||
if tensor is None:
|
||||
continue
|
||||
elif not isinstance(tensor, InferenceParameter):
|
||||
raise ValueError("Layer should be finalized, but {} ({}) is neither InferenceParameter or None".format(
|
||||
name, type(tensor)))
|
||||
elif check_device and tensor.device != torch.device(get_accelerator().current_device()):
|
||||
raise RuntimeError("Layer should be finalized, but {} is not on device {}".format(
|
||||
name,
|
||||
get_accelerator().current_device()))
|
||||
return True
|
||||
|
||||
@property
|
||||
def is_populated(self) -> bool:
|
||||
"""
|
||||
Returns whether or not all parameters have been populated by the checkpoint engine, but
|
||||
does not validat the parameters are on the correct device.
|
||||
"""
|
||||
return self._initialization_checker(check_device=False)
|
||||
|
||||
@property
|
||||
def is_initialized(self) -> bool:
|
||||
"""
|
||||
Returns whether or not all parameters have been initialized and transformed by
|
||||
the model and are located on the appropriate device. Once this returns True, all
|
||||
the `ParameterBase` instances ``InferenceParameter``s or explicitly set to ``None``.
|
||||
"""
|
||||
return self._initialization_checker()
|
||||
|
||||
@property
|
||||
def n_params(self) -> int:
|
||||
"""
|
||||
The number of parameters this container holds. This is a read-only value
|
||||
that is set by the metaclass.
|
||||
"""
|
||||
return self._n_params
|
||||
|
||||
@property
|
||||
def annotation_attrs(self) -> list:
|
||||
return self._annotation_attrs
|
||||
|
||||
@property
|
||||
def mapping_params(self) -> dict:
|
||||
return getattr(self.__class__, MAPPING_KEY, {})
|
||||
|
||||
@property
|
||||
def plist_helpers(self) -> list:
|
||||
return getattr(self.__class__, PLIST_HELPERS, [])
|
||||
|
||||
def direct_injection(self, name: str, tensor: InferenceParameter) -> None:
|
||||
|
||||
if name not in self._annotation_attrs:
|
||||
raise ValueError(f"Cannot directly inject {name}, not a valid parameter.")
|
||||
|
||||
setattr(self, name, tensor)
|
||||
self._finalized_params += 1
|
||||
|
||||
def set_dependency(self, dep_name: str, dep_value: torch.Tensor) -> None:
|
||||
"""
|
||||
Set dependency can be used for managing dependencies when a mapping is provided
|
||||
in the class definition for the layer. The dep_name here should have any prefix
|
||||
for transformer layers removed (such as model.layers.*.attn.qkv.weight -> attn.qkv.weight).
|
||||
|
||||
Args:
|
||||
dep_name (str): The name of the dependency to set.
|
||||
dep_value (torch.Tensor): The value to set the dependency to.
|
||||
"""
|
||||
|
||||
def get_dep_name_target(dep_name: str) -> str:
|
||||
"""
|
||||
Helper method for getting the target name for a dependency from the
|
||||
mapping params. Tries to match exact string first, then looks for
|
||||
wildcards and attempts regex matching. Will return empty string if
|
||||
no match found.
|
||||
"""
|
||||
if dep_name in self.mapping_params:
|
||||
# If we have an exact match, it's a direct mapping and we can
|
||||
# immediately set the value.
|
||||
return self.mapping_params[dep_name]
|
||||
|
||||
matched_targets = []
|
||||
for key, target in self.mapping_params.items():
|
||||
regex_key = key.replace("*", ".*")
|
||||
if re.match(regex_key, dep_name):
|
||||
matched_targets.append(target)
|
||||
if len(matched_targets) > 1:
|
||||
raise ValueError(f"Multiple targets matched for dependency {dep_name}: {matched_targets}")
|
||||
if matched_targets:
|
||||
return matched_targets[0]
|
||||
return ""
|
||||
|
||||
if dep_name in self.mapping_params:
|
||||
# If we have an exact match, it's a direct mapping and we can immediately set
|
||||
# the value.
|
||||
target = self.mapping_params[dep_name]
|
||||
|
||||
# Convert single targets to a list for consistency
|
||||
if isinstance(target, str):
|
||||
target = [target]
|
||||
|
||||
for target_name in target:
|
||||
# Double setting doesn't set the attribute correctly, so we do a getattr then setattr
|
||||
target_param_name, target_dependency_name = target_name.split(".")
|
||||
target_param = getattr(self, target_param_name)
|
||||
setattr(target_param, target_dependency_name, dep_value)
|
||||
return
|
||||
|
||||
# Otherwise we need to map to one of the parameter lists.
|
||||
for prefix, suffix, dests in self.plist_helpers:
|
||||
if dep_name.startswith(prefix) and dep_name.endswith(suffix):
|
||||
# We have a match, so we can set the value.
|
||||
target_idx = int(dep_name[len(prefix):-len(suffix)])
|
||||
|
||||
# Convert single targets to a list for consistency
|
||||
if isinstance(dests, str):
|
||||
dests = [dests]
|
||||
|
||||
for dest in dests:
|
||||
target_param_name, target_dependency_name = dest.split(".")
|
||||
target_param = getattr(self, target_param_name)
|
||||
target_dependency = getattr(target_param, target_dependency_name)
|
||||
target_dependency[target_idx] = dep_value
|
||||
return
|
||||
|
||||
# TODO: Refactor this with the help of cmikeh2
|
||||
# We should be able to combine this with the wildcard matching above.
|
||||
target = get_dep_name_target(dep_name)
|
||||
if target:
|
||||
# Convert single targets to a list for consistency
|
||||
if isinstance(target, str):
|
||||
target = [target]
|
||||
|
||||
for target_name in target:
|
||||
# Double setting doesn't set the attribute correctly, so we do a getattr then setattr
|
||||
target_param_name, target_dependency_name = target_name.split(".")
|
||||
target_param = getattr(self, target_param_name)
|
||||
setattr(target_param, target_dependency_name, dep_value)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
"Could not find a mapping for dependency \"{}\". Check that it is included in the ``MAPPING_PARAMS``. See docstring for more on ``MAPPING_PARAMS``"
|
||||
.format(dep_name))
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import Llama2Policy
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Llama model looks like this:
|
||||
|
||||
LlamaForCausalLM(
|
||||
(model): LlamaModel(
|
||||
(embed_tokens): Embedding(32000, 4096, padding_idx=0)
|
||||
(layers): ModuleList(
|
||||
(0-31): 32 x LlamaDecoderLayer(
|
||||
(self_attn): LlamaAttention(
|
||||
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(rotary_emb): LlamaRotaryEmbedding()
|
||||
)
|
||||
(mlp): LlamaMLP(
|
||||
(gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
|
||||
(up_proj): Linear(in_features=4096, out_features=11008, bias=False)
|
||||
(down_proj): Linear(in_features=11008, out_features=4096, bias=False)
|
||||
(act_fn): SiLUActivation()
|
||||
)
|
||||
(input_layernorm): LlamaRMSNorm()
|
||||
(post_attention_layernorm): LlamaRMSNorm()
|
||||
)
|
||||
)
|
||||
(norm): LlamaRMSNorm()
|
||||
)
|
||||
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class Llama2TransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Llama-2 model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: GatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate_proj.weight": "mlp_1_w.gate_params",
|
||||
"mlp.up_proj.weight": "mlp_1_w.up_params",
|
||||
"mlp.down_proj.weight": "mlp_2_w.params",
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class Llama2NonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Llama-2 model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .container import Llama2NonTransformerContainer, Llama2TransformerContainer
|
||||
|
||||
|
||||
class Llama2InferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[Llama2NonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[Llama2TransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
activation = self._config.hidden_act.lower()
|
||||
# llama model family is special and is always gated so force gated versions of relu, gelu, silu
|
||||
if activation == "gelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "relu":
|
||||
return ActivationType.ReGLU
|
||||
elif activation == "gegelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "silu":
|
||||
return ActivationType.SiGLU
|
||||
else:
|
||||
raise NotImplementedError(f"Activation {activation} not supported")
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
# Should be configurable in the future
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import Llama2NonTransformerContainer, Llama2TransformerContainer
|
||||
from .model import Llama2InferenceModel
|
||||
|
||||
|
||||
class Llama2Policy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Llama2InferenceModel:
|
||||
return Llama2InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [Llama2TransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(Llama2NonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(
|
||||
[f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import MistralPolicy
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from deepspeed.inference.v2.model_implementations.common_parameters import *
|
||||
from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Mistral model (mistralai/Mistral-7B-v0.1) looks like this:
|
||||
MistralForCausalLM(
|
||||
(model): MistralModel(
|
||||
(embed_tokens): Embedding(32000, 4096)
|
||||
(layers): ModuleList(
|
||||
(0-31): 32 x MistralDecoderLayer(
|
||||
(self_attn): MistralAttention(
|
||||
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
|
||||
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
|
||||
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(rotary_emb): MistralRotaryEmbedding()
|
||||
)
|
||||
(mlp): MistralMLP(
|
||||
(gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
|
||||
(up_proj): Linear(in_features=4096, out_features=14336, bias=False)
|
||||
(down_proj): Linear(in_features=14336, out_features=4096, bias=False)
|
||||
(act_fn): SiLUActivation()
|
||||
)
|
||||
(input_layernorm): MistralRMSNorm()
|
||||
(post_attention_layernorm): MistralRMSNorm()
|
||||
)
|
||||
)
|
||||
(norm): MistralRMSNorm()
|
||||
)
|
||||
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class MistralTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Mistral model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: GatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate_proj.weight": "mlp_1_w.gate_params",
|
||||
"mlp.up_proj.weight": "mlp_1_w.up_params",
|
||||
"mlp.down_proj.weight": "mlp_2_w.params",
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class MistralNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Mistral model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from ...model_implementations import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .container import MistralNonTransformerContainer, MistralTransformerContainer
|
||||
|
||||
|
||||
class MistralInferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Mistral models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[MistralNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[MistralTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
activation = self._config.hidden_act.lower()
|
||||
if activation == "gelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "relu":
|
||||
return ActivationType.ReGLU
|
||||
elif activation == "gegelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "silu":
|
||||
return ActivationType.SiGLU
|
||||
else:
|
||||
raise NotImplementedError(f"Activation {activation} not supported")
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
# Should be configurable in the future
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import MistralNonTransformerContainer, MistralTransformerContainer
|
||||
from .model import MistralInferenceModel
|
||||
|
||||
|
||||
class MistralPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> MistralInferenceModel:
|
||||
return MistralInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [MistralTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(MistralNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params([])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import MixtralPolicy
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from deepspeed.inference.v2.model_implementations.common_parameters import *
|
||||
from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
|
||||
|
||||
|
||||
class MixtralTransformerContainer(LayerContainer):
|
||||
|
||||
qkv_w: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
moe_gate: MoEGatingWeightParameter
|
||||
moe_mlp_1: UnfusedMoEGatedMLPParameter
|
||||
moe_mlp_2: UnfusedMoEMLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"block_sparse_moe.gate.weight": "moe_gate.params",
|
||||
"block_sparse_moe.experts.*.w1.weight": "moe_mlp_1.gating_experts",
|
||||
"block_sparse_moe.experts.*.w3.weight": "moe_mlp_1.up_experts",
|
||||
"block_sparse_moe.experts.*.w2.weight": "moe_mlp_2.experts",
|
||||
}
|
||||
|
||||
|
||||
class MixtralNonTransformerContainer(LayerContainer):
|
||||
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from ...model_implementations import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
from ..inference_model_base import (
|
||||
DSModelImplementationConfig,
|
||||
MPType,
|
||||
)
|
||||
|
||||
from .container import MixtralNonTransformerContainer, MixtralTransformerContainer
|
||||
|
||||
|
||||
class MixtralInferenceModel(DSMoETransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for Mixtral models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[MixtralNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[MixtralTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_position_embeddings
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
activation = self._config.hidden_act.lower()
|
||||
if activation == "gelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "relu":
|
||||
return ActivationType.ReGLU
|
||||
elif activation == "gegelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "silu":
|
||||
return ActivationType.SiGLU
|
||||
else:
|
||||
raise NotImplementedError(f"Activation {activation} not supported")
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
"""
|
||||
The positional embedding configuration for the model.
|
||||
"""
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Inherited from `DSMoETransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def n_experts(self) -> int:
|
||||
return self._config.num_local_experts
|
||||
|
||||
@property
|
||||
def n_top_k(self) -> int:
|
||||
return self._config.num_experts_per_tok
|
||||
|
||||
@property
|
||||
def normalize_expert_scores(self) -> bool:
|
||||
return True
|
||||
|
||||
"""
|
||||
Model implementation
|
||||
"""
|
||||
|
||||
def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig,
|
||||
base_mp_group: MPType) -> None:
|
||||
"""
|
||||
Base implementation for initialization. By default, this will initialize
|
||||
the traditional components of a transformer model:
|
||||
- Embedding
|
||||
- QKV projection
|
||||
- Self attention
|
||||
- Attention output projection
|
||||
- Feed forward network
|
||||
- Normalization
|
||||
- Unembedding
|
||||
|
||||
Arguments:
|
||||
config (DSModelImplementationConfig): Model-specific configuration. No assumptions
|
||||
should be made about this config that are not closely tied to the specific
|
||||
model implementation.
|
||||
engine_config (RaggedInferenceEngineConfig): Engine configuration.
|
||||
base_mp_group (MPType): Base communication group for Tensor-parallel inference.
|
||||
"""
|
||||
super().__init__(config, engine_config, base_mp_group)
|
||||
|
||||
self.make_norm_layer()
|
||||
self.make_qkv_layer()
|
||||
self.make_attn_layer()
|
||||
self.make_attn_out_layer()
|
||||
self.make_moe_layer()
|
||||
self.make_embedding_layer()
|
||||
self.make_unembedding_layer()
|
||||
self._kv_cache_config = None
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma)
|
||||
|
||||
hidden_states = self.moe(hidden_states, ragged_batch_info, cur_params.moe_gate, cur_params.moe_mlp_1,
|
||||
cur_params.moe_mlp_2)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import MixtralTransformerContainer, MixtralNonTransformerContainer
|
||||
from .model import MixtralInferenceModel
|
||||
|
||||
|
||||
class MixtralPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> MixtralInferenceModel:
|
||||
return MixtralInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [MixtralTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(MixtralNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params([])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import OPTPolicy
|
||||
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF OPT model looks like this:
|
||||
|
||||
OPTForCausalLM(
|
||||
(model): OPTModel(
|
||||
(decoder): OPTDecoder(
|
||||
(embed_tokens): Embedding(50272, 768, padding_idx=1)
|
||||
(embed_positions): OPTLearnedPositionalEmbedding(2050, 768)
|
||||
(final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
|
||||
(layers): ModuleList(
|
||||
(0-11): 12 x OPTDecoderLayer(
|
||||
(self_attn): OPTAttention(
|
||||
(k_proj): Linear(in_features=768, out_features=768, bias=True)
|
||||
(v_proj): Linear(in_features=768, out_features=768, bias=True)
|
||||
(q_proj): Linear(in_features=768, out_features=768, bias=True)
|
||||
(out_proj): Linear(in_features=768, out_features=768, bias=True)
|
||||
)
|
||||
(activation_fn): ReLU()
|
||||
(self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
|
||||
(fc1): Linear(in_features=768, out_features=3072, bias=True)
|
||||
(fc2): Linear(in_features=3072, out_features=768, bias=True)
|
||||
(final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
(lm_head): Linear(in_features=768, out_features=50272, bias=False)
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
|
||||
class OPTTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the OPT model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
qkv_b: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
attn_out_b: AttentionOutputParameter
|
||||
mlp_1_w: MLP1Parameter
|
||||
mlp_1_b: MLP1Parameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
mlp_2_b: MLP2Parameter
|
||||
attn_norm_beta: NormParameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_beta: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.q_proj.bias": "qkv_b.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.k_proj.bias": "qkv_b.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.v_proj.bias": "qkv_b.v_params",
|
||||
"self_attn.out_proj.weight": "attn_out_w.params",
|
||||
"self_attn.out_proj.bias": "attn_out_b.params",
|
||||
"fc1.weight": "mlp_1_w.params",
|
||||
"fc1.bias": "mlp_1_b.params",
|
||||
"fc2.weight": "mlp_2_w.params",
|
||||
"fc2.bias": "mlp_2_b.params",
|
||||
"self_attn_layer_norm.weight": "attn_norm_gamma.params",
|
||||
"self_attn_layer_norm.bias": "attn_norm_beta.params",
|
||||
"final_layer_norm.weight": "mlp_norm_gamma.params",
|
||||
"final_layer_norm.bias": "mlp_norm_beta.params",
|
||||
}
|
||||
|
||||
|
||||
class OPTNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the OPT model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_emb_pos: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm_w: NormParameter
|
||||
final_norm_b: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"*decoder.embed_tokens.weight": ["word_emb.params", "word_unembed.params"],
|
||||
"*decoder.embed_positions.weight": "word_emb_pos.params",
|
||||
"*decoder.final_layer_norm.weight": "final_norm_w.params",
|
||||
"*decoder.final_layer_norm.bias": "final_norm_b.params",
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from ...model_implementations import *
|
||||
from ...modules.configs import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
from .container import OPTNonTransformerContainer, OPTTransformerContainer
|
||||
|
||||
from ...modules.heuristics import instantiate_embed
|
||||
|
||||
|
||||
class OPTInferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for OPT models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[OPTNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[OPTTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.ffn_dim
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.RELU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.LayerNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.none
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return None
|
||||
|
||||
"""
|
||||
Overrides of ``DSTransformerModelBase`` methods
|
||||
"""
|
||||
|
||||
def make_embedding_layer(self) -> None:
|
||||
"""
|
||||
Performs setup and creates embedding DSModule. Since OPT includes trained
|
||||
positional embeddings, we will override the base model implementation.
|
||||
"""
|
||||
|
||||
embed_config = DSEmbeddingsConfig(max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
residual_dtype=self.activation_dtype,
|
||||
embedding_dim=self.model_dim,
|
||||
positional_embedding=True,
|
||||
positional_offset=2)
|
||||
|
||||
self.embed = instantiate_embed(embed_config, self._engine_config)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb, self._non_transformer.word_emb_pos)
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=cur_params.attn_out_b)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual,
|
||||
hidden_states,
|
||||
cur_params.mlp_norm_gamma,
|
||||
beta=cur_params.mlp_norm_beta)
|
||||
|
||||
# Should be configurable in the future
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=cur_params.mlp_1_b)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=cur_params.mlp_2_b)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual,
|
||||
hidden_states,
|
||||
next_params.attn_norm_gamma,
|
||||
beta=next_params.attn_norm_beta)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm_w,
|
||||
beta=self._non_transformer.final_norm_b)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual,
|
||||
None,
|
||||
self._transformer[0].attn_norm_gamma,
|
||||
beta=self._transformer[0].attn_norm_beta)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import OPTNonTransformerContainer, OPTTransformerContainer
|
||||
from .model import OPTInferenceModel
|
||||
|
||||
|
||||
class OPTPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> OPTInferenceModel:
|
||||
return OPTInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [OPTTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.decoder.layers', 'decoder.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(OPTNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(['lm_head.weight'])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import weakref
|
||||
from abc import abstractmethod
|
||||
from typing import Type
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.compat import get_annotations_from_namespace, get_annotations
|
||||
|
||||
# Currently have dependency loops for the type hints.
|
||||
InferenceModel = Type["InferenceModel"]
|
||||
LayerContainer = Type["LayerContainer"]
|
||||
|
||||
MAPPING_KEY = "PARAM_MAPPING"
|
||||
|
||||
|
||||
def make_param_getter(clsname, param):
|
||||
"""
|
||||
Normal getter implementation for a property.
|
||||
"""
|
||||
|
||||
def param_getter(self):
|
||||
return getattr(self, f"__{clsname}__{param}")
|
||||
|
||||
return param_getter
|
||||
|
||||
|
||||
def make_param_setter(clsname, param):
|
||||
"""
|
||||
Setter implementation that will call complete component to potentially
|
||||
finalize the parameter.
|
||||
"""
|
||||
|
||||
def param_setter(self, value):
|
||||
setattr(self, f"__{clsname}__{param}", value)
|
||||
self.dtype = value.dtype
|
||||
self.complete_component()
|
||||
|
||||
return param_setter
|
||||
|
||||
|
||||
def make_readonly_setter():
|
||||
"""
|
||||
Setter implementation that will raise an error if called.
|
||||
"""
|
||||
|
||||
def paramlist_setter(self, value):
|
||||
raise ValueError("Cannot set a ParametrizedList directly.")
|
||||
|
||||
return paramlist_setter
|
||||
|
||||
|
||||
class ParameterMetaclass(type):
|
||||
"""
|
||||
MetaClass for the ParameterBase base class. This class will parse the `src_params`
|
||||
attribute and create properties for each of the dependencies. A dependency can either
|
||||
be represented as a string, which is interpreted as a named Tensor, or a `ParametrizedList`
|
||||
subclass.
|
||||
"""
|
||||
|
||||
def __new__(cls, clsname, bases, attrs):
|
||||
|
||||
annotations = get_annotations_from_namespace(attrs)
|
||||
dependencies = {
|
||||
name: annotation
|
||||
for name, annotation in annotations.items() if issubclass(annotation, (torch.Tensor, ParametrizedList))
|
||||
}
|
||||
n_dependencies = len(dependencies)
|
||||
|
||||
# Create properties for each of our dependencies
|
||||
for d_name, d_type in dependencies.items():
|
||||
if issubclass(d_type, ParametrizedList):
|
||||
assert hasattr(
|
||||
d_type, "count_attr"
|
||||
), "ParametrizedList must have a count_attr attribute to access on the inference module."
|
||||
attrs[d_name] = property(make_param_getter(clsname, d_name), make_readonly_setter())
|
||||
else: # torch.Tensor
|
||||
attrs[d_name] = property(make_param_getter(clsname, d_name), make_param_setter(clsname, d_name))
|
||||
|
||||
new_cls = super().__new__(cls, clsname, bases, attrs)
|
||||
new_cls.n_dependencies = n_dependencies
|
||||
|
||||
return new_cls
|
||||
|
||||
def __call__(cls, *args, **kwargs):
|
||||
new_obj = super().__call__(*args, **kwargs)
|
||||
new_obj.__init__(*args, **kwargs)
|
||||
|
||||
setattr(new_obj, "dest_param", None)
|
||||
|
||||
# Initialize our dependences to None/empty `ParametrizedList`s
|
||||
for name, annotation in get_annotations(new_obj).items():
|
||||
if issubclass(annotation, ParametrizedList):
|
||||
#TODO(jeff): update assert with this, model implementation attribute does not align or missing wrt the ParametrizedList attributes
|
||||
assert hasattr(
|
||||
new_obj.inference_model, annotation.count_attr
|
||||
), f"new_obj={new_obj.__class__.__name__}, name={name}, annotation.count_attr={annotation.count_attr}"
|
||||
param_list = annotation(new_obj, getattr(new_obj.inference_model, annotation.count_attr))
|
||||
setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", param_list)
|
||||
else: # torch.Tensor
|
||||
setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", None)
|
||||
|
||||
return new_obj
|
||||
|
||||
|
||||
class ParameterBase(metaclass=ParameterMetaclass):
|
||||
"""
|
||||
A ParameterBase allows us to consolidate tracking the dependencies of loading a parameter from
|
||||
a checkpoint into a single object. This class should not be used directly, but rather subclassed
|
||||
and the `src_params` attribute set to a list of strings and/or `ParametrizedList`s.
|
||||
"""
|
||||
|
||||
# inference_model: InferenceModel
|
||||
"""
|
||||
Inference model that will provide context on how to shard and transform the parameter.
|
||||
"""
|
||||
|
||||
#completed_components: int
|
||||
"""
|
||||
How many of the layer dependencies have been met. This is used to determine when the parameter
|
||||
is ready to be finalized. A ParametrizedList counts as a single dependency for the purposes
|
||||
of this counter.
|
||||
"""
|
||||
|
||||
def __init__(self, model: InferenceModel, parent_container: LayerContainer) -> None:
|
||||
"""
|
||||
Direct constructor. This should not be called from client code.
|
||||
|
||||
Args:
|
||||
model (InferenceModel): Inference model that will be used to shard and transform the
|
||||
parameter in `finalize`.
|
||||
parent_container (LayerContainer): The parent container that this parameter is a member
|
||||
of. We will build a weakref to this container to call the finalization callback.
|
||||
"""
|
||||
self.inference_model = model
|
||||
self.completed_components = 0
|
||||
self.parent_container = weakref.ref(parent_container)
|
||||
|
||||
@abstractmethod
|
||||
def finalize(self) -> torch.Tensor:
|
||||
"""
|
||||
Finalize the parameter after all of its source parameters have been set. This method
|
||||
will be automatically called when all inputs have been set. It should return the Tensor
|
||||
with all transformations performed on it.
|
||||
"""
|
||||
pass
|
||||
|
||||
def complete_component(self) -> None:
|
||||
"""
|
||||
Mark a component as completed. This should be called by the relevant setter of a direct
|
||||
property or a ParametrizedList. This method will automatically call `finalize` when all
|
||||
dependencies have been met and then call the finalization callback on the parent container.
|
||||
|
||||
Once the finalization callback has been called, the parameter will be replaced with the
|
||||
`dst_param` attribute on the parent container, and this instance will be destroyed.
|
||||
"""
|
||||
self.completed_components += 1
|
||||
|
||||
if self.completed_components != self.n_dependencies:
|
||||
return
|
||||
|
||||
finalized_param = self.finalize()
|
||||
self.parent_container().finalization_callback(self, finalized_param)
|
||||
|
||||
|
||||
class ParametrizedList:
|
||||
"""
|
||||
A ParametrizedList is a list of parameters that are dependencies
|
||||
of a `ParameterBase` but may vary in length depending on the model
|
||||
configuration (rather than architecture). For example, a MoE layer
|
||||
may have different number of experts depending on the size of the model.
|
||||
|
||||
This class is used to manage these lists and provide integer indexing
|
||||
of a single component rather than accessing names directly. For example,
|
||||
it tends to be more natural to access the 8th expert with `experts[8]`
|
||||
rather than a name like `expert_8`, especially as an attribute.
|
||||
|
||||
To inherit from this class, set static variables `name` and `count_attr`.
|
||||
|
||||
```python
|
||||
class MyParametrizedList(ParametrizedList):
|
||||
count_attr: str = "my_list_count"
|
||||
```
|
||||
|
||||
In the above example, `my_list_count` should be an accessible attribute
|
||||
of the inference model (i.e. via `self.inference_model.my_list_count`).
|
||||
|
||||
NOTE: There are some APIs in which this type cannot be used as if it is
|
||||
just a list of Tensors. For example, `torch.cat(param_list)` will not work.
|
||||
However, you can make it compatible with a tuple wrapper:
|
||||
`torch.cat(tuple(param_list))`
|
||||
"""
|
||||
|
||||
n_params: int
|
||||
"""
|
||||
Number of params this list contains.
|
||||
"""
|
||||
|
||||
param: ParameterBase
|
||||
"""
|
||||
WeakRef to the owning parameter.
|
||||
"""
|
||||
|
||||
def __init__(self, param: ParameterBase, n_params: int) -> None:
|
||||
"""
|
||||
Constructor. Should not be called from client code.
|
||||
|
||||
Args:
|
||||
param (ParameterBase): The owning parameter.
|
||||
n_params (int): The number of parameters this list contains. This should be
|
||||
"""
|
||||
self.n_params = n_params
|
||||
self.set_params = 0
|
||||
self.param = weakref.ref(param)
|
||||
self._params = [None] * n_params
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self._params[index]
|
||||
|
||||
def __setitem__(self, index, value):
|
||||
if self._params[index] is not None:
|
||||
raise ValueError("Cannot set a parameter twice.")
|
||||
|
||||
self._params[index] = value
|
||||
self.set_params += 1
|
||||
|
||||
if self.set_params != self.n_params:
|
||||
return
|
||||
|
||||
self.param().complete_component()
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._params)
|
||||
|
||||
|
||||
def ParamList(attr: str):
|
||||
"""
|
||||
Helper to create a subclass of ParametrizedList with the desired `count_attr`.
|
||||
|
||||
In this manner, we can annotate the type of a Parameter dependency with the
|
||||
following:
|
||||
|
||||
```python
|
||||
class CustomParameter(ParameterBase):
|
||||
dependency_list: ParamList("dependencies_count_name")
|
||||
```
|
||||
|
||||
where "dependencies_count_name" is the name of the attribute on the inference model.
|
||||
"""
|
||||
|
||||
class ParametrizedListInstance(ParametrizedList):
|
||||
count_attr: str = attr
|
||||
|
||||
return ParametrizedListInstance
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import PhiPolicy
|
||||
@@ -0,0 +1,91 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Phi-2 model looks like this:
|
||||
|
||||
PhiForCausalLM(
|
||||
(model): PhiModel(
|
||||
(embed_tokens): Embedding(51200, 2560)
|
||||
(embed_dropout): Dropout(p=0.0, inplace=False)
|
||||
(layers): ModuleList(
|
||||
(0-31): 32 x PhiDecoderLayer(
|
||||
(self_attn): PhiAttention(
|
||||
(q_proj): Linear(in_features=2560, out_features=2560, bias=True)
|
||||
(k_proj): Linear(in_features=2560, out_features=2560, bias=True)
|
||||
(v_proj): Linear(in_features=2560, out_features=2560, bias=True)
|
||||
(dense): Linear(in_features=2560, out_features=2560, bias=True)
|
||||
(rotary_emb): PhiRotaryEmbedding()
|
||||
)
|
||||
(mlp): PhiMLP(
|
||||
(activation_fn): NewGELUActivation()
|
||||
(fc1): Linear(in_features=2560, out_features=10240, bias=True)
|
||||
(fc2): Linear(in_features=10240, out_features=2560, bias=True)
|
||||
)
|
||||
(input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
|
||||
(resid_dropout): Dropout(p=0.1, inplace=False)
|
||||
)
|
||||
)
|
||||
(final_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
|
||||
)
|
||||
(lm_head): Linear(in_features=2560, out_features=51200, bias=True)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class PhiTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Phi model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
qkv_b: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
attn_out_b: AttentionOutputParameter
|
||||
mlp_1_w: MLP1Parameter
|
||||
mlp_1_b: MLP1Parameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
mlp_2_b: MLP2Parameter
|
||||
ln_gamma: NormParameter
|
||||
ln_beta: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.q_proj.bias": "qkv_b.q_params",
|
||||
"self_attn.k_proj.bias": "qkv_b.k_params",
|
||||
"self_attn.v_proj.bias": "qkv_b.v_params",
|
||||
"self_attn.dense.weight": "attn_out_w.params",
|
||||
"self_attn.dense.bias": "attn_out_b.params",
|
||||
"mlp.fc1.weight": "mlp_1_w.params",
|
||||
"mlp.fc1.bias": "mlp_1_b.params",
|
||||
"mlp.fc2.weight": "mlp_2_w.params",
|
||||
"mlp.fc2.bias": "mlp_2_b.params",
|
||||
"input_layernorm.weight": "ln_gamma.params",
|
||||
"input_layernorm.bias": "ln_beta.params",
|
||||
}
|
||||
|
||||
|
||||
class PhiNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Phi model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed_w: UnembedParameter
|
||||
word_unembed_b: UnembedParameter
|
||||
final_norm_gamma: NormParameter
|
||||
final_norm_beta: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.final_layernorm.weight": "final_norm_gamma.params",
|
||||
"model.final_layernorm.bias": "final_norm_beta.params",
|
||||
"lm_head.weight": "word_unembed_w.params",
|
||||
"lm_head.bias": "word_unembed_b.params",
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .containers import PhiNonTransformerContainer, PhiTransformerContainer
|
||||
|
||||
|
||||
class PhiInferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[PhiNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[PhiTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties inherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties inherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.GELU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.LayerNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
rotary_dim = int(self._config.partial_rotary_factor * self.head_size)
|
||||
return RotateHalfConfig(rotate_dim=rotary_dim, theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
attn_ln_out = hidden_states
|
||||
attn_hidden_state = self.qkv(attn_ln_out, cur_params.qkv_w, b=cur_params.qkv_b)
|
||||
attn_hidden_state = self.attn(attn_hidden_state, kv_cache, ragged_batch_info)
|
||||
attention_output = self.attn_out(attn_hidden_state, cur_params.attn_out_w, b=cur_params.attn_out_b)
|
||||
|
||||
mlp_ln_out = hidden_states
|
||||
mlp_hidden_state = self.mlp_1(mlp_ln_out, cur_params.mlp_1_w, b=cur_params.mlp_1_b)
|
||||
mlp_output = self.mlp_2(mlp_hidden_state, cur_params.mlp_2_w, b=cur_params.mlp_2_b)
|
||||
|
||||
mlp_output.add_(attention_output)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(mlp_output, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, mlp_output = self.norm(residual, mlp_output, next_params.ln_gamma, beta=next_params.ln_beta)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(mlp_output)
|
||||
|
||||
return residual, mlp_output
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed_w,
|
||||
ragged_batch_info,
|
||||
bias=self._non_transformer.word_unembed_b,
|
||||
gamma=self._non_transformer.final_norm_gamma,
|
||||
beta=self._non_transformer.final_norm_beta)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual,
|
||||
None,
|
||||
gamma=self._transformer[0].ln_gamma,
|
||||
beta=self._transformer[0].ln_beta)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,32 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .containers import PhiNonTransformerContainer, PhiTransformerContainer
|
||||
from .model import PhiInferenceModel
|
||||
|
||||
|
||||
class PhiPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> PhiInferenceModel:
|
||||
return PhiInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
trans_container_cls = PhiTransformerContainer
|
||||
transformer_containers = [trans_container_cls(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(PhiNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(
|
||||
[f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import Phi3Policy
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Phi-3 model looks like this:
|
||||
|
||||
Phi3ForCausalLM(
|
||||
(model): Phi3Model(
|
||||
(embed_tokens): Embedding(32064, 3072)
|
||||
(embed_dropout): Dropout(p=0.0, inplace=False)
|
||||
(layers): ModuleList(
|
||||
(0-31): 32 x Phi3DecoderLayer(
|
||||
(self_attn): Phi3Attention(
|
||||
(o_proj): Linear(in_features=3072, out_features=3072, bias=False)
|
||||
(qkv_proj): Linear(in_features=3072, out_features=9216, bias=False)
|
||||
(rotary_emb): Phi3RotaryEmbedding()
|
||||
)
|
||||
(mlp): PhiMLP(
|
||||
(gate_up_proj): Linear(in_features=3072, out_features=16384, bias=False)
|
||||
(down_proj): Linear(in_features=16384, out_features=3072, bias=False)
|
||||
(activation_fn): SiLU()
|
||||
)
|
||||
(input_layernorm): Phi3RMSNorm((3072,), eps=1e-05)
|
||||
(resid_attn_dropout): Dropout(p=0.0)
|
||||
(resid_mlp_dropout): Dropout(p=0.0)
|
||||
(post_attention_layernorm): Phi3RMSNorm((3072,), eps=1e-05)
|
||||
)
|
||||
)
|
||||
(final_layernorm): Phi3RMSNorm((3072,), eps=1e-05)
|
||||
)
|
||||
(lm_head): Linear(in_features=3072, out_features=32064, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class Phi3TransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Phi model.
|
||||
"""
|
||||
qkv_w: FusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: FusedGatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.qkv_proj.weight": "qkv_w.params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate_up_proj.weight": "mlp_1_w.params",
|
||||
"mlp.down_proj.weight": "mlp_2_w.params",
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class Phi3NonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Phi model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed_w: UnembedParameter
|
||||
final_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm_gamma.params",
|
||||
"lm_head.weight": "word_unembed_w.params",
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .containers import Phi3NonTransformerContainer, Phi3TransformerContainer
|
||||
|
||||
|
||||
class Phi3InferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[Phi3NonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[Phi3TransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties inherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties inherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
if self._config.torch_dtype == torch.float16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.torch_dtype == torch.bfloat16:
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
activation = self._config.hidden_act.lower()
|
||||
if activation == "gelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "relu":
|
||||
return ActivationType.ReGLU
|
||||
elif activation == "gegelu":
|
||||
return ActivationType.GEGLU
|
||||
elif activation == "silu":
|
||||
return ActivationType.SiGLU
|
||||
else:
|
||||
raise NotImplementedError(f"Activation {activation} not supported")
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed_w,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm_gamma)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, gamma=self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .containers import Phi3NonTransformerContainer, Phi3TransformerContainer
|
||||
from .model import Phi3InferenceModel
|
||||
|
||||
|
||||
class Phi3Policy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Phi3InferenceModel:
|
||||
return Phi3InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [Phi3TransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(Phi3NonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params([])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import QwenPolicy
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Qwen model looks like this:
|
||||
|
||||
QWenLMHeadModel(
|
||||
(transformer): QWenModel(
|
||||
(wte): Embedding(151936, 4096)
|
||||
(drop): Dropout(p=0.0, inplace=False)
|
||||
(rotary_emb): RotaryEmbedding()
|
||||
(h): ModuleList(
|
||||
(0-31): 32 x QWenBlock(
|
||||
(ln_1): RMSNorm()
|
||||
(attn): QWenAttention(
|
||||
(c_attn): Linear(in_features=4096, out_features=12288, bias=True)
|
||||
(c_proj): Linear(in_features=4096, out_features=4096, bias=False)
|
||||
(attn_dropout): Dropout(p=0.0, inplace=False)
|
||||
)
|
||||
(ln_2): RMSNorm()
|
||||
(mlp): QWenMLP(
|
||||
(w1): Linear(in_features=4096, out_features=11008, bias=False)
|
||||
(w2): Linear(in_features=4096, out_features=11008, bias=False)
|
||||
(c_proj): Linear(in_features=11008, out_features=4096, bias=False)
|
||||
)
|
||||
)
|
||||
)
|
||||
(ln_f): RMSNorm()
|
||||
)
|
||||
(lm_head): Linear(in_features=4096, out_features=151936, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class QwenTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Qwen model.
|
||||
"""
|
||||
qkv_w: FusedQKVParameter
|
||||
qkv_b: FusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: GatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"attn.c_attn.weight": "qkv_w.params",
|
||||
"attn.c_attn.bias": "qkv_b.params",
|
||||
"attn.c_proj.weight": "attn_out_w.params",
|
||||
"mlp.w1.weight": "mlp_1_w.up_params",
|
||||
"mlp.w2.weight": "mlp_1_w.gate_params",
|
||||
"mlp.c_proj.weight": "mlp_2_w.params",
|
||||
"ln_1.weight": "attn_norm_gamma.params",
|
||||
"ln_2.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class QwenNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Qwen model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"transformer.wte.weight": "word_emb.params",
|
||||
"transformer.ln_f.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...modules import heuristics
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .container import QwenNonTransformerContainer, QwenTransformerContainer
|
||||
|
||||
|
||||
class QwenInferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[QwenNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[QwenTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size // 2
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.hidden_size // self._config.kv_channels
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
autoset_precision = self._config.bf16 + self._config.fp16 == 0
|
||||
if autoset_precision:
|
||||
return DtypeEnum.fp16
|
||||
if self._config.fp16:
|
||||
return DtypeEnum.fp16
|
||||
elif self._config.bf16:
|
||||
# TODO(ZonePG): bf16 inference results may be different from huggingface bf16,
|
||||
# because in rms_norm, Qwen still use float() instead of bf16
|
||||
return DtypeEnum.bf16
|
||||
else:
|
||||
raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.SiGLU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rotary_emb_base)
|
||||
|
||||
def make_norm_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the normalization layer for the model. This sets the `self.norm` attribute.
|
||||
|
||||
TODO(cmikeh2): In the future we'll distinguish between the different norm objects,
|
||||
but for now we'll just use the same one for all of them.
|
||||
"""
|
||||
norm_config = DSNormConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
type=self.norm_type,
|
||||
channels=self.model_dim,
|
||||
residual_dtype=self.activation_dtype,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
eps=self._config.layer_norm_epsilon,
|
||||
)
|
||||
|
||||
self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
# Should be configurable in the future
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import QwenNonTransformerContainer, QwenTransformerContainer
|
||||
from .model import QwenInferenceModel
|
||||
|
||||
|
||||
class QwenPolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> QwenInferenceModel:
|
||||
return QwenInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [QwenTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['transformer.h'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(QwenNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(['transformer.rotary_emb.inv_freq'])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import Qwen2Policy
|
||||
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Qwen2 model looks like this:
|
||||
|
||||
Qwen2ForCausalLM(
|
||||
(model): Qwen2Model(
|
||||
(embed_tokens): Embedding(151936, 1024)
|
||||
(layers): ModuleList(
|
||||
(0-23): 24 x Qwen2DecoderLayer(
|
||||
(self_attn): Qwen2SdpaAttention(
|
||||
(q_proj): Linear(in_features=1024, out_features=1024, bias=True)
|
||||
(k_proj): Linear(in_features=1024, out_features=1024, bias=True)
|
||||
(v_proj): Linear(in_features=1024, out_features=1024, bias=True)
|
||||
(o_proj): Linear(in_features=1024, out_features=1024, bias=False)
|
||||
(rotary_emb): Qwen2RotaryEmbedding()
|
||||
)
|
||||
(mlp): Qwen2MLP(
|
||||
(gate_proj): Linear(in_features=1024, out_features=2816, bias=False)
|
||||
(up_proj): Linear(in_features=1024, out_features=2816, bias=False)
|
||||
(down_proj): Linear(in_features=2816, out_features=1024, bias=False)
|
||||
(act_fn): SiLU()
|
||||
)
|
||||
(input_layernorm): Qwen2RMSNorm()
|
||||
(post_attention_layernorm): Qwen2RMSNorm()
|
||||
)
|
||||
)
|
||||
(norm): Qwen2RMSNorm()
|
||||
)
|
||||
(lm_head): Linear(in_features=1024, out_features=151936, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class Qwen2TransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Qwen2 model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
qkv_b: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
mlp_1_w: GatedMLPParameter
|
||||
mlp_2_w: MLP2Parameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.q_proj.bias": "qkv_b.q_params",
|
||||
"self_attn.k_proj.bias": "qkv_b.k_params",
|
||||
"self_attn.v_proj.bias": "qkv_b.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate_proj.weight": "mlp_1_w.gate_params",
|
||||
"mlp.up_proj.weight": "mlp_1_w.up_params",
|
||||
"mlp.down_proj.weight": "mlp_2_w.params",
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class Qwen2NonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Qwen2 model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from .. import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...modules import heuristics
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
|
||||
from .container import Qwen2NonTransformerContainer, Qwen2TransformerContainer
|
||||
|
||||
|
||||
class Qwen2InferenceModel(DSTransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for ragged batching for Llama-2 models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[Qwen2NonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[Qwen2TransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_seq_length
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
# TODO(ZonePG): bf16 inference results may be different from huggingface bf16,
|
||||
# because in rms_norm, Qwen still use float() instead of bf16
|
||||
# if self._config.torch_dtype == torch.float16:
|
||||
# return DtypeEnum.fp16
|
||||
# elif self._config.torch_dtype == torch.bfloat16:
|
||||
# return DtypeEnum.bf16
|
||||
# else:
|
||||
# raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
return DtypeEnum.fp16
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.SiGLU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
def make_norm_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the normalization layer for the model. This sets the `self.norm` attribute.
|
||||
|
||||
TODO(cmikeh2): In the future we'll distinguish between the different norm objects,
|
||||
but for now we'll just use the same one for all of them.
|
||||
"""
|
||||
norm_config = DSNormConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
type=self.norm_type,
|
||||
channels=self.model_dim,
|
||||
residual_dtype=self.activation_dtype,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
eps=self._config.rms_norm_eps,
|
||||
)
|
||||
|
||||
self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config)
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
# Should be configurable in the future
|
||||
hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None)
|
||||
hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states,
|
||||
wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import Qwen2NonTransformerContainer, Qwen2TransformerContainer
|
||||
from .model import Qwen2InferenceModel
|
||||
|
||||
|
||||
class Qwen2Policy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Qwen2InferenceModel:
|
||||
return Qwen2InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [Qwen2TransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(Qwen2NonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params(
|
||||
[f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .policy import Qwen2MoePolicy
|
||||
@@ -0,0 +1,103 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
|
||||
from ..common_parameters import *
|
||||
from ..layer_container_base import LayerContainer
|
||||
'''
|
||||
# HF Qwen2-57B-A14B model looks like this:
|
||||
|
||||
Qwen2MoeForCausalLM(
|
||||
(model): Qwen2MoeModel(
|
||||
(embed_tokens): Embedding(151936, 3584)
|
||||
(layers): ModuleList(
|
||||
(0-27): 28 x Qwen2MoeDecoderLayer(
|
||||
(self_attn): Qwen2MoeSdpaAttention(
|
||||
(q_proj): Linear(in_features=3584, out_features=3584, bias=True)
|
||||
(k_proj): Linear(in_features=3584, out_features=512, bias=True)
|
||||
(v_proj): Linear(in_features=3584, out_features=512, bias=True)
|
||||
(o_proj): Linear(in_features=3584, out_features=3584, bias=False)
|
||||
(rotary_emb): Qwen2MoeRotaryEmbedding()
|
||||
)
|
||||
(mlp): Qwen2MoeSparseMoeBlock(
|
||||
(gate): Linear(in_features=3584, out_features=64, bias=False)
|
||||
(experts): ModuleList(
|
||||
(0-63): 64 x Qwen2MoeMLP(
|
||||
(gate_proj): Linear(in_features=3584, out_features=2560, bias=False)
|
||||
(up_proj): Linear(in_features=3584, out_features=2560, bias=False)
|
||||
(down_proj): Linear(in_features=2560, out_features=3584, bias=False)
|
||||
(act_fn): SiLU()
|
||||
)
|
||||
)
|
||||
(shared_expert): Qwen2MoeMLP(
|
||||
(gate_proj): Linear(in_features=3584, out_features=20480, bias=False)
|
||||
(up_proj): Linear(in_features=3584, out_features=20480, bias=False)
|
||||
(down_proj): Linear(in_features=20480, out_features=3584, bias=False)
|
||||
(act_fn): SiLU()
|
||||
)
|
||||
(shared_expert_gate): Linear(in_features=3584, out_features=1, bias=False)
|
||||
)
|
||||
(input_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06)
|
||||
(post_attention_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06)
|
||||
)
|
||||
)
|
||||
(norm): Qwen2MoeRMSNorm((3584,), eps=1e-06)
|
||||
)
|
||||
(lm_head): Linear(in_features=3584, out_features=151936, bias=False)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
class Qwen2MoeTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Transformer layer container for the Qwen2Moe model.
|
||||
"""
|
||||
qkv_w: UnfusedQKVParameter
|
||||
qkv_b: UnfusedQKVParameter
|
||||
attn_out_w: AttentionOutputParameter
|
||||
moe_gate: MoEGatingWeightParameter
|
||||
moe_mlp_1: UnfusedMoEGatedMLPParameter
|
||||
moe_mlp_2: UnfusedMoEMLP2Parameter
|
||||
shared_moe_mlp_1: GatedMLPParameter
|
||||
shared_moe_mlp_2: MLP2Parameter
|
||||
shared_moe_gate: MoEGatingWeightParameter
|
||||
attn_norm_gamma: NormParameter
|
||||
mlp_norm_gamma: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"self_attn.q_proj.weight": "qkv_w.q_params",
|
||||
"self_attn.k_proj.weight": "qkv_w.k_params",
|
||||
"self_attn.v_proj.weight": "qkv_w.v_params",
|
||||
"self_attn.q_proj.bias": "qkv_b.q_params",
|
||||
"self_attn.k_proj.bias": "qkv_b.k_params",
|
||||
"self_attn.v_proj.bias": "qkv_b.v_params",
|
||||
"self_attn.o_proj.weight": "attn_out_w.params",
|
||||
"mlp.gate.weight": "moe_gate.params",
|
||||
"mlp.experts.*.gate_proj.weight": "moe_mlp_1.gating_experts",
|
||||
"mlp.experts.*.up_proj.weight": "moe_mlp_1.up_experts",
|
||||
"mlp.experts.*.down_proj.weight": "moe_mlp_2.experts",
|
||||
"mlp.shared_expert.gate_proj.weight": "shared_moe_mlp_1.gate_params",
|
||||
"mlp.shared_expert.up_proj.weight": "shared_moe_mlp_1.up_params",
|
||||
"mlp.shared_expert.down_proj.weight": "shared_moe_mlp_2.params",
|
||||
"mlp.shared_expert_gate.weight": "shared_moe_gate.params",
|
||||
"input_layernorm.weight": "attn_norm_gamma.params",
|
||||
"post_attention_layernorm.weight": "mlp_norm_gamma.params",
|
||||
}
|
||||
|
||||
|
||||
class Qwen2MoeNonTransformerContainer(LayerContainer):
|
||||
"""
|
||||
Non-Transformer layer container for the Qwen2Moe model.
|
||||
"""
|
||||
word_emb: EmbeddingParameter
|
||||
word_unembed: UnembedParameter
|
||||
final_norm: NormParameter
|
||||
|
||||
PARAM_MAPPING = {
|
||||
"model.embed_tokens.weight": "word_emb.params",
|
||||
"model.norm.weight": "final_norm.params",
|
||||
"lm_head.weight": "word_unembed.params",
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from ...allocator import empty_from
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ...inference_utils import ActivationType, DtypeEnum
|
||||
from ...model_implementations import *
|
||||
from ...modules.configs import *
|
||||
from ...modules.interfaces import *
|
||||
from ...modules import heuristics
|
||||
from ...ragged import RaggedBatchWrapper
|
||||
from ..inference_model_base import (
|
||||
DSModelImplementationConfig,
|
||||
MPType,
|
||||
)
|
||||
|
||||
from .container import Qwen2MoeNonTransformerContainer, Qwen2MoeTransformerContainer
|
||||
|
||||
|
||||
class Qwen2MoeInferenceModel(DSMoETransformerModelBase):
|
||||
"""
|
||||
Inference model implementation for Qwen2MoE models.
|
||||
"""
|
||||
|
||||
_non_transformer: Optional[Qwen2MoeNonTransformerContainer]
|
||||
"""
|
||||
Embed + unembed container. Specializing the type annotation.
|
||||
"""
|
||||
|
||||
_transformer: Optional[Iterable[Qwen2MoeTransformerContainer]]
|
||||
"""
|
||||
Per-layer transformer container. Specializing the type annotation.
|
||||
"""
|
||||
"""
|
||||
Properties ineherited from `DSInferenceModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def max_sequence_length(self) -> int:
|
||||
return self._config.max_position_embeddings
|
||||
|
||||
"""
|
||||
Properties ineherited from `DSTransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def model_dim(self) -> int:
|
||||
return self._config.hidden_size
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self._config.vocab_size
|
||||
|
||||
@property
|
||||
def head_size(self) -> int:
|
||||
return self.model_dim // self.n_heads
|
||||
|
||||
@property
|
||||
def n_heads(self) -> int:
|
||||
return self._config.num_attention_heads
|
||||
|
||||
@property
|
||||
def intermediate_dim(self) -> int:
|
||||
return self._config.shared_expert_intermediate_size
|
||||
|
||||
@property
|
||||
def n_heads_kv(self) -> int:
|
||||
return self._config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def activation_dtype(self) -> DtypeEnum:
|
||||
# TODO(ZonePG): bf16 inference results may be different from huggingface bf16,
|
||||
# because in rms_norm, Qwen still use float() instead of bf16
|
||||
# if self._config.torch_dtype == torch.float16:
|
||||
# return DtypeEnum.fp16
|
||||
# elif self._config.torch_dtype == torch.bfloat16:
|
||||
# return DtypeEnum.bf16
|
||||
# else:
|
||||
# raise NotImplementedError("Only fp16 and bf16 are supported")
|
||||
return DtypeEnum.fp16
|
||||
|
||||
@property
|
||||
def mlp_activation_fn(self) -> ActivationType:
|
||||
return ActivationType.SiGLU
|
||||
|
||||
@property
|
||||
def norm_type(self) -> NormTypeEnum:
|
||||
return NormTypeEnum.RMSNorm
|
||||
|
||||
@property
|
||||
def positional_embedding_type(self) -> PositionalEmbeddingType:
|
||||
return PositionalEmbeddingType.rotate_half
|
||||
|
||||
@property
|
||||
def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
|
||||
return RotateHalfConfig(theta_base=self._config.rope_theta)
|
||||
|
||||
"""
|
||||
Inherited from `DSMoETransformerModelBase`
|
||||
"""
|
||||
|
||||
@property
|
||||
def n_experts(self) -> int:
|
||||
return self._config.num_experts
|
||||
|
||||
@property
|
||||
def n_top_k(self) -> int:
|
||||
return self._config.num_experts_per_tok
|
||||
|
||||
@property
|
||||
def normalize_expert_scores(self) -> bool:
|
||||
return self._config.norm_topk_prob
|
||||
|
||||
def make_moe_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the MoE layer for the model. This sets the `self.moe` attribute.
|
||||
"""
|
||||
sharded_dim = sharded_intermediate_dim(self.intermediate_dim // self.n_top_k, self.tp_size, self.tp_rank)
|
||||
|
||||
moe_config = DSMoEConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
model_dim=self.model_dim,
|
||||
intermediate_features=sharded_dim,
|
||||
activation=self.mlp_activation_fn,
|
||||
n_experts=self.n_experts,
|
||||
top_k=self.n_top_k,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
normalize_scores=self.normalize_expert_scores,
|
||||
)
|
||||
|
||||
self.moe = heuristics.instantiate_moe(moe_config, self._engine_config)
|
||||
|
||||
######### MLP 1 #########
|
||||
def make_shared_expert_mlp_1_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the first MLP in the feedforward network.
|
||||
This sets the `self.mlp_1` attribute.
|
||||
"""
|
||||
shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=self.model_dim,
|
||||
out_channels=shard_size,
|
||||
activation=self.mlp_activation_fn,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.shared_expert_mlp_1 = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
######### MLP 2 #########
|
||||
def make_shared_expert_mlp_2_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the second MLP in the feedforward network.
|
||||
This sets the `self.mlp_2` attribute.
|
||||
"""
|
||||
shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=shard_size,
|
||||
out_channels=self.model_dim,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.shared_expert_mlp_2 = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
######### MLP 2 #########
|
||||
def make_shared_expert_gate_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the linear projection layer for the second MLP in the feedforward network.
|
||||
This sets the `self.mlp_2` attribute.
|
||||
"""
|
||||
shard_size = sharded_intermediate_dim(self.model_dim, self.tp_size, self.tp_rank)
|
||||
|
||||
linear_config = DSLinearConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
in_channels=shard_size,
|
||||
out_channels=8,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
)
|
||||
|
||||
self.shared_expert_gate = heuristics.instantiate_linear(linear_config, self._engine_config)
|
||||
|
||||
def make_norm_layer(self) -> None:
|
||||
"""
|
||||
Instantiates the normalization layer for the model. This sets the `self.norm` attribute.
|
||||
|
||||
TODO(cmikeh2): In the future we'll distinguish between the different norm objects,
|
||||
but for now we'll just use the same one for all of them.
|
||||
"""
|
||||
norm_config = DSNormConfig(
|
||||
max_tokens=self._engine_config.state_manager.max_ragged_batch_size,
|
||||
type=self.norm_type,
|
||||
channels=self.model_dim,
|
||||
residual_dtype=self.activation_dtype,
|
||||
input_dtype=self.activation_dtype,
|
||||
output_dtype=self.activation_dtype,
|
||||
eps=self._config.rms_norm_eps,
|
||||
)
|
||||
|
||||
self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config)
|
||||
|
||||
"""
|
||||
Model implementation
|
||||
"""
|
||||
|
||||
def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig,
|
||||
base_mp_group: MPType) -> None:
|
||||
"""
|
||||
Base implementation for initialization. By default, this will initialize
|
||||
the traditional components of a transformer model:
|
||||
- Embedding
|
||||
- QKV projection
|
||||
- Self attention
|
||||
- Attention output projection
|
||||
- Feed forward network
|
||||
- Normalization
|
||||
- Unembedding
|
||||
|
||||
Arguments:
|
||||
config (DSModelImplementationConfig): Model-specific configuration. No assumptions
|
||||
should be made about this config that are not closely tied to the specific
|
||||
model implementation.
|
||||
engine_config (RaggedInferenceEngineConfig): Engine configuration.
|
||||
base_mp_group (MPType): Base communication group for Tensor-parallel inference.
|
||||
"""
|
||||
super().__init__(config, engine_config, base_mp_group)
|
||||
|
||||
self.make_norm_layer()
|
||||
self.make_qkv_layer()
|
||||
self.make_attn_layer()
|
||||
self.make_attn_out_layer()
|
||||
self.make_moe_layer()
|
||||
self.make_shared_expert_mlp_1_layer()
|
||||
self.make_shared_expert_mlp_2_layer()
|
||||
self.make_shared_expert_gate_layer()
|
||||
self.make_embedding_layer()
|
||||
self.make_unembedding_layer()
|
||||
self._kv_cache_config = None
|
||||
|
||||
"""
|
||||
Forward implementations
|
||||
"""
|
||||
|
||||
def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs the embedding lookup prior to running the transformer of the model.
|
||||
|
||||
Arguments:
|
||||
ragged_batch (RaggedBatchWrapper): The batch to embed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The embedded batch.
|
||||
"""
|
||||
embed = self.embed(ragged_batch, self._non_transformer.word_emb)
|
||||
|
||||
if embed.shape[-1] != self.model_dim:
|
||||
raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")
|
||||
|
||||
return embed
|
||||
|
||||
def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor,
|
||||
ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead
|
||||
optimization to fuse the layer norm of the next layer into the current layer.
|
||||
|
||||
Arguments:
|
||||
layer_idx (int): The index of the layer to execute.
|
||||
residual (torch.Tensor): The residual tensor from the previous layer.
|
||||
hidden_states (torch.Tensor): The hidden states from the previous layer. This is the
|
||||
hidden states after pre normalization.
|
||||
ragged_batch_info (RaggedBatchWrapper): The batch metadata.
|
||||
"""
|
||||
# TODO(cmikeh2): Distribute ragged_batch_info to all modules
|
||||
|
||||
cur_params = self._transformer[layer_idx]
|
||||
kv_cache = self.state_manager.get_cache(layer_idx)
|
||||
|
||||
hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b)
|
||||
hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info)
|
||||
hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None)
|
||||
|
||||
shared_expert_output = self.shared_expert_mlp_1(hidden_states, cur_params.shared_moe_mlp_1, b=None)
|
||||
shared_expert_output = self.shared_expert_mlp_2(shared_expert_output, cur_params.shared_moe_mlp_2, b=None)
|
||||
shared_expert_gate_output = self.shared_expert_gate(hidden_states, cur_params.shared_moe_gate, b=None)[..., :1]
|
||||
# shared_expert_gate_output shape[-1] is 1
|
||||
shared_expert_output.mul_(torch.sigmoid(shared_expert_gate_output))
|
||||
hidden_states = self.moe(hidden_states, ragged_batch_info, cur_params.moe_gate, cur_params.moe_mlp_1,
|
||||
cur_params.moe_mlp_2)
|
||||
hidden_states.add_(shared_expert_output)
|
||||
|
||||
if self.tp_size > 1:
|
||||
dist.all_reduce(hidden_states, group=self._base_mp_group)
|
||||
|
||||
if layer_idx != self.num_layers - 1:
|
||||
next_params = self._transformer[layer_idx + 1]
|
||||
residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None)
|
||||
else:
|
||||
# On last layer, we just need to perform the residual add. Adding into the residual
|
||||
# here is safe.
|
||||
residual.add_(hidden_states)
|
||||
|
||||
return residual, hidden_states
|
||||
|
||||
def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor:
|
||||
"""
|
||||
Performs unembedding of the hidden states to logits. This will only sample the final
|
||||
token of each sequence.
|
||||
"""
|
||||
logits = self.unembed(hidden_states,
|
||||
self._non_transformer.word_unembed,
|
||||
ragged_batch_info,
|
||||
gamma=self._non_transformer.final_norm)
|
||||
|
||||
if self.tp_size > 1:
|
||||
comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1]))
|
||||
full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size))
|
||||
|
||||
dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group)
|
||||
|
||||
full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size))
|
||||
|
||||
return full_logits
|
||||
else:
|
||||
return logits
|
||||
|
||||
def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
|
||||
|
||||
residual = self._forward_embed(wrapped_batch)
|
||||
|
||||
residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None)
|
||||
|
||||
for layer_idx in range(self.num_layers):
|
||||
residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch)
|
||||
|
||||
return self._forward_unembed(residual, wrapped_batch)
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...config_v2 import RaggedInferenceEngineConfig
|
||||
from ..inference_policy_base import ContainerMap, InferenceV2Policy
|
||||
from .container import Qwen2MoeNonTransformerContainer, Qwen2MoeTransformerContainer
|
||||
from .model import Qwen2MoeInferenceModel
|
||||
|
||||
|
||||
class Qwen2MoePolicy(InferenceV2Policy):
|
||||
|
||||
def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Qwen2MoeInferenceModel:
|
||||
return Qwen2MoeInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group)
|
||||
|
||||
def build_container_map(self) -> ContainerMap:
|
||||
map = ContainerMap()
|
||||
|
||||
transformer_containers = [Qwen2MoeTransformerContainer(self.model) for _ in range(self.model.num_layers)]
|
||||
|
||||
map.set_transformer_params(['model.layers'], transformer_containers)
|
||||
|
||||
map.set_non_transformer_params(Qwen2MoeNonTransformerContainer(self.model))
|
||||
|
||||
map.set_unmapped_params([])
|
||||
|
||||
return map
|
||||
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .attn import *
|
||||
from .attn_out import *
|
||||
from .embedding import *
|
||||
from .mlp import *
|
||||
from .qkv import *
|
||||
from .types import *
|
||||
from .unembed import *
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
def get_local_heads(shard_rank: int,
|
||||
num_shards: int,
|
||||
n_heads_q: int,
|
||||
n_heads_kv: Optional[int] = None) -> Tuple[int, int]:
|
||||
"""
|
||||
Helper to determine the number of local heads of a given shard.
|
||||
|
||||
Args:
|
||||
shard_rank (int): The rank of the shard.
|
||||
num_shards (int): The total number of shards that attention is distributed over.
|
||||
n_heads_q (int): The number of query heads.
|
||||
n_heads_kv (int): The number of key/value heads. If not passed, it is assumed that
|
||||
the number of query and key/value heads are the same.
|
||||
"""
|
||||
if n_heads_q < num_shards:
|
||||
raise ValueError("There must be at least as many attention heads as there are shards.")
|
||||
|
||||
if n_heads_kv is None or n_heads_kv == n_heads_q:
|
||||
# MHA attention
|
||||
base_heads = n_heads_q // num_shards
|
||||
extra_heads = n_heads_q % num_shards
|
||||
|
||||
if shard_rank < extra_heads:
|
||||
return (base_heads + 1), (base_heads + 1)
|
||||
else:
|
||||
return base_heads, base_heads
|
||||
else:
|
||||
# GQA attention
|
||||
if n_heads_q % n_heads_kv != 0:
|
||||
raise ValueError("Must be an even ratio between query and key/value heads.")
|
||||
|
||||
if n_heads_kv < num_shards and num_shards % n_heads_kv != 0:
|
||||
raise ValueError(
|
||||
"If splitting a group across multiple shards, we must be able to distribute the groups evenly.")
|
||||
|
||||
if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0:
|
||||
raise ValueError("If parallelizing groups, must be able to evenly distribute them.")
|
||||
|
||||
q_ratio = n_heads_q // n_heads_kv
|
||||
|
||||
if n_heads_kv >= num_shards:
|
||||
local_kv_heads = n_heads_kv // num_shards
|
||||
local_q_heads = local_kv_heads * q_ratio
|
||||
return local_q_heads, local_kv_heads
|
||||
else:
|
||||
group_sharding_size = num_shards // n_heads_kv
|
||||
group_rank_idx = shard_rank % group_sharding_size
|
||||
|
||||
base_heads = q_ratio // group_sharding_size
|
||||
extra_heads = q_ratio % group_sharding_size
|
||||
|
||||
if group_rank_idx < extra_heads:
|
||||
return (base_heads + 1), 1
|
||||
else:
|
||||
return base_heads, 1
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType
|
||||
from .utils import shard_param, get_shard_endpoints
|
||||
|
||||
|
||||
def shard_attn_out_param(param: torch.Tensor,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
head_size: int,
|
||||
n_heads_q: Optional[int] = None,
|
||||
n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Utility method for sharding an attention output parameter.
|
||||
"""
|
||||
if len(param.shape) == 1:
|
||||
# We will do the bias addition on the 0th rank only rather than scale the parameter and
|
||||
# implicitly reconstruct this in the distributed reduce.
|
||||
return param if shard_rank == 0 else None
|
||||
|
||||
assert n_heads_kv is None or (n_heads_q is not None
|
||||
and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q"
|
||||
|
||||
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
|
||||
|
||||
if mha_sharding:
|
||||
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size)
|
||||
else:
|
||||
assert param.shape[0] == head_size * n_heads_q, "GQA param shape is not correct"
|
||||
|
||||
# 32 KV heads, 16 shards for example
|
||||
even_kv_sharding = n_heads_kv % num_shards == 0
|
||||
|
||||
# 8 KV heads, 16 shards for example
|
||||
even_kv_distribution = num_shards % n_heads_kv == 0
|
||||
|
||||
assert even_kv_sharding or even_kv_distribution, "No partitioning algorithm for this yet."
|
||||
|
||||
if even_kv_sharding:
|
||||
# Same as original sharding scenario
|
||||
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size)
|
||||
else:
|
||||
# We will first do a sharding on the KV and Q to map to the one KV shard per group of Q.
|
||||
q_sharding_degree = num_shards // n_heads_kv
|
||||
|
||||
kv_head = shard_rank // q_sharding_degree
|
||||
|
||||
q_sharding_rank = shard_rank % q_sharding_degree
|
||||
q_factor = n_heads_q // n_heads_kv
|
||||
|
||||
q_chunk = param[..., q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size]
|
||||
|
||||
return shard_param(q_chunk,
|
||||
ShardingType.INNER_DIMENSION,
|
||||
q_sharding_rank,
|
||||
q_sharding_degree,
|
||||
granularity=head_size)
|
||||
|
||||
|
||||
def attn_out_in_features(out_features: int,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
head_size: int,
|
||||
n_heads_q: Optional[int] = None,
|
||||
n_heads_kv: Optional[int] = None) -> int:
|
||||
"""
|
||||
Helper to calculate the expected output projection dimension of a QKV projection matrix.
|
||||
|
||||
Args:
|
||||
in_features (int): The model dimension.
|
||||
shard_rank (int): Which rank to return the corresponding size for.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
head_size (int): The size of each attention head.
|
||||
n_heads_q (int): The number of query heads on the model. This only needs to be passed if the number
|
||||
of query and key/value heads are different. If passed without n_heads_kv, default
|
||||
MHA partitioning will be used.
|
||||
n_heads_kv (int): The number of key and value heads on the model. This only needs to be passed
|
||||
if the number of query and key/value heads are different. This argument cannot be passed without
|
||||
also passing n_heads_q (we want to explicitly opt into GQA sharding).
|
||||
"""
|
||||
assert n_heads_kv is None or (n_heads_q is not None
|
||||
and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q"
|
||||
|
||||
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
|
||||
|
||||
if mha_sharding:
|
||||
endpoints = get_shard_endpoints(out_features, shard_rank, num_shards, granularity=head_size)
|
||||
return endpoints[1] - endpoints[0]
|
||||
else:
|
||||
if n_heads_kv >= num_shards:
|
||||
assert n_heads_kv % num_shards == 0, "No partitioning algorithm for this yet."
|
||||
n_local_groups = n_heads_kv // num_shards
|
||||
group_size = n_heads_q // n_heads_kv
|
||||
|
||||
return n_local_groups * head_size * group_size
|
||||
else:
|
||||
assert num_shards % n_heads_kv == 0, "No partitioning algorithm for this yet."
|
||||
q_split_degree = num_shards // n_heads_kv
|
||||
q_split_rank = shard_rank % q_split_degree
|
||||
split_granularity = (n_heads_q // n_heads_kv) * head_size
|
||||
|
||||
q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size)
|
||||
|
||||
return q_endpoints[1] - q_endpoints[0]
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType
|
||||
from .utils import shard_param, get_shard_endpoints
|
||||
|
||||
|
||||
def shard_embedding_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor:
|
||||
"""
|
||||
Utility method for sharding an embedding parameter.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
"""
|
||||
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards)
|
||||
|
||||
|
||||
def sharded_embedding_dim(embedding_size: int, shard_rank: int, num_shards: int) -> int:
|
||||
"""
|
||||
Utility method for getting the size of the embedding dimension of a sharded embedding.
|
||||
|
||||
Args:
|
||||
embedding_size (int): The size of the embedding.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
"""
|
||||
start_idx, end_idx = get_shard_endpoints(embedding_size, shard_rank, num_shards)
|
||||
return end_idx - start_idx
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
|
||||
from .utils import shard_param, get_shard_endpoints
|
||||
|
||||
|
||||
def shard_mlp_1_param(param: torch.Tensor,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
gated: bool = False,
|
||||
is_moe: bool = False) -> torch.Tensor:
|
||||
"""
|
||||
Utility method for sharding an MLP 1 parameter. Both biases and weights are supported, as well
|
||||
as for fused weights for MoE.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
gated (bool): Whether or not the parameter is from a gated MLP.
|
||||
"""
|
||||
bias_dims = 2 if is_moe else 1
|
||||
|
||||
if gated:
|
||||
return shard_param(param,
|
||||
ShardingType.OUTER_DIMENSION,
|
||||
shard_rank,
|
||||
num_shards,
|
||||
granularity=DEFAULT_SHARD_GRANULARITY * 2,
|
||||
bias_dims=bias_dims)
|
||||
else:
|
||||
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, bias_dims=bias_dims)
|
||||
|
||||
|
||||
def shard_mlp_2_param(param: torch.Tensor,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
is_moe: bool = False) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Utility method for sharding an MLP 2 parameter.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
is_moe (bool): Whether or not the parameter is from a MoE model.
|
||||
"""
|
||||
bias_dim_size = 2 if is_moe else 1
|
||||
|
||||
if len(param.shape) == bias_dim_size:
|
||||
# We will do the bias addition on the 0th rank only rather than scale the parameter and
|
||||
# implicitly reconstruct this in the distributed reduce.
|
||||
return param if shard_rank == 0 else None
|
||||
|
||||
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards)
|
||||
|
||||
|
||||
def sharded_intermediate_dim(intermediate_size: int, num_shards: int, shard_rank: int) -> int:
|
||||
"""
|
||||
Utility method for getting the size of the intermediate dimension of a sharded MLP.
|
||||
|
||||
Args:
|
||||
intermediate_size (int): The size of the intermediate dimension.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
"""
|
||||
endpoints = get_shard_endpoints(intermediate_size, shard_rank, num_shards)
|
||||
return endpoints[1] - endpoints[0]
|
||||
@@ -0,0 +1,167 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType
|
||||
from .utils import shard_param, get_shard_endpoints
|
||||
|
||||
|
||||
def shard_qkv_param(param: torch.Tensor,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
head_size: int,
|
||||
n_heads_q: Optional[int] = None,
|
||||
n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Utility method for sharding a QKV parameter. Both biases and weights are supported. It is assumed
|
||||
that the layout of the parameter is such that all Q heads, all K heads, and all V heads
|
||||
are contiguous with respect to each other.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
head_size (int): The size of each head.
|
||||
n_heads_q (int): The number of query heads. This only needs to be passed if the number
|
||||
of query and key/value heads are different. If passed without n_heads_kv, default
|
||||
MHA partitioning will be used.
|
||||
n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number
|
||||
of query and key/value heads are different. This argument should not be passed without
|
||||
n_heads_q (we want to explicitly opt into GQA sharding).
|
||||
"""
|
||||
if n_heads_kv is not None and n_heads_q is None:
|
||||
raise ValueError("n_heads_kv should not be passed without n_heads_q")
|
||||
|
||||
if param is None:
|
||||
raise ValueError("param should not be None")
|
||||
if n_heads_q is None:
|
||||
# Guaranteed to be in MHA
|
||||
if param.shape[0] // 3 % head_size != 0:
|
||||
raise ValueError("MHA param shape is not correct")
|
||||
n_heads_q = param.shape[0] // head_size // 3
|
||||
mha_sharding = True
|
||||
elif n_heads_kv is None:
|
||||
mha_sharding = True
|
||||
else:
|
||||
mha_sharding = n_heads_q == n_heads_kv
|
||||
|
||||
if n_heads_q < num_shards:
|
||||
raise ValueError("There must be at least as many query heads as there are shards.")
|
||||
|
||||
if mha_sharding:
|
||||
return shard_param(param,
|
||||
ShardingType.OUTER_DIMENSION,
|
||||
shard_rank,
|
||||
num_shards,
|
||||
num_concatenated_matrices=3,
|
||||
granularity=head_size)
|
||||
else:
|
||||
if n_heads_q % n_heads_kv != 0:
|
||||
raise ValueError("Must be an even ratio between query and key/value heads.")
|
||||
|
||||
if param.shape[0] != head_size * (n_heads_q + 2 * n_heads_kv):
|
||||
raise ValueError("GQA param shape is not correct")
|
||||
|
||||
# 32 KV heads, 16 shards for example
|
||||
if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0:
|
||||
raise ValueError("Currently do not support uneven partitioning of KV heads for GQA.")
|
||||
|
||||
# 8 KV heads, 16 shards for example
|
||||
if n_heads_kv < num_shards and num_shards % n_heads_kv != 0:
|
||||
raise ValueError("Currently do not support distributing KV heads across different numbers of shards.")
|
||||
else:
|
||||
even_kv_sharding = n_heads_kv >= num_shards
|
||||
|
||||
q_param = param[:head_size * n_heads_q]
|
||||
kv_param = param[head_size * n_heads_q:]
|
||||
|
||||
if even_kv_sharding:
|
||||
# This is equivalent to the original sharding algorithm since n_heads_q = C * n_heads_kv.
|
||||
# If n_heads_kv % num_shards == 0, then n_heads_q % num_shards == 0.
|
||||
q_param = shard_param(q_param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=head_size)
|
||||
kv_param = shard_param(kv_param,
|
||||
ShardingType.OUTER_DIMENSION,
|
||||
shard_rank,
|
||||
num_shards,
|
||||
num_concatenated_matrices=2,
|
||||
granularity=head_size)
|
||||
return torch.cat([q_param, kv_param], dim=0)
|
||||
else:
|
||||
# We will first do a sharding on the KV and Q to map to the one KV shard per group of Q.
|
||||
q_sharding_degree = num_shards // n_heads_kv
|
||||
|
||||
kv_head = shard_rank // q_sharding_degree
|
||||
k_param = kv_param[kv_head * head_size:(kv_head + 1) * head_size]
|
||||
v_param = kv_param[(n_heads_kv + kv_head) * head_size:(n_heads_kv + kv_head + 1) * head_size]
|
||||
|
||||
q_sharding_rank = shard_rank % q_sharding_degree
|
||||
q_factor = n_heads_q // n_heads_kv
|
||||
|
||||
q_chunk = q_param[q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size]
|
||||
|
||||
q_param = shard_param(q_chunk,
|
||||
ShardingType.OUTER_DIMENSION,
|
||||
q_sharding_rank,
|
||||
q_sharding_degree,
|
||||
granularity=head_size)
|
||||
|
||||
return torch.cat([q_param, k_param, v_param], dim=0)
|
||||
|
||||
|
||||
def qkv_out_features(in_features: int,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
head_size: int,
|
||||
n_heads_q: Optional[int] = None,
|
||||
n_heads_kv: Optional[int] = None) -> int:
|
||||
"""
|
||||
Helper to calculate the expected output projection dimension of a QKV projection matrix.
|
||||
|
||||
Args:
|
||||
in_features (int): The model dimension.
|
||||
shard_rank (int): Which rank to return the corresponding size for.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
head_size (int): The size of each head.
|
||||
n_heads_q (int): The number of query heads. This only needs to be passed if the number
|
||||
of query and key/value heads are different. If passed without n_heads_kv, default
|
||||
MHA partitioning will be used.
|
||||
n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number
|
||||
of query and key/value heads are different. This argument cannot be passed without also
|
||||
passing n_heads_q (we want to explicitly opt into GQA sharding).
|
||||
"""
|
||||
if n_heads_kv is not None and n_heads_q is None:
|
||||
raise ValueError("n_heads_kv should not be passed without n_heads_q")
|
||||
|
||||
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
|
||||
|
||||
if n_heads_q is not None and in_features != head_size * n_heads_q:
|
||||
raise ValueError("in_features is not consistent with n_heads_q and head_size")
|
||||
|
||||
if mha_sharding:
|
||||
endpoints = get_shard_endpoints(in_features, shard_rank, num_shards, granularity=head_size)
|
||||
return (endpoints[1] - endpoints[0]) * 3
|
||||
else:
|
||||
if n_heads_kv >= num_shards:
|
||||
if n_heads_kv % num_shards != 0:
|
||||
raise ValueError("The KV heads must be evenly distributed across the shards.")
|
||||
|
||||
n_local_groups = n_heads_kv // num_shards
|
||||
group_size = n_heads_q // n_heads_kv
|
||||
|
||||
return n_local_groups * head_size * (2 + group_size)
|
||||
else:
|
||||
if num_shards % n_heads_kv != 0:
|
||||
raise ValueError("A shared KV head must always partition across the same number of shards.")
|
||||
|
||||
q_split_degree = num_shards // n_heads_kv
|
||||
q_split_rank = shard_rank % q_split_degree
|
||||
split_granularity = (n_heads_q // n_heads_kv) * head_size
|
||||
|
||||
q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size)
|
||||
|
||||
return (q_endpoints[1] - q_endpoints[0]) + 2 * head_size
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from enum import Enum
|
||||
|
||||
DEFAULT_SHARD_GRANULARITY = 32
|
||||
|
||||
|
||||
class ShardingType(Enum):
|
||||
# Inner dimension sharding corresponds to splitting the Tensor along the K-dimension
|
||||
# of a matrix multiplication. This would be used for attention_output or MLP2.
|
||||
INNER_DIMENSION = 1
|
||||
|
||||
# Outer dimension sharding corresponds to splitting the Tensor along the N-dimension
|
||||
# of a matrix multiplication. This would be used for the QKV and MLP1 projections.
|
||||
OUTER_DIMENSION = 0
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType
|
||||
from .utils import shard_param, get_shard_endpoints
|
||||
|
||||
|
||||
def shard_unembed_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor:
|
||||
"""
|
||||
Utility method for sharding an unembed parameter. We shard unembeddings on the vocab dimension
|
||||
with the expectation of an all-gather to produce the full results.
|
||||
|
||||
TODO(cmikeh2): Really ideal would be if MII could have access to the comm and we would do
|
||||
an A2A and sharded sampling.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The sharded parameter of shape [sharded_vocab_size, model_dim]
|
||||
"""
|
||||
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=1)
|
||||
|
||||
|
||||
def sharded_unembed_dim(vocab_size: int, shard_rank: int, num_shards: int) -> int:
|
||||
"""
|
||||
Utility method for determining the sharded vocab size of a sharded unembed parameter.
|
||||
|
||||
Args:
|
||||
vocab_size (int): The size of the vocabulary.
|
||||
shard_rank (int): Which shard of the partitioned tensor to return.
|
||||
num_shards (int): The total number of shards the parameter is distributed across.
|
||||
"""
|
||||
start_idx, end_idx = get_shard_endpoints(vocab_size, shard_rank, num_shards, granularity=1)
|
||||
return end_idx - start_idx
|
||||
@@ -0,0 +1,104 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
|
||||
|
||||
|
||||
def get_shard_endpoints(dim_size: int,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
granularity: int = DEFAULT_SHARD_GRANULARITY) -> Tuple[int, int]:
|
||||
"""
|
||||
Given a dimension to shard with size dim_size, return the start and end indices of the slice
|
||||
that belong to the given rank.
|
||||
|
||||
The typical use of this is as an internal helper function, so see if there is a higher level
|
||||
API that better suits the application.
|
||||
|
||||
Args:
|
||||
dim_size (int): The size of the dimension to shard.
|
||||
shard_rank (int): The rank of the shard to return.
|
||||
num_shards (int): Total number of shards the dimension will be distributed across.
|
||||
granularity (int): The minimum alignment of the shard endpoints. This is used to support
|
||||
non-even head counts as well as align dimensions to cleaner GEMM boundaries.
|
||||
"""
|
||||
assert dim_size % granularity == 0, "Dimension size must be divisible by granularity"
|
||||
|
||||
total_chunks = dim_size // granularity
|
||||
base_chunks_per_rank = total_chunks // num_shards
|
||||
remainder_chunks = total_chunks % num_shards
|
||||
|
||||
start_chunk_id = shard_rank * base_chunks_per_rank + min(shard_rank, remainder_chunks)
|
||||
end_chunk_id = start_chunk_id + base_chunks_per_rank + (1 if shard_rank < remainder_chunks else 0)
|
||||
|
||||
return start_chunk_id * granularity, end_chunk_id * granularity
|
||||
|
||||
|
||||
def shard_param(param: Optional[torch.Tensor],
|
||||
shard_mode: ShardingType,
|
||||
shard_rank: int,
|
||||
num_shards: int,
|
||||
num_concatenated_matrices: int = 1,
|
||||
granularity: int = 32,
|
||||
bias_dims: int = 1) -> torch.Tensor:
|
||||
"""
|
||||
Utility for sharding a parameter. This will return the slice of the parameter that should
|
||||
exist on the given shard_rank given the sharding configuration. The workflow here is
|
||||
to find the minimum bounded Tensor to shard, get the slicing endpoints, and then concatenate
|
||||
as needed.
|
||||
|
||||
The typical use of this is as an internal helper function, so see if there is a higher level
|
||||
API that better suits the application.
|
||||
|
||||
Args:
|
||||
param (torch.Tensor): The parameter to shard.
|
||||
shard_mode (ShardingType): The type of sharding to apply. See ShardingType for more context.
|
||||
shard_rank (int): The rank of the shard to return.
|
||||
num_shards (int): Total number of shards the parameter will be distrbuted across.
|
||||
num_concatenated_matrices (int): The number of matrices that have been concatenated together in the original
|
||||
parameter. An example of this is a fused QKV projection matrix, where the `num_concatenated_matrices`
|
||||
argument would be 3.
|
||||
granularity (int): The minimum alignment of the shard endpoints. For attention projection matrices, this
|
||||
should be set to the head size to support non-even sharding.
|
||||
bias_dims (int): The number of dimensions that are considered bias dimensions. This is used to support
|
||||
sharding of MoE and non-MoE biases on the same codepath.
|
||||
"""
|
||||
assert shard_rank < num_shards, "Shard rank must be less than num_shards"
|
||||
|
||||
# Easier to hide this inside of the sharding logic than to add checks in every model
|
||||
# implementation.
|
||||
if param is None:
|
||||
return None
|
||||
|
||||
if num_shards == 1:
|
||||
# Trivial case of no sharding.
|
||||
return param
|
||||
|
||||
if shard_mode == ShardingType.OUTER_DIMENSION:
|
||||
|
||||
def get_matrices(dim_idx: int) -> torch.Tensor:
|
||||
dim_size = param.size(dim_idx) // num_concatenated_matrices
|
||||
start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
|
||||
return torch.chunk(param, num_concatenated_matrices, dim=dim_idx), start_channel_id, end_channel_id
|
||||
|
||||
if param.ndim == bias_dims:
|
||||
# Special case for bias parameters.
|
||||
matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-1)
|
||||
return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)
|
||||
else:
|
||||
# General case for weight parameters. This assumes MoE parameters are stored in the format of
|
||||
# [num_experts, out_features, in_features]
|
||||
matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-2)
|
||||
return torch.cat([mat[..., start_channel_id:end_channel_id, :] for mat in matrices], dim=-2)
|
||||
|
||||
elif shard_mode == ShardingType.INNER_DIMENSION:
|
||||
dim_size = param.size(-1) // num_concatenated_matrices
|
||||
start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
|
||||
matrices = torch.chunk(param, num_concatenated_matrices, dim=-1)
|
||||
return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)
|
||||
Reference in New Issue
Block a user