megatron_patch/model/llama2/transformer.py [574:962]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - gather_output=False) else: assert attention_type == AttnType.cross_attn if self.group_query_attention: raise NotImplementedError("Grouped query attention not implemented for cross-attention.") assert query_projection_size == kv_projection_size self.query = tensor_parallel.ColumnParallelLinear( config.hidden_size, query_projection_size, config=config, init_method=config.init_method, bias=config.add_bias_linear, gather_output=False) self.key_value = tensor_parallel.ColumnParallelLinear( config.hidden_size, 2 * kv_projection_size, config=config, init_method=config.init_method, bias=config.add_bias_linear, gather_output=False) self.core_attention = CoreAttention(self.layer_number, config, self.attn_mask_type) self.checkpoint_core_attention = config.recompute_granularity == 'selective' if self.use_flash_attn: self.core_attention_flash = FlashSelfAttention( causal=True, attention_dropout=config.attention_dropout ) # Output. self.dense = tensor_parallel.RowParallelLinear( query_projection_size, config.hidden_size, config=config, init_method=config.output_layer_init_method, bias=args.add_bias_linear, input_is_parallel=True, skip_bias_add=True) if args.use_llama2_rotary_position_embeddings: self.use_llama2_rotary_position_embeddings = True self.seq_length = args.seq_length rotary_dim = args.hidden_size // args.num_attention_heads \ if args.kv_channels is None else args.kv_channels if args.rotary_percent < 1.0: rotary_dim = int(rotary_dim * args.rotary_percent) # partial rotary embeddings, which is better than full rotary # Wang and Komatsuzaki et al # https://github.com/kingoflolz/mesh-transformer-jax/ self.rotary_emb = RotaryEmbedding( rotary_dim, args.max_position_embeddings, args.rotary_base, args.rotary_scale_factor ) else: self.use_llama2_rotary_position_embeddings = False def _checkpointed_attention_forward(self, query_layer, key_layer, value_layer, attention_mask, rotary_pos_emb=None): """Forward method with activation checkpointing.""" def custom_forward(*inputs): query_layer = inputs[0] key_layer = inputs[1] value_layer = inputs[2] attention_mask = inputs[3] output_ = self.core_attention(query_layer, key_layer, value_layer, attention_mask) return output_ q_pos_emb, k_pos_emb = (None, None) if rotary_pos_emb is None \ else rotary_pos_emb hidden_states = tensor_parallel.checkpoint( custom_forward, False, query_layer, key_layer, value_layer, attention_mask, q_pos_emb, k_pos_emb) return hidden_states def _allocate_memory(self, inference_max_sequence_len, batch_size, num_attention_heads): return torch.empty( inference_max_sequence_len, batch_size, num_attention_heads, self.hidden_size_per_attention_head, dtype=self.params_dtype, device=torch.cuda.current_device()) def forward(self, hidden_states, attention_mask, encoder_output=None, inference_params=None, rotary_pos_emb=None, position_ids=None): # hidden_states: [sq, b, h] # ================================================= # Pre-allocate memory for key-values for inference. # ================================================= is_first_step = False if inference_params: if self.layer_number not in inference_params.key_value_memory_dict: inf_max_seq_len = inference_params.max_sequence_length inf_max_batch_size = inference_params.max_batch_size inference_key_memory = self._allocate_memory( inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition) inference_value_memory = self._allocate_memory( inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition) inference_params.key_value_memory_dict[self.layer_number] = ( inference_key_memory, inference_value_memory) is_first_step = True else: inference_key_memory, inference_value_memory = \ inference_params.key_value_memory_dict[self.layer_number] # ===================== # Query, Key, and Value # ===================== if self.attention_type == AttnType.self_attn: # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] mixed_x_layer, _ = self.query_key_value(hidden_states) # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn] new_tensor_shape = mixed_x_layer.size()[:-1] + ( self.num_query_groups_per_partition, ( (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) * self.hidden_size_per_attention_head ), ) mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) # [sq, b, ng, (np/ng + 2) * hn] --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query_layer, key_layer, value_layer) = torch.split( mixed_x_layer, [ ( self.num_attention_heads_per_partition // self.num_query_groups_per_partition * self.hidden_size_per_attention_head ), self.hidden_size_per_attention_head, self.hidden_size_per_attention_head ], dim=3) # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] - query_layer = query_layer.contiguous().view(query_layer.size(0), query_layer.size(1), -1, self.hidden_size_per_attention_head) else: # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)] mixed_kv_layer, _ = self.key_value(encoder_output) # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn] new_tensor_shape = mixed_kv_layer.size()[:-1] + \ (self.num_attention_heads_per_partition, 2 * self.hidden_size_per_attention_head) mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn] (key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_kv_layer, 2) # Attention head [sq, b, h] --> [sq, b, hp] query_layer, _ = self.query(hidden_states) # [sq, b, hp] --> [sq, b, np, hn] new_tensor_shape = query_layer.size()[:-1] + \ (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) query_layer = query_layer.view(*new_tensor_shape) # ================================== # Adjust key and value for inference # ================================== # duplicate the pos_emb for self attention if rotary_pos_emb is not None: if isinstance(rotary_pos_emb, tuple): rotary_pos_emb = rotary_pos_emb else: rotary_pos_emb = ((rotary_pos_emb,) * 2) if inference_params: if self.use_llama2_rotary_position_embeddings: kv_seq_len = key_layer.shape[0] kv_seq_len += inference_params.sequence_len_offset value_layer = value_layer.transpose(0, 1).transpose(1, 2) query_layer = query_layer.transpose(0, 1).transpose(1, 2) key_layer = key_layer.transpose(0, 1).transpose(1, 2) cos, sin = self.rotary_emb(value_layer, kv_seq_len) query_layer, key_layer = apply_llama2_rotary_pos_emb( query_layer, key_layer, cos, sin, position_ids) value_layer = value_layer.transpose(1, 2).transpose(0, 1) query_layer = query_layer.transpose(1, 2).transpose(0, 1) key_layer = key_layer.transpose(1, 2).transpose(0, 1) batch_start = inference_params.batch_size_offset batch_end = batch_start + key_layer.size(1) assert batch_end <= inference_key_memory.size(1) sequence_start = inference_params.sequence_len_offset sequence_end = sequence_start + key_layer.size(0) assert sequence_end <= inference_key_memory.size(0) # Copy key and values. inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = key_layer inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = value_layer key_layer = inference_key_memory[ :sequence_end, batch_start:batch_end, ...] value_layer = inference_value_memory[ :sequence_end, batch_start:batch_end, ...] # adjust the key rotary positional embedding if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb # need to cross check this condition during inference # if not set_inference_key_value_memory: if not is_first_step: # In inference, we compute one token at a time. # Select the correct positional embedding # (only the last token in the sequence) q_pos_emb = q_pos_emb[sequence_end - 1 : sequence_end] else: # In the first forward pass of inference, # we use the entire provided prefix. # q_pos_emb here has the rope embeddings of the entire # prefix + to-be-generated output so # we slice to just the prefix. q_pos_emb = q_pos_emb[:sequence_end, :, :, :] k_pos_emb = k_pos_emb[:sequence_end, :, :, :] rotary_pos_emb = (q_pos_emb, k_pos_emb) # ================================== # core attention computation # ================================== # expand the key_layer and value_layer [sk, b, ng, hn] -> [sk, b, np, hn] key_layer = key_layer.repeat_interleave( self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim = 2 ) value_layer = value_layer.repeat_interleave( self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim = 2 ) # apply relative positional encoding (rotary embedding) if self.use_llama2_rotary_position_embeddings: kv_seq_len = key_layer.shape[0] value_layer = value_layer.transpose(0, 1).transpose(1, 2) query_layer = query_layer.transpose(0, 1).transpose(1, 2) key_layer = key_layer.transpose(0, 1).transpose(1, 2) cos, sin = self.rotary_emb(value_layer, kv_seq_len) query_layer, key_layer = apply_llama2_rotary_pos_emb( query_layer, key_layer, cos, sin, position_ids) value_layer = value_layer.transpose(1, 2).transpose(0, 1) query_layer = query_layer.transpose(1, 2).transpose(0, 1) key_layer = key_layer.transpose(1, 2).transpose(0, 1) else: # apply relative positional encoding (rotary embedding) if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb) key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb) # TODO, can apply positional embedding to value_layer so it has # absolute positional embedding. # otherwise, only relative positional embedding takes effect # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb) if not self.use_flash_attn: if self.checkpoint_core_attention: context_layer = self._checkpointed_attention_forward( query_layer, key_layer, value_layer, attention_mask) else: context_layer = self.core_attention( query_layer, key_layer, value_layer, attention_mask) else: q, k, v = [rearrange(x, 's b ... -> b s ...').contiguous() for x in (query_layer, key_layer, value_layer)] if not self.sequence_parallel: with tensor_parallel.get_cuda_rng_tracker().fork(): context_layer = self.core_attention_flash(q, k, v) else: context_layer = self.core_attention_flash(q, k, v) context_layer = rearrange(context_layer, 'b s h d -> s b (h d)').contiguous() # ================= # Output. [sq, b, h] # ================= output, bias = self.dense(context_layer) return output, bias def bias_dropout_add(x, bias, residual, prob, training): # type: (Tensor, Optional[Tensor], Tensor, float, bool) -> Tensor if bias is not None: x = x + bias out = torch.nn.functional.dropout(x, p=prob, training=training) out = residual + out return out def get_bias_dropout_add(training): def _bias_dropout_add(x, bias, residual, prob): return bias_dropout_add(x, bias, residual, prob, training) return _bias_dropout_add @torch.jit.script def bias_dropout_add_fused_train(x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, True) @torch.jit.script def bias_dropout_add_fused_inference(x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, False) class ParallelTransformerLayer(MegatronModule): """A single transformer layer. Transformer layer takes input with size [s, b, h] and returns an output of the same size. """ def __init__(self, config, layer_number, layer_type=LayerType.encoder, self_attn_mask_type=AttnMaskType.padding, drop_path_rate=0.): # retriever=None): args = get_args() super(ParallelTransformerLayer, self).__init__() self.layer_number = layer_number self.layer_type = layer_type self.apply_residual_connection_post_norm \ = config.apply_residual_connection_post_layernorm self.bf16 = config.bf16 self.fp32_residual_connection = config.fp32_residual_connection # Normalize the input data. self.input_norm = get_norm(config) # Self attention. self.self_attention = ParallelAttention( config, layer_number, attention_type=AttnType.self_attn, attn_mask_type=self_attn_mask_type) self.hidden_dropout = config.hidden_dropout self.bias_dropout_fusion = config.bias_dropout_fusion self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None # Normalize the attention output self.post_attention_norm = get_norm(config) # Cross attention. if self.layer_type in (LayerType.decoder, LayerType.retro_decoder, LayerType.retro_decoder_with_retriever, LayerType.retro_encoder): self.inter_attention = ParallelAttention( config, layer_number, attention_type=AttnType.cross_attn) # Normalize the attention output. self.post_inter_attention_norm = get_norm(config) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - megatron_patch/model/qwen1_5_megablocks/transformer.py [465:852]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - gather_output=False) else: assert attention_type == AttnType.cross_attn if self.group_query_attention: raise NotImplementedError("Grouped query attention not implemented for cross-attention.") assert query_projection_size == kv_projection_size self.query = tensor_parallel.ColumnParallelLinear( config.hidden_size, query_projection_size, config=config, init_method=config.init_method, bias=config.add_bias_linear, gather_output=False) self.key_value = tensor_parallel.ColumnParallelLinear( config.hidden_size, 2 * kv_projection_size, config=config, init_method=config.init_method, bias=config.add_bias_linear, gather_output=False) self.core_attention = CoreAttention(self.layer_number, config, self.attn_mask_type) self.checkpoint_core_attention = config.recompute_granularity == 'selective' if self.use_flash_attn: self.core_attention_flash = FlashSelfAttention( causal=True, attention_dropout=config.attention_dropout ) # Output. self.dense = tensor_parallel.RowParallelLinear( query_projection_size, config.hidden_size, config=config, init_method=config.output_layer_init_method, bias=args.add_bias_linear, input_is_parallel=True, skip_bias_add=True) if args.use_llama2_rotary_position_embeddings: self.use_llama2_rotary_position_embeddings = True self.seq_length = args.seq_length rotary_dim = args.hidden_size // args.num_attention_heads \ if args.kv_channels is None else args.kv_channels if args.rotary_percent < 1.0: rotary_dim = int(rotary_dim * args.rotary_percent) # partial rotary embeddings, which is better than full rotary # Wang and Komatsuzaki et al # https://github.com/kingoflolz/mesh-transformer-jax/ self.rotary_emb = RotaryEmbedding( rotary_dim, args.max_position_embeddings, args.rotary_base, args.rotary_scale_factor ) else: self.use_llama2_rotary_position_embeddings = False def _checkpointed_attention_forward(self, query_layer, key_layer, value_layer, attention_mask, rotary_pos_emb=None): """Forward method with activation checkpointing.""" def custom_forward(*inputs): query_layer = inputs[0] key_layer = inputs[1] value_layer = inputs[2] attention_mask = inputs[3] output_ = self.core_attention(query_layer, key_layer, value_layer, attention_mask) return output_ q_pos_emb, k_pos_emb = (None, None) if rotary_pos_emb is None \ else rotary_pos_emb hidden_states = tensor_parallel.checkpoint( custom_forward, False, query_layer, key_layer, value_layer, attention_mask, q_pos_emb, k_pos_emb) return hidden_states def _allocate_memory(self, inference_max_sequence_len, batch_size, num_attention_heads): return torch.empty( inference_max_sequence_len, batch_size, num_attention_heads, self.hidden_size_per_attention_head, dtype=self.params_dtype, device=torch.cuda.current_device()) def forward(self, hidden_states, attention_mask, encoder_output=None, inference_params=None, rotary_pos_emb=None, position_ids=None): # hidden_states: [sq, b, h] # ================================================= # Pre-allocate memory for key-values for inference. # ================================================= is_first_step = False if inference_params: if self.layer_number not in inference_params.key_value_memory_dict: inf_max_seq_len = inference_params.max_sequence_length inf_max_batch_size = inference_params.max_batch_size inference_key_memory = self._allocate_memory( inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition) inference_value_memory = self._allocate_memory( inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition) inference_params.key_value_memory_dict[self.layer_number] = ( inference_key_memory, inference_value_memory) is_first_step = True else: inference_key_memory, inference_value_memory = \ inference_params.key_value_memory_dict[self.layer_number] # ===================== # Query, Key, and Value # ===================== if self.attention_type == AttnType.self_attn: # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] mixed_x_layer, _ = self.query_key_value(hidden_states) # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn] new_tensor_shape = mixed_x_layer.size()[:-1] + ( self.num_query_groups_per_partition, ( (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) * self.hidden_size_per_attention_head ), ) mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) # [sq, b, ng, (np/ng + 2) * hn] --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query_layer, key_layer, value_layer) = torch.split( mixed_x_layer, [ ( self.num_attention_heads_per_partition // self.num_query_groups_per_partition * self.hidden_size_per_attention_head ), self.hidden_size_per_attention_head, self.hidden_size_per_attention_head ], dim=3) # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] - query_layer = query_layer.contiguous().view(query_layer.size(0), query_layer.size(1), -1, self.hidden_size_per_attention_head) else: # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)] mixed_kv_layer, _ = self.key_value(encoder_output) # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn] new_tensor_shape = mixed_kv_layer.size()[:-1] + \ (self.num_attention_heads_per_partition, 2 * self.hidden_size_per_attention_head) mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn] (key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_kv_layer, 2) # Attention head [sq, b, h] --> [sq, b, hp] query_layer, _ = self.query(hidden_states) # [sq, b, hp] --> [sq, b, np, hn] new_tensor_shape = query_layer.size()[:-1] + \ (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) query_layer = query_layer.view(*new_tensor_shape) # ================================== # Adjust key and value for inference # ================================== # duplicate the pos_emb for self attention if rotary_pos_emb is not None: if isinstance(rotary_pos_emb, tuple): rotary_pos_emb = rotary_pos_emb else: rotary_pos_emb = ((rotary_pos_emb,) * 2) if inference_params: if self.use_llama2_rotary_position_embeddings: kv_seq_len = key_layer.shape[0] kv_seq_len += inference_params.sequence_len_offset value_layer = value_layer.transpose(0, 1).transpose(1, 2) query_layer = query_layer.transpose(0, 1).transpose(1, 2) key_layer = key_layer.transpose(0, 1).transpose(1, 2) cos, sin = self.rotary_emb(value_layer, kv_seq_len) query_layer, key_layer = apply_llama2_rotary_pos_emb( query_layer, key_layer, cos, sin, position_ids) value_layer = value_layer.transpose(1, 2).transpose(0, 1) query_layer = query_layer.transpose(1, 2).transpose(0, 1) key_layer = key_layer.transpose(1, 2).transpose(0, 1) batch_start = inference_params.batch_size_offset batch_end = batch_start + key_layer.size(1) assert batch_end <= inference_key_memory.size(1) sequence_start = inference_params.sequence_len_offset sequence_end = sequence_start + key_layer.size(0) assert sequence_end <= inference_key_memory.size(0) # Copy key and values. inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = key_layer inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = value_layer key_layer = inference_key_memory[ :sequence_end, batch_start:batch_end, ...] value_layer = inference_value_memory[ :sequence_end, batch_start:batch_end, ...] # adjust the key rotary positional embedding if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb # need to cross check this condition during inference # if not set_inference_key_value_memory: if not is_first_step: # In inference, we compute one token at a time. # Select the correct positional embedding # (only the last token in the sequence) q_pos_emb = q_pos_emb[sequence_end - 1 : sequence_end] else: # In the first forward pass of inference, # we use the entire provided prefix. # q_pos_emb here has the rope embeddings of the entire # prefix + to-be-generated output so # we slice to just the prefix. q_pos_emb = q_pos_emb[:sequence_end, :, :, :] k_pos_emb = k_pos_emb[:sequence_end, :, :, :] rotary_pos_emb = (q_pos_emb, k_pos_emb) # ================================== # core attention computation # ================================== # expand the key_layer and value_layer [sk, b, ng, hn] -> [sk, b, np, hn] key_layer = key_layer.repeat_interleave( self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim = 2 ) value_layer = value_layer.repeat_interleave( self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim = 2 ) # apply relative positional encoding (rotary embedding) if self.use_llama2_rotary_position_embeddings: kv_seq_len = key_layer.shape[0] value_layer = value_layer.transpose(0, 1).transpose(1, 2) query_layer = query_layer.transpose(0, 1).transpose(1, 2) key_layer = key_layer.transpose(0, 1).transpose(1, 2) cos, sin = self.rotary_emb(value_layer, kv_seq_len) query_layer, key_layer = apply_llama2_rotary_pos_emb( query_layer, key_layer, cos, sin, position_ids) value_layer = value_layer.transpose(1, 2).transpose(0, 1) query_layer = query_layer.transpose(1, 2).transpose(0, 1) key_layer = key_layer.transpose(1, 2).transpose(0, 1) else: # apply relative positional encoding (rotary embedding) if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb) key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb) # TODO, can apply positional embedding to value_layer so it has # absolute positional embedding. # otherwise, only relative positional embedding takes effect # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb) if not self.use_flash_attn: if self.checkpoint_core_attention: context_layer = self._checkpointed_attention_forward( query_layer, key_layer, value_layer, attention_mask) else: context_layer = self.core_attention( query_layer, key_layer, value_layer, attention_mask) else: q, k, v = [rearrange(x, 's b ... -> b s ...').contiguous() for x in (query_layer, key_layer, value_layer)] if not self.sequence_parallel: with tensor_parallel.get_cuda_rng_tracker().fork(): context_layer = self.core_attention_flash(q, k, v) else: context_layer = self.core_attention_flash(q, k, v) context_layer = rearrange(context_layer, 'b s h d -> s b (h d)').contiguous() # ================= # Output. [sq, b, h] # ================= output, bias = self.dense(context_layer) return output, bias def bias_dropout_add(x, bias, residual, prob, training): # type: (Tensor, Optional[Tensor], Tensor, float, bool) -> Tensor if bias is not None: x = x + bias out = torch.nn.functional.dropout(x, p=prob, training=training) out = residual + out return out def get_bias_dropout_add(training): def _bias_dropout_add(x, bias, residual, prob): return bias_dropout_add(x, bias, residual, prob, training) return _bias_dropout_add @torch.jit.script def bias_dropout_add_fused_train(x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, True) @torch.jit.script def bias_dropout_add_fused_inference(x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, False) class ParallelTransformerLayer(MegatronModule): """A single transformer layer. Transformer layer takes input with size [s, b, h] and returns an output of the same size. """ def __init__(self, config, layer_number, layer_type=LayerType.encoder, self_attn_mask_type=AttnMaskType.padding, drop_path_rate=0.): args = get_args() super(ParallelTransformerLayer, self).__init__() self.layer_number = layer_number self.layer_type = layer_type self.apply_residual_connection_post_norm \ = config.apply_residual_connection_post_layernorm self.bf16 = config.bf16 self.fp32_residual_connection = config.fp32_residual_connection # Normalize the input data. self.input_norm = get_norm(config) # Self attention. self.self_attention = ParallelAttention( config, layer_number, attention_type=AttnType.self_attn, attn_mask_type=self_attn_mask_type) self.hidden_dropout = config.hidden_dropout self.bias_dropout_fusion = config.bias_dropout_fusion self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None # Normalize the attention output self.post_attention_norm = get_norm(config) # Cross attention. if self.layer_type in (LayerType.decoder, LayerType.retro_decoder, LayerType.retro_decoder_with_retriever, LayerType.retro_encoder): self.inter_attention = ParallelAttention( config, layer_number, attention_type=AttnType.cross_attn) # Normalize the attention output. self.post_inter_attention_norm = get_norm(config) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -