chore: import upstream snapshot with attribution

deepspeed/runtime/data_pipeline/data_routing/utils.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+
+def bsh_decoder_gather(reserved_length, hidden_states, mask):
+    # random-layer-token-drop
+    rand_list = []
+    part_hidden_states = []  #  batch, seq, hidden ## different from megatron
+    for k in range(hidden_states.size(0)):
+        B_tmp = torch.randperm(hidden_states.size(1), device=hidden_states.device)[:reserved_length]
+        B = B_tmp.sort()[0]
+        rand_list.append(B)
+        part_hidden_states.append(hidden_states[k:k + 1, B, :])
+
+    part_hidden_states = torch.cat(part_hidden_states, dim=0)
+    part_mask = mask[:, :, :reserved_length, :reserved_length]
+    return part_hidden_states, rand_list, part_mask
+
+
+def bsh_decoder_scatter(hidden_states, part_hidden_states, rand_list):
+    for k in range(hidden_states.size(0)):
+        hidden_states[k, rand_list[k], :] = part_hidden_states[k, :, :]
+    return hidden_states