wehub-resource-sync
114 lines
5.5 KiB
Python
114 lines
5.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from deepspeed.utils import logger
|
|
from torch import Tensor
|
|
from torch.nn import Module
|
|
from ..constants import *
|
|
from deepspeed.ops.random_ltd.dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens
|
|
|
|
|
|
#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586
|
|
class RandomLayerTokenDrop(Module):
|
|
"""
|
|
A layer wrapper for random LTD
|
|
"""
|
|
|
|
def __init__(self, layer: Module):
|
|
super(RandomLayerTokenDrop, self).__init__()
|
|
self.random_ltd_layer = layer
|
|
self.reserved_length = None #config['max_value']
|
|
self.random_ltd_scheduler = None
|
|
self.max_length = None
|
|
self.reserved_length = -1
|
|
self.curr_seq = -1
|
|
self.batch_first = False
|
|
|
|
def init_config(self, config, scheduler, random_ltd_layer_id):
|
|
self.random_ltd_scheduler = scheduler
|
|
self.random_ltd_layer_id = random_ltd_layer_id
|
|
self.max_length = self.random_ltd_scheduler.state[RANDOM_LTD_MAX_VALUE]
|
|
|
|
self.mask_name = config[RANDOM_LTD_MODEL_MASK_NAME]
|
|
self.micro_bs = config[RANDOM_LTD_MICRO_BATCH_SIZE]
|
|
self.random_ltd_num_layer = self.random_ltd_scheduler.random_ltd_layer_num
|
|
hs_order = config[RANDOM_LTD_HIDDEN_STATE_ORDER]
|
|
self.model_type = config[RANDOM_LTD_MODEL_TYPE]
|
|
|
|
if hs_order == 'batch_seq_dim':
|
|
self.get_hidden_tensor_shape = self.get_bsh
|
|
self.batch_first = True
|
|
elif hs_order == 'seq_batch_dim':
|
|
self.get_hidden_tensor_shape = self.get_sbh
|
|
self.batch_first = False
|
|
else:
|
|
logger.warning(
|
|
"************For now, we only support batch_seq_dim or seq_batch_dim inputs. You can easily \
|
|
your own input dimension orders************")
|
|
raise NotImplementedError
|
|
|
|
if self.model_type == 'encoder':
|
|
self.index_generator = bert_sample_tokens
|
|
elif self.model_type == 'decoder':
|
|
self.index_generator = gpt_sample_tokens
|
|
else:
|
|
logger.warning("************For now, we only support encoder-only or decoder-only models************")
|
|
raise NotImplementedError
|
|
|
|
def get_bsh(self, hidden_stats):
|
|
self.curr_seq, self.curr_micro_batch = hidden_stats.size()[1], hidden_stats.size()[0]
|
|
|
|
def get_sbh(self, hidden_stats):
|
|
self.curr_seq, self.curr_micro_batch = hidden_stats.size()[0], hidden_stats.size()[1]
|
|
|
|
def forward(self, hidden_states, **kwargs) -> Tensor:
|
|
if self.random_ltd_scheduler is not None:
|
|
self.reserved_length = self.random_ltd_scheduler.get_current_seq()
|
|
self.get_hidden_tensor_shape(hidden_states)
|
|
if self.training and self.random_ltd_scheduler is not None and self.reserved_length < self.curr_seq:
|
|
if self.mask_name is not None:
|
|
mask = kwargs[self.mask_name]
|
|
else:
|
|
mask = None
|
|
if self.random_ltd_layer_id == 0:
|
|
sampled_indices, part_attention_mask = self.index_generator(self.reserved_length,\
|
|
self.curr_seq, \
|
|
self.curr_micro_batch, \
|
|
self.random_ltd_num_layer, \
|
|
hidden_states.device, mask)
|
|
self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] = sampled_indices
|
|
self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] = part_attention_mask
|
|
else:
|
|
sampled_indices = self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX]
|
|
part_attention_mask = self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK]
|
|
|
|
hidden_states, part_hidden_states = GatherTokens.apply(hidden_states,
|
|
sampled_indices[self.random_ltd_layer_id, :, :],
|
|
self.batch_first)
|
|
if self.mask_name is not None:
|
|
if self.model_type == 'encoder':
|
|
kwargs[self.mask_name] = part_attention_mask[self.random_ltd_layer_id]
|
|
else:
|
|
kwargs[self.mask_name] = part_attention_mask
|
|
|
|
outputs = self.random_ltd_layer(part_hidden_states, **kwargs)
|
|
|
|
if isinstance(outputs, tuple):
|
|
hidden_states = ScatterTokens.apply(hidden_states, outputs[0],
|
|
sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first)
|
|
my_list = list(outputs)
|
|
my_list[0] = hidden_states
|
|
return tuple(my_list)
|
|
elif isinstance(outputs, Tensor):
|
|
hidden_states = ScatterTokens.apply(hidden_states, outputs,
|
|
sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first)
|
|
return hidden_states
|
|
else:
|
|
logger.warning("************For now, we only support tuple and tensor output. \
|
|
You need to adjust the output according to the layer in your model************")
|
|
raise NotImplementedError
|
|
else:
|
|
return self.random_ltd_layer(hidden_states, **kwargs)
|