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