model_utils/model_proteinglm_clm.py

import math
import copy
import torch
from torch.nn import functional as F
import torch.nn as nn
import contextlib

from sat import mpu
from sat.transformer_defaults import standard_attention, attention_fn_default
from sat.mpu.utils import split_tensor_along_last_dim, divide
from sat.mpu.layers import ColumnParallelLinear
from sat.model.base_model import BaseModel, BaseMixin
from sat.model.position_embedding import RotaryEmbedding
from sat.model.position_embedding import apply_rotary_pos_emb_index
from sat.ops import LayerNorm


class RotaryEmbeddingMixin(BaseMixin):
    def __init__(
        self,
        fp16,
        hidden_size,
        num_attention_heads,
        model_parallel_size,
        rotary_embedding_2d=True,
    ):
        super().__init__()
        hidden_size_per_attention_head = divide(hidden_size, num_attention_heads)
        self.hidden_size_per_attention_head = hidden_size_per_attention_head
        self.rotary_embedding_2d = rotary_embedding_2d
        self.num_attention_heads_per_partition = divide(num_attention_heads, model_parallel_size)
        self.rotary_emb = RotaryEmbedding(
            # hidden_size_per_attention_head,
            hidden_size_per_attention_head // 2
            if rotary_embedding_2d
            else hidden_size_per_attention_head,
            base=10000,
            precision=torch.half if fp16 else torch.bfloat16,
            learnable=False,
            device=torch.cuda.current_device(),
        )


    def attention_forward(self, hidden_states, mask, **kw_args):
        attn = self.transformer.layers[kw_args["layer_id"]].attention
        attention_fn = attention_fn_default
        if "attention_fn" in attn.hooks:
            attention_fn = attn.hooks["attention_fn"]

        # [seq, b, 3 * hn * np]
        mixed_raw_layer = attn.query_key_value(hidden_states)

        # [seq, b, (np * 3 * hn)] --> [seq, b, np, 3 * hn]
        new_tensor_shape = mixed_raw_layer.size()[:-1] + (
            self.num_attention_heads_per_partition,
            3 * self.hidden_size_per_attention_head,
        )
        mixed_raw_layer = mixed_raw_layer.view(*new_tensor_shape)

        # [sq, b, np, hn]
        (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_raw_layer, 3)
        # print(key_layer.shape)
        dropout_fn = attn.attention_dropout if attn.training else None
        if self.rotary_embedding_2d:
            q1, q2 = query_layer.chunk(2, dim=(query_layer.ndim - 1))
            k1, k2 = key_layer.chunk(2, dim=(key_layer.ndim - 1))
            cos, sin = self.rotary_emb(q1, seq_len=kw_args["position_ids"].max() + 1)
            position_ids, block_position_ids = \
                kw_args["position_ids"][:, 0, :].transpose(0, 1).contiguous(), \
                kw_args["position_ids"][:, 1, :].transpose(0, 1).contiguous()
            q1, k1 = apply_rotary_pos_emb_index(q1, k1, cos, sin, position_ids)
            q2, k2 = apply_rotary_pos_emb_index(q2, k2, cos, sin, block_position_ids)
            query_layer = torch.concat([q1, q2], dim=(q1.ndim - 1))
            key_layer = torch.concat([k1, k2], dim=(k1.ndim - 1))
        else:
            kw_args["position_ids"] = kw_args["position_ids"].transpose(0, 1)
            cos, sin = self.rotary_emb(value_layer, seq_len=kw_args["position_ids"].max() + 1)
            query_layer, key_layer = apply_rotary_pos_emb_index(query_layer, key_layer, cos, sin, kw_args["position_ids"])

        context_layer = attention_fn(query_layer, key_layer, value_layer, mask, dropout_fn, **kw_args)
        output = attn.dense(context_layer)

        if attn.training:
            output = attn.output_dropout(output)

        return output


