Method forward

codegeex/oneflow/codegeex_model.py:530–570

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

        self.mlp = MLP(self.hidden_size)

    def forward(
        self,
        hidden_states,
        attention_mask,
        layer_past=None,
        get_key_value=False,
        prompt_length=None,
        context_length=None,
        layer_id=0,
    ):
        # hidden_states: [b, s, h]
        # Use FP32 for Layernorm
        # layernorm_output = self.input_layernorm(hidden_states.float()).half()
        layernorm_output = self.input_layernorm(hidden_states)

        # Self attention.
        attention_output, attention_mask = self.attention(layernorm_output,
                                          attention_mask,
                                          layer_past=layer_past,
                                          get_key_value=get_key_value,
                                          prompt_length=prompt_length,
                                          context_length=context_length,
                                          layer_id=layer_id)

        if get_key_value:
            attention_output, presents = attention_output

        # Residual connection.
        residual = hidden_states
        layernorm_input = attention_output + residual
        
        # Use FP32 for Layernorm
        # layernorm_output = self.post_attention_layernorm(layernorm_input.float()).half()
        layernorm_output = self.post_attention_layernorm(layernorm_input)
        mlp_output = self.mlp(layernorm_output)
        output = mlp_output + layernorm_input

        if get_key_value:
            output = [output, presents]

        return output, attention_mask


class TopQueryLayer(torch.nn.Module):