(self, hidden_size, num_attention_heads, attention_dropout_prob, output_dropout_prob, init_method,
output_layer_init_method=None)
| 58 | """Parallel cross-attention layer for Transformer""" |
| 59 | |
| 60 | def __init__(self, hidden_size, num_attention_heads, attention_dropout_prob, output_dropout_prob, init_method, |
| 61 | output_layer_init_method=None): |
| 62 | super(ParallelCrossAttention, self).__init__() |
| 63 | # Set output layer initialization if not provided. |
| 64 | if output_layer_init_method is None: |
| 65 | output_layer_init_method = init_method |
| 66 | # Per attention head and per partition values. |
| 67 | world_size = get_model_parallel_world_size() |
| 68 | self.hidden_size_per_partition = divide(hidden_size, world_size) |
| 69 | self.hidden_size_per_attention_head = divide(hidden_size, |
| 70 | num_attention_heads) |
| 71 | self.num_attention_heads_per_partition = divide(num_attention_heads, |
| 72 | world_size) |
| 73 | # Strided linear layer. |
| 74 | self.query = ColumnParallelLinear(hidden_size, hidden_size, |
| 75 | gather_output=False, |
| 76 | init_method=init_method) |
| 77 | self.key_value = ColumnParallelLinear(hidden_size, 2 * hidden_size, |
| 78 | stride=2, |
| 79 | gather_output=False, |
| 80 | init_method=init_method) |
| 81 | # Dropout. Note that for a single iteration, this layer will generate |
| 82 | # different outputs on different number of parallel partitions but |
| 83 | # on average it should not be partition dependent. |
| 84 | self.attention_dropout = torch.nn.Dropout(attention_dropout_prob) |
| 85 | |
| 86 | # Output. |
| 87 | self.dense = RowParallelLinear(hidden_size, |
| 88 | hidden_size, |
| 89 | input_is_parallel=True, |
| 90 | init_method=output_layer_init_method) |
| 91 | self.output_dropout = torch.nn.Dropout(output_dropout_prob) |
| 92 | |
| 93 | if deepspeed.checkpointing.is_configured(): |
| 94 | global get_cuda_rng_tracker, checkpoint |
| 95 | get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker |
| 96 | checkpoint = deepspeed.checkpointing.checkpoint |
| 97 | |
| 98 | def _transpose_for_scores(self, tensor): |
| 99 | """Transpose a 3D tensor [b, s, np*hn] into a 4D tensor with |
nothing calls this directly
no test coverage detected