class GEGLU(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.activation_fn = F.gelu

    def forward(self, x):
        # dim=-1 breaks in jit for pt<1.10
        x1, x2 = x.chunk(2, dim=(x.ndim - 1))
        return x1 * self.activation_fn(x2)


class DeepNormWithGLUMixin(BaseMixin):
    def __init__(self, num_layers, hidden_size, inner_hidden_size=None):
        super().__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        if inner_hidden_size is None:
            inner_hidden_size = 4 * hidden_size * 2 // 3
        self.inner_hidden_size = inner_hidden_size

    def reinit(self):
        for layer in self.transformer.layers:
            del layer.mlp.dense_h_to_4h
            layer.mlp.dense_h_to_4h = ColumnParallelLinear(
                self.hidden_size,
                2 * self.inner_hidden_size,
                gather_output=False,
                bias=True,
                params_dtype=torch.half,
                module=self,
                name="dense_h_to_4h",
                skip_init=True,
            )
            del layer.mlp.activation_func
            layer.mlp.activation_func = GEGLU()

    def layer_forward(self, hidden_states, mask, *args, **kw_args):
        """
        hidden_states: [seq_len, batch, hidden_size]
        mask: [(1, 1), seq_len, seq_len]
        """
        layer = self.transformer.layers[kw_args["layer_id"]]
        # Layer norm at the begining of the transformer layer.

        attention_input = layer.input_layernorm(hidden_states)

        # Self attention.
        attention_output = layer.attention(attention_input, mask, **kw_args)

        # Residual connection.
        alpha = (2 * self.num_layers) ** 0.5
        hidden_states = attention_input * alpha + attention_output

        mlp_input = layer.post_attention_layernorm(hidden_states)

        # MLP.
        mlp_output = layer.mlp(mlp_input, **kw_args)

        # Second residual connection.
        output = mlp_input * alpha + mlp_output

        return output


class SelfAttentionWithFP32SoftmaxMixin(BaseMixin):
    def __init__(self, fp16, hidden_size, num_attention_heads, model_parallel_size):
        super().__init__()
        self.hidden_size_per_attention_head = divide(hidden_size, num_attention_heads)
        self.hidden_size_per_partition = divide(hidden_size, model_parallel_size)
        self.scale_mask_softmax = None
        self.fp16 = fp16

    @staticmethod
    def attention_mask_func(attention_scores, attention_mask):
        attention_scores.masked_fill_(attention_mask, -10000.0)
        return attention_scores

    def attention_fn(
        self,
        query_layer,
        key_layer,
        value_layer,
        attention_mask,
        attention_dropout=None,
        log_attention_weights=None,
        scaling_attention_score=True,
        mems=None,
        **kwargs
    ):

        mem = mems[kwargs["layer_id"]] if mems is not None else None

        # seqlen, batch, head, hidden_size
        seq_len, b, nh, hidden_size = key_layer.shape

        # stack, seqlen, b, head, hidden
        # b, seqlen, stack, head, hidden
        cache_kv = (
            torch.stack((key_layer, value_layer))
            .permute(2, 1, 0, 3, 4)
            .detach()
            .contiguous()
            .view(b, seq_len, nh * hidden_size * 2)
        )
        kwargs["output_this_layer"]["mem_kv"] = cache_kv

        if mem is not None:  # the first time, mem is None
            # might change batch_size
            # b, seqlen, stack, head, hidden -> stack, seqlen, b, head, hidden
            mem = mem.expand(b, -1, -1).reshape(b, mem.shape[1], 2, nh, hidden_size).permute(2, 1, 0, 3, 4)
            memk, memv = mem[0], mem[1]
            key_layer = torch.cat((memk, key_layer), dim=0)
            value_layer = torch.cat((memv, value_layer), dim=0)


        # check if use flash attention
        is_low_triangle = (attention_mask == ~torch.ones_like(attention_mask, dtype=torch.bool).tril()).all()
        is_full = (attention_mask is None) or (attention_mask == 0).all()
        if int(torch.__version__.split('.')[0]) >= 2 and (is_full or is_low_triangle):
            # Pytorch 2.0 attention uses very much memory if attention_mask is float, and has NaN bug if attention_mask is None.
            dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p
            #[b, np, sq, hn]
            query_layer, key_layer, value_layer = query_layer.permute(1,2,0,3).contiguous(), key_layer.permute(1,2,0,3).contiguous(), value_layer.permute(1,2,0,3).contiguous()
            batch_size, num_query_heads = query_layer.shape[:2] # [b, np, s, hn]
            num_kv_heads = key_layer.shape[1] # [b, np, s, hn]
            key_layer = key_layer.unsqueeze(2).expand(-1, -1, num_query_heads//num_kv_heads, -1, -1).contiguous().view(batch_size, num_query_heads, *key_layer.shape[2:])
            value_layer = value_layer.unsqueeze(2).expand(-1, -1, num_query_heads//num_kv_heads, -1, -1).contiguous().view(batch_size, num_query_heads, *value_layer.shape[2:])

            if dropout_p > 0 and mpu.get_cuda_rng_tracker is not None:
                context = mpu.get_cuda_rng_tracker().fork()
            else:
                context = contextlib.nullcontext()

            with context:
                context_layer = torch.nn.functional.scaled_dot_product_attention(
                    query_layer, key_layer, value_layer, 
                    attn_mask=None,
                    dropout_p=dropout_p,
                    is_causal=not is_full
                )


            #[sq, b, np, hn]
            context_layer = context_layer.permute(2, 0, 1, 3).contiguous()

            # [sq, b, np, hn] --> [sq, b, hp]
            new_context_layer_shape = context_layer.size()[:-2] + (-1,)
            context_layer = context_layer.view(*new_context_layer_shape)
            return context_layer

        else:
            # standard attention
            # [b, np, sq, sk]
            output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))

            query_key_layer_scaling_coeff = float(kwargs["layer_id"] + 1)


            if scaling_attention_score:
                query_layer = query_layer / (math.sqrt(self.hidden_size_per_attention_head) * query_key_layer_scaling_coeff)
            # ===================================
            # Raw attention scores. [b, np, s, s]
            # ===================================
            # [sq, b, np, hn] -> [sq, b * np, hn]
            query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
            # [sk, b, np, hn] -> [sk, b * np, hn]
            key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)

            matmul_result = torch.empty(
                output_size[0] * output_size[1],
                output_size[2],
                output_size[3],
                dtype=query_layer.dtype,
                device=torch.cuda.current_device(),
            )

            matmul_result = torch.baddbmm(
                matmul_result,
                query_layer.transpose(0, 1),  # [b * np, sq, hn]
                key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
                beta=0.0,
                alpha=1.0,
            )

            # change view to [b, np, sq, sk]
            attention_scores = matmul_result.view(*output_size)
            
            if not (attention_mask.shape[-2] == 1 and (attention_mask > 0).all()):
                # if auto-regressive, skip
                attention_scores.masked_fill_(attention_mask.bool(), -float("inf"))

            attention_scores = attention_scores.float()
            attention_scores = attention_scores * query_key_layer_scaling_coeff
        

            attention_probs = F.softmax(attention_scores, dim=-1)

            if self.fp16:
                attention_probs = attention_probs.half()
            else:
                attention_probs = attention_probs.bfloat16()

            if attention_dropout is not None:
                if mpu.get_cuda_rng_tracker() is not None:
                    with mpu.get_cuda_rng_tracker().fork():
                        attention_probs = attention_dropout(attention_probs)
                else:
                    attention_probs = attention_dropout(attention_probs)
                
            # =========================
            # Context layer. [sq, b, hp]
            # =========================
            
            # value_layer -> context layer.
            # [sk, b, np, hn] --> [b, np, sq, hn]

            # context layer shape: [b, np, sq, hn]
            output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))

            # change view [sk, b * np, hn]
            value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)

            # change view [b * np, sq, sk]
            attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
            # matmul: [b * np, sq, hn]
                
            context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))

            # change view [b, np, sq, hn]
            context_layer = context_layer.view(*output_size)

            # [b, np, sq, hn] --> [sq, b, np, hn]
            context_layer = context_layer.permute(2, 0, 1, 3).contiguous()

            # [sq, b, np, hn] --> [sq, b, hp]
            new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
            context_layer = context_layer.view(*new_context_layer_shape)
            return context_layer



class FinalForwardMixin(BaseMixin):
    def __init__(self):
        super().__init__()

    def final_forward(self, logits, **kw_args):
        return F.linear(logits, self.transformer.word_embeddings.weight).transpose(0, 1).contiguous()


class UntieFinalForwardMixin(BaseMixin):
    def __init__(self, hidden_size, vocab_size, untie_head_num, layernorm_epsilon=1.0e-5):
        super().__init__()

        self.lm_head = nn.ModuleList()
        for i in range(untie_head_num):
            self.lm_head.append(
                ColumnParallelLinear(
                hidden_size,
                2 * hidden_size,
                gather_output=True,
                bias=False,
                module=self,
                name=f"lm_head.{i}",
                )
             ) # Setting bias to False always to keep it consistent with embedding tying that also does not have a bias.

        self.head_layernorm = nn.ModuleList()
        for i in range(untie_head_num):
            self.head_layernorm.append(
                LayerNorm(
                hidden_size,
                eps=layernorm_epsilon
                )
            )
        self.activation_func=GEGLU()


    def final_forward(self, logits, **kwargs):
        logits = self.lm_head[1](logits)
        logits = self.activation_func(logits)
        logits = self.head_layernorm[1](logits)
        return F.linear(logits, self.transformer.word_embeddings.weight).transpose(0, 1).contiguous()


class NonePositionEmbedding(BaseMixin):
    def __init__(self):
        super().__init__()

    def position_embedding_forward(self, position_ids, output_cross_layer, **kw_args):
        return None


class WordEmbedding(BaseMixin):
    def __init__(self):
        super().__init__()

    def word_embedding_forward(self, input_ids, output_cross_layer, **kw_args):
        return self.transformer.word_embeddings(input_ids).transpose(0, 1)


class ProteinGLMForGeneration(BaseModel):
    def __init__(self, args, transformer=None, **kwargs):
        super().__init__(
            args,
            transformer=transformer,
            **kwargs
        )
        self.add_mixin("glu-deepnorm", DeepNormWithGLUMixin(args.num_layers, args.hidden_size, args.inner_hidden_size))
        self.add_mixin(
            "fp32-softmax",
            SelfAttentionWithFP32SoftmaxMixin(args.fp16, args.hidden_size, args.num_attention_heads, args.model_parallel_size),
        )
        if args.untie_head:
            self.add_mixin("final-forward", UntieFinalForwardMixin(args.hidden_size, args.vocab_size, args.head_num))
        else:
            self.add_mixin("final-forward", FinalForwardMixin())
        self.add_mixin("non-position-embedding", NonePositionEmbedding())
        del self.transformer.position_embeddings
        self.add_mixin("word-embedding", WordEmbedding())
        self.add_mixin(
            "rotary-embedding",
            RotaryEmbeddingMixin(
                args.fp16,
                args.hidden_size,
                args.num_attention_heads,
                args.model_parallel_size,
                args.rotary_embedding_2d
            ),
        )
        self.get_mixin("glu-deepnorm").reinit()

    @classmethod
    def add_model_specific_args(cls, parser):
        group = parser.add_argument_group('ProteinGLMForGeneration', 'ProteinGLMForGeneration Configurations')
        group.add_argument('--untie-head', action='store_true', help='untie-heads')
        group.add_argument('--head-num', default=1, type=int, help='head>1')
        group.add_argument('--infer-type', default=1, type=int, help='1 for Generation')
        group.add_argument('--rotary-embedding-2d', action='store_true',
                help='If set, use 2D rotary embedding for ProtenGLM.') 
        return super().add_model_specific_args(parser)