megatron_patch/model/qwen1_5_megablocks/transformer.py [868:1175]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if layer_type == LayerType.retro_decoder_with_retriever: self.retriever = ParallelTransformer( config=config, model_type=ModelType.retro_encoder, self_attn_mask_type=AttnMaskType.padding, pre_process=True, post_process=False, ) self._retriever_key = 'retriever' else: self.retriever = None def default_decoder_cross_attention(self, encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func): '''Cross attention for a standard encoder-decoder model.''' # Attention. attention_output, attention_bias = \ self.inter_attention(norm_output, enc_dec_attn_mask, encoder_output=encoder_output) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) # Bias-dropout-add. with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) # Normalize. norm_output = self.post_inter_attention_norm(norm_input) return norm_input, norm_output def retro_encoder_cross_attention(self, retriever_output, norm_input, norm_output, bias_dropout_add_func): """Cross attention for Retro encoder. Notation: ns : Sequence length. bs : Batch size. d : Hidden size. l : Number of chunks per sample (i.e., seq_length/chunk_length). k : Number of neighbors. r : Number of retrieved tokens (neighbors + continuation). """ ns, bs, d = norm_output.shape # [r, bs * l * k, d] # Divide sequence dimension into chunks. chunked_outputs = norm_output.reshape(self.retro_retrieved_length, -1, self.retro_num_neighbors, d) chunked_outputs_before_norm = \ norm_input.reshape(self.retro_retrieved_length, -1, self.retro_num_neighbors, d) # [r, bs*l, k, d] # Per-chunk attention. norm_inputs = [] norm_outputs = [] for k in range(self.retro_num_neighbors): # Attention. chunked_output = chunked_outputs[:,:,k].contiguous() attention_output, attention_bias = \ self.inter_attention( chunked_output, # Q (neighbor embedding) None, encoder_output=retriever_output) # K, V (hidden act) # Residual connection. if self.apply_residual_connection_post_norm: residual = chunked_output else: residual = chunked_outputs_before_norm[:,:,k] # Re-enable torch grad to enable fused optimization. with torch.enable_grad(): norm_input = bias_dropout_add_func( attention_output, None if attention_bias is None else attention_bias.expand_as(residual), residual, self.hidden_dropout) norm_inputs.append(norm_input) # Layer norm. norm_output = self.post_inter_attention_norm(norm_input) norm_outputs.append(norm_output) # Concatenate layer norms. # norm_input : [r, k * bs * l, d] # norm_output : [r, k * bs * l, d] norm_input = torch.stack(norm_inputs, dim=1).reshape(ns, bs, d) norm_output = torch.stack(norm_outputs, dim=1).reshape(ns, bs, d) return norm_input, norm_output def retro_decoder_cross_attention(self, retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func): """Cross attention for Retro decoder. Notation: ns : Sequence length. bs : Batch size. d : Hidden size. l : Number of chunks per sample (i.e., seq_length/chunk_length). m : Number of tokens per chunk. k : Number of neighbors. r : Number of retrieved tokens (neighbors + continuation). """ ns, bs, d = norm_output.shape l = int(np.ceil(ns / self.retro_chunk_length)) # Retrieve neighbors. if self.layer_type == LayerType.retro_decoder_with_retriever: first_ns = ns % self.retro_chunk_length if first_ns > 0: raise Exception("test this case.") first_chunk, rest_chunk = \ norm_output[:first_ns], norm_output[first_ns:] first_chunk = torch.nn.functional.pad( first_chunk, (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns), 'constant', 0) chunked_output = \ torch.cat((first_chunk, rest_chunk), dim=0) # [l * m, bs, d] else: chunked_output = norm_output # [l * m, bs, d] chunked_output = chunked_output \ .reshape(l, self.retro_chunk_length, bs, d) \ .permute(1, 2, 0, 3) \ .reshape(self.retro_chunk_length, bs * l, d) \ .contiguous() # Get Encoder Output retriever_output = self.retriever( hidden_states=retriever_input, attention_mask=retriever_attn_mask, retriever_output=chunked_output, retriever_attn_mask=retriever_attn_mask, inference_params=inference_params) # [r, k * bs * l , d] retriever_output = retriever_output.reshape( self.retro_retrieved_length * self.retro_num_neighbors, bs * l, d) # [r * k, bs * l, d] # Chunks. pad = (ns - 1) % self.retro_chunk_length attending_chunks = norm_output[pad:] padded_chunks = torch.nn.functional.pad( attending_chunks, (0, 0, 0, 0, 0, self.retro_chunk_length - 1), 'constant', 0) padded_chunked_output = padded_chunks \ .reshape(l, self.retro_chunk_length, bs, d) \ .permute(1, 2, 0, 3) padded_chunked_output = padded_chunked_output.reshape( self.retro_chunk_length, bs * l, d).contiguous() # Encoder output. attention_output, attention_bias = \ self.inter_attention(padded_chunked_output, None, encoder_output=retriever_output) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input # Re-enable torch grad to enable fused optimization. with torch.enable_grad(): norm_input = bias_dropout_add_func( attention_output, None if attention_bias is None else attention_bias.expand_as(attention_output), torch.zeros_like(attention_output), self.hidden_dropout) norm_input = norm_input \ .reshape(self.retro_chunk_length, bs, l, d) \ .permute(2, 0, 1, 3) # [l, m, bs, d] norm_input = norm_input.reshape(self.retro_chunk_length * l, bs, d) norm_input = torch.nn.functional.pad( norm_input, (0, 0, 0, 0, pad, 0), 'constant', 0)[:ns] # [ns, b, d] norm_input = norm_input + residual # Layer norm post the decoder attention norm_output = self.post_inter_attention_norm(norm_input) return retriever_output, norm_input, norm_output def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, retriever_input=None, retriever_output=None, retriever_attn_mask=None, inference_params=None, rotary_pos_emb=None, position_ids=None): # hidden_states: [s, b, h] # Layer norm at the beginning of the transformer layer. norm_output = self.input_norm(hidden_states) # Self attention. attention_output, attention_bias = \ self.self_attention( norm_output, attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, position_ids=position_ids ) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = hidden_states if self.drop_path is None: # jit scripting for a nn.module (with dropout) is not # trigerring the fusion kernel. For now, we use two # different nn.functional routines to account for varying # dropout semantics during training and inference phases. if self.bias_dropout_fusion: if self.training: bias_dropout_add_func = bias_dropout_add_fused_train else: bias_dropout_add_func = bias_dropout_add_fused_inference else: bias_dropout_add_func = get_bias_dropout_add(self.training) if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) else: out = torch.nn.functional.dropout(attention_output + attention_bias, p=self.hidden_dropout, training=self.training) norm_input = residual + self.drop_path(out) # Layer norm post the self attention. norm_output = self.post_attention_norm(norm_input) # Cross attention. if self.layer_type == LayerType.encoder: pass elif self.layer_type == LayerType.decoder: norm_input, norm_output = \ self.default_decoder_cross_attention( encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type == LayerType.retro_encoder: norm_input, norm_output = \ self.retro_encoder_cross_attention( retriever_output, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type in (LayerType.retro_decoder, LayerType.retro_decoder_with_retriever): retriever_output, norm_input, norm_output = \ self.retro_decoder_cross_attention( retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func) else: raise Exception("Unsupported layer type, '%s'." % self.layer_type.name) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - megatron_patch/model/qwen_vl/transformer.py [983:1290]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if layer_type == LayerType.retro_decoder_with_retriever: self.retriever = ParallelTransformer( config=config, model_type=ModelType.retro_encoder, self_attn_mask_type=AttnMaskType.padding, pre_process=True, post_process=False, ) self._retriever_key = 'retriever' else: self.retriever = None def default_decoder_cross_attention(self, encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func): '''Cross attention for a standard encoder-decoder model.''' # Attention. attention_output, attention_bias = \ self.inter_attention(norm_output, enc_dec_attn_mask, encoder_output=encoder_output) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) # Bias-dropout-add. with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) # Normalize. norm_output = self.post_inter_attention_norm(norm_input) return norm_input, norm_output def retro_encoder_cross_attention(self, retriever_output, norm_input, norm_output, bias_dropout_add_func): """Cross attention for Retro encoder. Notation: ns : Sequence length. bs : Batch size. d : Hidden size. l : Number of chunks per sample (i.e., seq_length/chunk_length). k : Number of neighbors. r : Number of retrieved tokens (neighbors + continuation). """ ns, bs, d = norm_output.shape # [r, bs * l * k, d] # Divide sequence dimension into chunks. chunked_outputs = norm_output.reshape(self.retro_retrieved_length, -1, self.retro_num_neighbors, d) chunked_outputs_before_norm = \ norm_input.reshape(self.retro_retrieved_length, -1, self.retro_num_neighbors, d) # [r, bs*l, k, d] # Per-chunk attention. norm_inputs = [] norm_outputs = [] for k in range(self.retro_num_neighbors): # Attention. chunked_output = chunked_outputs[:,:,k].contiguous() attention_output, attention_bias = \ self.inter_attention( chunked_output, # Q (neighbor embedding) None, encoder_output=retriever_output) # K, V (hidden act) # Residual connection. if self.apply_residual_connection_post_norm: residual = chunked_output else: residual = chunked_outputs_before_norm[:,:,k] # Re-enable torch grad to enable fused optimization. with torch.enable_grad(): norm_input = bias_dropout_add_func( attention_output, None if attention_bias is None else attention_bias.expand_as(residual), residual, self.hidden_dropout) norm_inputs.append(norm_input) # Layer norm. norm_output = self.post_inter_attention_norm(norm_input) norm_outputs.append(norm_output) # Concatenate layer norms. # norm_input : [r, k * bs * l, d] # norm_output : [r, k * bs * l, d] norm_input = torch.stack(norm_inputs, dim=1).reshape(ns, bs, d) norm_output = torch.stack(norm_outputs, dim=1).reshape(ns, bs, d) return norm_input, norm_output def retro_decoder_cross_attention(self, retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func): """Cross attention for Retro decoder. Notation: ns : Sequence length. bs : Batch size. d : Hidden size. l : Number of chunks per sample (i.e., seq_length/chunk_length). m : Number of tokens per chunk. k : Number of neighbors. r : Number of retrieved tokens (neighbors + continuation). """ ns, bs, d = norm_output.shape l = int(np.ceil(ns / self.retro_chunk_length)) # Retrieve neighbors. if self.layer_type == LayerType.retro_decoder_with_retriever: first_ns = ns % self.retro_chunk_length if first_ns > 0: raise Exception("test this case.") first_chunk, rest_chunk = \ norm_output[:first_ns], norm_output[first_ns:] first_chunk = torch.nn.functional.pad( first_chunk, (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns), 'constant', 0) chunked_output = \ torch.cat((first_chunk, rest_chunk), dim=0) # [l * m, bs, d] else: chunked_output = norm_output # [l * m, bs, d] chunked_output = chunked_output \ .reshape(l, self.retro_chunk_length, bs, d) \ .permute(1, 2, 0, 3) \ .reshape(self.retro_chunk_length, bs * l, d) \ .contiguous() # Get Encoder Output retriever_output = self.retriever( hidden_states=retriever_input, attention_mask=retriever_attn_mask, retriever_output=chunked_output, retriever_attn_mask=retriever_attn_mask, inference_params=inference_params) # [r, k * bs * l , d] retriever_output = retriever_output.reshape( self.retro_retrieved_length * self.retro_num_neighbors, bs * l, d) # [r * k, bs * l, d] # Chunks. pad = (ns - 1) % self.retro_chunk_length attending_chunks = norm_output[pad:] padded_chunks = torch.nn.functional.pad( attending_chunks, (0, 0, 0, 0, 0, self.retro_chunk_length - 1), 'constant', 0) padded_chunked_output = padded_chunks \ .reshape(l, self.retro_chunk_length, bs, d) \ .permute(1, 2, 0, 3) padded_chunked_output = padded_chunked_output.reshape( self.retro_chunk_length, bs * l, d).contiguous() # Encoder output. attention_output, attention_bias = \ self.inter_attention(padded_chunked_output, None, encoder_output=retriever_output) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = norm_input # Re-enable torch grad to enable fused optimization. with torch.enable_grad(): norm_input = bias_dropout_add_func( attention_output, None if attention_bias is None else attention_bias.expand_as(attention_output), torch.zeros_like(attention_output), self.hidden_dropout) norm_input = norm_input \ .reshape(self.retro_chunk_length, bs, l, d) \ .permute(2, 0, 1, 3) # [l, m, bs, d] norm_input = norm_input.reshape(self.retro_chunk_length * l, bs, d) norm_input = torch.nn.functional.pad( norm_input, (0, 0, 0, 0, pad, 0), 'constant', 0)[:ns] # [ns, b, d] norm_input = norm_input + residual # Layer norm post the decoder attention norm_output = self.post_inter_attention_norm(norm_input) return retriever_output, norm_input, norm_output def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, retriever_input=None, retriever_output=None, retriever_attn_mask=None, inference_params=None, rotary_pos_emb=None, position_ids=None): # hidden_states: [s, b, h] # Layer norm at the beginning of the transformer layer. norm_output = self.input_norm(hidden_states) # Self attention. attention_output, attention_bias = \ self.self_attention( norm_output, attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, position_ids=position_ids ) # Residual connection. if self.apply_residual_connection_post_norm: residual = norm_output else: residual = hidden_states if self.drop_path is None: # jit scripting for a nn.module (with dropout) is not # trigerring the fusion kernel. For now, we use two # different nn.functional routines to account for varying # dropout semantics during training and inference phases. if self.bias_dropout_fusion: if self.training: bias_dropout_add_func = bias_dropout_add_fused_train else: bias_dropout_add_func = bias_dropout_add_fused_inference else: bias_dropout_add_func = get_bias_dropout_add(self.training) if attention_bias is not None: attention_bias = attention_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( attention_output, attention_bias, residual, self.hidden_dropout) else: out = torch.nn.functional.dropout(attention_output + attention_bias, p=self.hidden_dropout, training=self.training) norm_input = residual + self.drop_path(out) # Layer norm post the self attention. norm_output = self.post_attention_norm(norm_input) # Cross attention. if self.layer_type == LayerType.encoder: pass elif self.layer_type == LayerType.decoder: norm_input, norm_output = \ self.default_decoder_cross_attention( encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type == LayerType.retro_encoder: norm_input, norm_output = \ self.retro_encoder_cross_attention( retriever_output, norm_input, norm_output, bias_dropout_add_func) elif self.layer_type in (LayerType.retro_decoder, LayerType.retro_decoder_with_retriever): retriever_output, norm_input, norm_output = \ self.retro_decoder_cross_attention( retriever_input, retriever_output, retriever_attn_mask, norm_input, norm_output, inference_params, bias_dropout_add_func) else: raise Exception("Unsupported layer type, '%s'." % self.layer_type.name) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -