Method forward

codegeex/megatron/model/transformer.py:987–1062

(
            self,
            hidden_states,
            query_hidden_state,
            attention_mask,
            layer_past=None,
            get_key_value=False,
            prompt_length=None,
            context_length=None,
    )

        self.input_tensor = input_tensor
        
    def forward(
            self,
            hidden_states,
            query_hidden_state,
            attention_mask,
            layer_past=None,
            get_key_value=False,
            prompt_length=None,
            context_length=None,
    ):

        # Checks
        if layer_past is not None:
            assert get_key_value, \
                'for not None values in layer_past, ' \
                'expected get_key_value to be set'
        if get_key_value:
            assert not self.checkpoint_activations, \
                'get_key_value does not work with ' \
                'activation checkpointing'

        # data format change to avoid explicit tranposes : [b s h] --> [s b h]
        hidden_states = hidden_states.transpose(0, 1).contiguous()
        query_hidden_state = query_hidden_state.transpose(0, 1).contiguous()

        if self.checkpoint_activations:
            hidden_states = self._checkpointed_forward(hidden_states,
                                                       attention_mask)
        else:
            if get_key_value:
                presents = []
            for index in range(self.num_layers):
                layer = self._get_layer(index)
                past = None
                if layer_past is not None:
                    past = layer_past[index]
                hidden_states = layer(hidden_states,
                                      attention_mask,
                                      layer_past=past,
                                      get_key_value=get_key_value,
                                      prompt_length=prompt_length,
                                      context_length=context_length)
                if get_key_value:
                    hidden_states, present = hidden_states
                    presents.append(present)

        if self.ln_fp16:
            hidden_states_ = self.final_layernorm(hidden_states)
        else:
            hidden_states_ = self.final_layernorm(hidden_states.float()).half()

        #################################
        # top query layer
        #################################
        past = None
        if layer_past is not None:
            past = layer_past[self.num_layers]
        hidden_states = self.topQueryLayer(hidden_states_,