wehub-resource-sync
120 lines
5.3 KiB
Python
120 lines
5.3 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from .base import *
|
|
from .features.megatron import MegatronContainer
|
|
from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference
|
|
import torch
|
|
from ..policy import TransformerPolicy
|
|
from packaging import version as pkg_version
|
|
|
|
|
|
class DS_MegatronGPTContainer(MegatronContainer, BaseTransformerContainer):
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
|
|
# All model specific things should be defined here instead of the base class.
|
|
|
|
def create_module(self, config=None):
|
|
_config = config if config is not None else self.ds_model_config
|
|
self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group)
|
|
self.module.config.scale_attention = self.scale_attention
|
|
|
|
if self.megatron_v2:
|
|
self.module.config.rotate_half = True
|
|
self.module.config.rotate_every_two = False
|
|
|
|
return self.module
|
|
|
|
|
|
# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp
|
|
# TODO: Generalize MoE overall goal, expand beyond Megatron
|
|
class MegatronLayerPolicy(TransformerPolicy):
|
|
_orig_layer_class = None
|
|
version = 0
|
|
moe_type = 'standard'
|
|
megatron_v2 = True
|
|
use_mup = False
|
|
|
|
def __init__(self, client_module, inference=True):
|
|
super().__init__(inference, megatron_v2=MegatronLayerPolicy.megatron_v2, use_mup=MegatronLayerPolicy.use_mup)
|
|
self.client_module = client_module
|
|
# we use megatron version to differentiate between the old and new
|
|
# megatron-lm source code
|
|
if MegatronLayerPolicy._orig_layer_class is None:
|
|
if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"):
|
|
MegatronLayerPolicy._orig_layer_class = None
|
|
else:
|
|
try:
|
|
from megatron.model.transformer import ParallelTransformerLayer
|
|
MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer
|
|
MegatronLayerPolicy.version = 1
|
|
except ImportError:
|
|
MegatronLayerPolicy._orig_layer_class = None
|
|
|
|
def get_hidden_heads(self):
|
|
if MegatronLayerPolicy.version == 0:
|
|
return self.client_module.attention.query_key_value.weight.shape[1], \
|
|
self.client_module.attention.num_attention_heads, \
|
|
self.client_module.input_layernorm.eps, \
|
|
DEFAULT_INTERMEDIATE_SIZE
|
|
else:
|
|
return self.client_module.self_attention.query_key_value.weight.shape[1], \
|
|
self.client_module.self_attention.num_attention_heads, \
|
|
self.client_module.input_layernorm.eps, \
|
|
DEFAULT_INTERMEDIATE_SIZE
|
|
|
|
def attention(self, enable_training=False):
|
|
if self.inference:
|
|
if MegatronLayerPolicy.version == 0:
|
|
attention = self.client_module.attention
|
|
else:
|
|
attention = self.client_module.self_attention
|
|
else:
|
|
return None
|
|
|
|
return attention.query_key_value.weight, \
|
|
attention.query_key_value.bias, \
|
|
attention.dense.weight, \
|
|
attention.dense.bias
|
|
|
|
def mlp(self, moe_type='standard', enable_training=False):
|
|
from deepspeed.moe.utils import has_moe_layers
|
|
moe, _ = has_moe_layers(self.client_module)
|
|
|
|
if moe:
|
|
moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \
|
|
self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts
|
|
num_experts = len(moe_experts)
|
|
if moe_type == 'standard':
|
|
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
|
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
|
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
|
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)]
|
|
else:
|
|
|
|
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
|
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
|
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
|
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \
|
|
self.client_module.mlp.mlp.dense_h_to_4h.weight, \
|
|
self.client_module.mlp.mlp.dense_h_to_4h.bias, \
|
|
self.client_module.mlp.mlp.dense_4h_to_h.weight, \
|
|
self.client_module.mlp.mlp.dense_4h_to_h.bias, \
|
|
self.client_module.mlp.coefficient.weight
|
|
|
|
else:
|
|
return self.client_module.mlp.dense_h_to_4h.weight, \
|
|
self.client_module.mlp.dense_h_to_4h.bias, \
|
|
self.client_module.mlp.dense_4h_to_h.weight, \
|
|
self.client_module.mlp.dense_4h_to_h.bias
|
|
|
|
def layernorm(self):
|
|
return self.client_module.post_attention_layernorm.weight, \
|
|
self.client_module.post_attention_layernorm.bias, \
|
|
self.client_module.input_layernorm.weight, \
|
|
self.client_module.input_layernorm.bias
|