Parallel top query self-attention layer abstract class. Self-attention layer takes input with size [b, s, h] and returns output of the same size.
| 323 | |
| 324 | |
| 325 | class ParallelTopQuerySelfAttention(MegatronModule): |
| 326 | """Parallel top query self-attention layer abstract class. |
| 327 | |
| 328 | Self-attention layer takes input with size [b, s, h] |
| 329 | and returns output of the same size. |
| 330 | """ |
| 331 | |
| 332 | def __init__(self, init_method, |
| 333 | output_layer_init_method, layer_number): |
| 334 | super(ParallelTopQuerySelfAttention, self).__init__() |
| 335 | args = get_args() |
| 336 | self.fp16 = args.fp16 |
| 337 | self.attention_softmax_in_fp32 = args.attention_softmax_in_fp32 |
| 338 | self.layer_number = max(1, layer_number) |
| 339 | |
| 340 | if hasattr(args, 'attention_upweight_top'): |
| 341 | self.attention_upweight = args.attention_upweight_top |
| 342 | else: |
| 343 | self.attention_upweight = None |
| 344 | # Per attention head and per partition values. |
| 345 | world_size = mpu.get_model_parallel_world_size() |
| 346 | self.hidden_size_per_partition = mpu.divide(args.hidden_size, |
| 347 | world_size) |
| 348 | self.hidden_size_per_attention_head = mpu.divide( |
| 349 | args.hidden_size, args.num_attention_heads) |
| 350 | self.num_attention_heads_per_partition = mpu.divide( |
| 351 | args.num_attention_heads, world_size) |
| 352 | |
| 353 | self.query = mpu.ColumnParallelLinear( |
| 354 | args.hidden_size, |
| 355 | args.hidden_size, |
| 356 | gather_output=False, |
| 357 | init_method=init_method) |
| 358 | |
| 359 | self.key = mpu.ColumnParallelLinear( |
| 360 | args.hidden_size, |
| 361 | args.hidden_size, |
| 362 | gather_output=False, |
| 363 | init_method=init_method) |
| 364 | |
| 365 | self.value = mpu.ColumnParallelLinear( |
| 366 | args.hidden_size, |
| 367 | args.hidden_size, |
| 368 | gather_output=False, |
| 369 | init_method=init_method) |
| 370 | |
| 371 | self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) |
| 372 | self.softmax = torch.nn.Softmax(dim=-1) |
| 373 | |
| 374 | # Dropout. Note that for a single iteration, this layer will generate |
| 375 | # different outputs on different number of parallel partitions but |
| 376 | # on average it should not be partition dependent. |
| 377 | self.attention_dropout = torch.nn.Dropout(args.attention_dropout) |
| 378 | |
| 379 | # Output. |
| 380 | self.dense = mpu.RowParallelLinear( |
| 381 | args.hidden_size, |
| 382 | args.hidden_size, |