chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,113 @@
+# 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)