mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-28 04:55:25 +00:00
[shardformer] add gpt2 test and layer class refactor (#4041)
* add gpt2 test and layer class refactor * add dropout in gpt2 policy
This commit is contained in:
FoolPlayer
Frank Lee
@@ -1,126 +1,101 @@
|
||||
from typing import Any, Callable, Dict, List, Tuple, Type
|
||||
from typing import Type, Union
|
||||
|
||||
import torch.nn as nn
|
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2Block, GPT2Model
|
||||
|
||||
import colossalai.shardformer.layer.layers as col_nn
|
||||
import colossalai.shardformer.layer as col_nn
|
||||
from colossalai.shardformer.layer.dropout import Dropout1D
|
||||
|
||||
from .basepolicy import Argument, Col_Layer, Layer, Policy, Row_Layer
|
||||
from ..utils import getattr_, setattr_
|
||||
from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
|
||||
|
||||
|
||||
class GPT2Policy(Policy):
|
||||
|
||||
@staticmethod
|
||||
def argument_policy(config, world_size):
|
||||
def preprocess(self):
|
||||
# reshape the embedding layer
|
||||
r"""
|
||||
Reshape the Embedding layer to make the embedding dimension divisible by world_size
|
||||
"""
|
||||
vocab_size = self.model.config.vocab_size
|
||||
world_size = self.shard_config.tensor_parallel_size
|
||||
if vocab_size % world_size != 0:
|
||||
new_vocab_size = vocab_size + world_size - vocab_size % world_size
|
||||
self.model.resize_token_embeddings(new_vocab_size)
|
||||
return self.model
|
||||
|
||||
def module_policy(self):
|
||||
return {
|
||||
GPT2Model:
|
||||
Argument(attr_dict={}, param_funcs=[
|
||||
GPT2Policy.embedding,
|
||||
]),
|
||||
ModulePolicyDescription(attribute_replacement={},
|
||||
param_replacement=[],
|
||||
sub_module_replacement=[
|
||||
SubModuleReplacementDescription(
|
||||
suffix="wte",
|
||||
target_module=col_nn.VocabParallelEmbedding1D,
|
||||
),
|
||||
]),
|
||||
GPT2Block:
|
||||
Argument(
|
||||
attr_dict={
|
||||
# 1. reduce hidden size
|
||||
"attn.embed_dim": config.hidden_size // world_size,
|
||||
"attn.split_size": config.hidden_size // world_size,
|
||||
"crossattention.embed_dim": config.hidden_size // world_size,
|
||||
"crossattention.split_size": config.hidden_size // world_size,
|
||||
# 2. reduce number of heads
|
||||
"attn.num_heads": config.num_attention_heads // world_size,
|
||||
"crossattention.num_heads": config.num_attention_heads // world_size,
|
||||
},
|
||||
param_funcs=[
|
||||
GPT2Policy.attn_in,
|
||||
GPT2Policy.attn_out,
|
||||
GPT2Policy.mlp_in,
|
||||
GPT2Policy.mlp_out,
|
||||
]),
|
||||
ModulePolicyDescription(attribute_replacement={
|
||||
"attn.embed_dim": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
|
||||
"attn.split_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
|
||||
"attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
|
||||
},
|
||||
param_replacement=[],
|
||||
sub_module_replacement=[
|
||||
SubModuleReplacementDescription(
|
||||
suffix="attn.c_attn",
|
||||
target_module=col_nn.LinearConv1D_Col,
|
||||
kwargs={
|
||||
"n_cast": 3,
|
||||
},
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="attn.c_proj",
|
||||
target_module=col_nn.LinearConv1D_Row,
|
||||
kwargs={
|
||||
"n_cast": 1,
|
||||
},
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="mlp.c_fc",
|
||||
target_module=col_nn.LinearConv1D_Col,
|
||||
kwargs={
|
||||
"n_cast": 1,
|
||||
},
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="mlp.c_proj",
|
||||
target_module=col_nn.LinearConv1D_Row,
|
||||
kwargs={
|
||||
"n_cast": 1,
|
||||
},
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="attn.attn_dropout",
|
||||
target_module=col_nn.Dropout1D,
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="attn.resid_dropout",
|
||||
target_module=col_nn.Dropout1D,
|
||||
),
|
||||
SubModuleReplacementDescription(
|
||||
suffix="mlp.dropout",
|
||||
target_module=col_nn.Dropout1D,
|
||||
),
|
||||
])
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def attn_in() -> List:
|
||||
return [
|
||||
Col_Layer(suffix="attn.c_attn",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
n_cast=3,
|
||||
reversed=True,
|
||||
replace_layer=col_nn.Linear1D_Col),
|
||||
Col_Layer(suffix="crossattention.c_attn",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
n_cast=2,
|
||||
reversed=True,
|
||||
ignore=True,
|
||||
replace_layer=col_nn.Linear1D_Col),
|
||||
Col_Layer(suffix="crossattention.q_attn",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
reversed=True,
|
||||
ignore=True,
|
||||
replace_layer=col_nn.Linear1D_Col)
|
||||
]
|
||||
def new_model_class(self):
|
||||
|
||||
@staticmethod
|
||||
def attn_out() -> List:
|
||||
return [
|
||||
Row_Layer(suffix="attn.c_proj",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
reversed=True,
|
||||
replace_layer=col_nn.Linear1D_Row),
|
||||
Row_Layer(suffix="crossattention.c_proj",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
reversed=True,
|
||||
ignore=True,
|
||||
replace_layer=col_nn.Linear1D_Row)
|
||||
]
|
||||
return self.model
|
||||
|
||||
@staticmethod
|
||||
def mlp_in() -> List:
|
||||
return [
|
||||
Col_Layer(suffix="mlp.c_fc", weight="weight", bias="bias", reversed=True,
|
||||
replace_layer=col_nn.Linear1D_Col),
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def mlp_out() -> List:
|
||||
return [
|
||||
Row_Layer(suffix="mlp.c_proj",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
reversed=True,
|
||||
replace_layer=col_nn.Linear1D_Row)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def embedding() -> List:
|
||||
return [Col_Layer(suffix="wte", weight="weight", replace_layer=col_nn.VocabParallelEmbedding1D)]
|
||||
def postprocess(self):
|
||||
return self.model
|
||||
|
||||
|
||||
from transformers import GPT2LMHeadModel
|
||||
# GPT2Model
|
||||
class GPT2ModelPolicy(GPT2Policy):
|
||||
|
||||
|
||||
class GPT2LMHeadModelPolicy(GPT2Policy):
|
||||
|
||||
@staticmethod
|
||||
def argument_policy(config, world_size):
|
||||
base_argument = GPT2Policy.argument_policy(config, world_size)
|
||||
argument = {
|
||||
GPT2LMHeadModel: Argument(attr_dict={}, param_funcs=[
|
||||
GPT2LMHeadModelPolicy.unembedding,
|
||||
]),
|
||||
}
|
||||
argument.update(base_argument)
|
||||
return argument
|
||||
|
||||
@staticmethod
|
||||
def unembedding() -> List:
|
||||
return [
|
||||
Col_Layer(suffix="lm_head",
|
||||
weight="weight",
|
||||
bias="bias",
|
||||
replace_layer=col_nn.Linear1D_Col,
|
||||
gather_output=True)
|
||||
]
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
Reference in New Issue
Block a user