(self, hidden_size, num_attention_heads,
attention_dropout_prob, output_dropout_prob,
init_method, layer_id, hidden_size_per_attention_head=None, output_layer_init_method=None, bias=True, qkv_bias=False, num_multi_query_heads=0, row_parallel_linear_final_bias=True,
hooks={}, transformer_pointer=None, params_dtype=torch.float, skip_init=False, device=torch.device('cpu'))
| 33 | |
| 34 | class SelfAttention(torch.nn.Module): |
| 35 | def __init__(self, hidden_size, num_attention_heads, |
| 36 | attention_dropout_prob, output_dropout_prob, |
| 37 | init_method, layer_id, hidden_size_per_attention_head=None, output_layer_init_method=None, bias=True, qkv_bias=False, num_multi_query_heads=0, row_parallel_linear_final_bias=True, |
| 38 | hooks={}, transformer_pointer=None, params_dtype=torch.float, skip_init=False, device=torch.device('cpu')): |
| 39 | super(SelfAttention, self).__init__() |
| 40 | # Set output layer initialization if not provided. |
| 41 | if output_layer_init_method is None: |
| 42 | output_layer_init_method = init_method |
| 43 | self.hooks = hooks |
| 44 | self.layer_id = layer_id |
| 45 | # Per attention head and per partition values. |
| 46 | world_size = get_model_parallel_world_size() |
| 47 | self.hidden_size = hidden_size |
| 48 | self.num_attention_heads = num_attention_heads |
| 49 | self.num_multi_query_heads = num_multi_query_heads |
| 50 | if hidden_size_per_attention_head is None: |
| 51 | self.hidden_size_per_attention_head = divide(hidden_size, num_attention_heads) |
| 52 | else: |
| 53 | self.hidden_size_per_attention_head = hidden_size_per_attention_head |
| 54 | self.num_attention_heads_per_partition = divide(num_attention_heads, world_size) |
| 55 | self.num_multi_query_heads_per_partition = divide(num_multi_query_heads, world_size) |
| 56 | self.inner_hidden_size = num_attention_heads * self.hidden_size_per_attention_head |
| 57 | self.hidden_size_per_partition = self.hidden_size_per_attention_head * self.num_attention_heads_per_partition |
| 58 | |
| 59 | # Strided linear layer. |
| 60 | if num_multi_query_heads == 0: |
| 61 | qkv_size = 3 * self.inner_hidden_size |
| 62 | self.stride = 3 |
| 63 | else: # multi-query |
| 64 | qkv_size = self.inner_hidden_size + self.hidden_size_per_attention_head * self.num_multi_query_heads * 2 |
| 65 | self.stride = [self.num_attention_heads_per_partition, self.num_multi_query_heads_per_partition, self.num_multi_query_heads_per_partition] |
| 66 | self.query_key_value = ColumnParallelLinear( |
| 67 | hidden_size, |
| 68 | qkv_size, |
| 69 | stride=self.stride, |
| 70 | gather_output=False, |
| 71 | init_method=init_method, |
| 72 | bias=bias or qkv_bias, |
| 73 | params_dtype=params_dtype, |
| 74 | module=self, |
| 75 | name="query_key_value", |
| 76 | skip_init=skip_init, |
| 77 | device=device |
| 78 | ) |
| 79 | self.attention_dropout = torch.nn.Dropout(attention_dropout_prob) |
| 80 | |
| 81 | self.dense = RowParallelLinear( |
| 82 | self.inner_hidden_size, |
| 83 | hidden_size, |
| 84 | input_is_parallel=True, |
| 85 | init_method=output_layer_init_method, |
| 86 | bias=bias, |
| 87 | params_dtype=params_dtype, |
| 88 | module=self, |
| 89 | name="dense", |
| 90 | skip_init=skip_init, |
| 91 | device=device, |
| 92 | final_bias=row_parallel_linear_final_bias |
nothing calls this directly
no test coverage detected