wehub-resource-sync
28 lines
955 B
Python
28 lines
955 B
Python
# 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
|