modules/SwissArmyTransformer/sat/model/official/cait_model.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from sat.model.base_model import BaseMixin, BaseModel, non_conflict
from sat.model.official.vit_model import ViTModel, ClsMixin
from sat.model.mixins import BaseMixin
from sat import mpu

class AttnMixin(BaseMixin):
    def __init__(self, num_heads, num_layers):
        super().__init__()
        self.num_layers = num_layers
        self.proj_l = nn.ModuleList([nn.Linear(num_heads, num_heads) for i in range(num_layers)])
        self.proj_w = nn.ModuleList([nn.Linear(num_heads, num_heads) for i in range(num_layers)])

    def attention_fn(self, query_layer, key_layer, value_layer, attention_mask,
                       attention_dropout=None, log_attention_weights=None, scaling_attention_score=True, **kwargs):
        # adapted from https://github.com/THUDM/sat/blob/main/sat/mpu/transformer.py#L47
        if scaling_attention_score:
            query_layer = query_layer * (query_layer.shape[-1]**-0.5) # / math.sqrt(query_layer.shape[-1])
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        if log_attention_weights is not None:
            attention_scores += log_attention_weights
        
        attention_scores = self.proj_l[kwargs['layer_id']](attention_scores.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

        if not (attention_mask.shape[-2] == 1 and (attention_mask > 0).all()):
            # if auto-regressive, skip
            attention_scores = torch.mul(attention_scores, attention_mask) - \
                            10000.0 * (1.0 - attention_mask)

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

        attention_probs = self.proj_w[kwargs['layer_id']](attention_probs.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

        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 = torch.matmul(attention_probs, value_layer)
        return context_layer

    def reinit(self, parent_model=None):
        # init with identity matrix so that pretrained weights with standard_attn can be reused
        for i in range(self.num_layers):
            nn.init.eye_(self.proj_l[i].weight)
            nn.init.eye_(self.proj_w[i].weight)

class EncForward(BaseMixin):
    def __init__(self, dim, num_layers, init_values=1e-4):
        super().__init__()
        self.gamma_1 = nn.ParameterList([nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) for i in range(num_layers)])
        self.gamma_2 = nn.ParameterList([nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) for i in range(num_layers)])
    def layer_forward(self, hidden_states, mask, *args, **kw_args):
        layer = self.transformer.layers[kw_args['layer_id']]

        # Layer norm at the begining of the transformer layer.
        layernorm_output1 = layer.input_layernorm(hidden_states)
        # Self attention.
        attention_output = layer.attention(layernorm_output1, mask, **kw_args)

        # Residual connection.
        layernorm_input = hidden_states + self.gamma_1[kw_args['layer_id']] * attention_output
        # Layer norm post the self attention.
        layernorm_output = layer.post_attention_layernorm(layernorm_input)

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

        # Second residual connection.
        output = layernorm_input + self.gamma_2[kw_args['layer_id']] * mlp_output

        return output

from sat.model.transformer import standard_attention
from sat.mpu.utils import split_tensor_along_last_dim

class DecForward(BaseMixin):
    def __init__(self, dim, num_layers, init_values=1e-4):
        super().__init__()
        self.gamma_1 = nn.ParameterList([nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) for i in range(num_layers)])
        self.gamma_2 = nn.ParameterList([nn.Parameter(init_values * torch.ones((dim)), requires_grad=True) for i in range(num_layers)])
    
    def position_embedding_forward(self, position_ids, **kwargs):
        return 0

    def layer_forward(self, hidden_states, mask, *args, **kw_args):
        '''
            hidden_states: [batch, seq_len, hidden_size]
            mask: [(1, 1), seq_len, seq_len]
        '''
        layer = self.transformer.layers[kw_args['layer_id']]
        encoder_outputs = kw_args['encoder_outputs']
        assert encoder_outputs is not None
        # Layer norm at the begining of the transformer layer.
        u = torch.cat([hidden_states, encoder_outputs], 1)
        layernorm_output1 = layer.input_layernorm(u)
        assert 'cross_attention_mask' in kw_args
        # Cross attention
        attention_output = layer.cross_attention(layernorm_output1, **kw_args)
        # Residual connection.
        layernorm_input = hidden_states + self.gamma_1[kw_args['layer_id']] * attention_output
        # Layer norm post the cross attention
        layernorm_output = layer.post_cross_attention_layernorm(layernorm_input)

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

        # Second residual connection.
        output = layernorm_input + self.gamma_2[kw_args['layer_id']] * mlp_output

        return output

    def cross_attention_forward(self, hidden_states, cross_attention_mask, encoder_outputs, **kw_args):
        # adapted from https://github.com/THUDM/sat/blob/d8c9d1e0a9bb2af1e1d26a68b35f16d84aafcc2f/sat/mpu/transformer.py#L216
        # if you want to use a customized attention_fn, just inherit the attention mixin for this mixin and use self.attention_fn instead of self.hooks['attention_fn']
        layer = self.transformer.layers[kw_args['layer_id']].cross_attention
        
        attention_fn = standard_attention
        
        mixed_query_layer = layer.query(hidden_states[:, :1])
        mixed_x_layer = layer.key_value(hidden_states)
        (mixed_key_layer, mixed_value_layer) = split_tensor_along_last_dim(mixed_x_layer, 2)

        dropout_fn = layer.attention_dropout if layer.training else None
        # Reshape and transpose [b, np, s, hn]
        query_layer = layer._transpose_for_scores(mixed_query_layer)
        key_layer = layer._transpose_for_scores(mixed_key_layer)
        value_layer = layer._transpose_for_scores(mixed_value_layer)

        context_layer = attention_fn(query_layer, key_layer, value_layer, cross_attention_mask, dropout_fn,
                                        cross_attention=True, **kw_args)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (layer.hidden_size_per_partition,)
        # [b, s, hp]
        context_layer = context_layer.view(*new_context_layer_shape)

        # Output. [b, s, h]
        output = layer.dense(context_layer)
        if layer.training:
            output = layer.output_dropout(output)

        return output

class CaiTEncoder(ViTModel):
    def __init__(self, args, transformer=None, layernorm_epsilon=1e-6, use_final_layernorm=False):
        super().__init__(args, transformer=transformer, layernorm_epsilon=layernorm_epsilon, use_final_layernorm=use_final_layernorm)
        self.del_mixin('cls')
        self.add_mixin('attn', AttnMixin(args.num_attention_heads, args.num_layers))
        self.add_mixin('enc_forward', EncForward(args.hidden_size, args.num_layers, init_values=args.init_scale))
    
    @classmethod
    def add_model_specific_args(cls, parser):
        group = parser.add_argument_group('CaiT-enc', 'CaiT encoder Configurations')
        group.add_argument('--init-scale', type=float, default=1e-4)
        return super().add_model_specific_args(parser)

class CaiTDecoder(BaseModel):
    def __init__(self, args, transformer=None, layernorm_epsilon=1e-6):
        super().__init__(args, is_decoder=True, transformer=transformer, layernorm_epsilon=layernorm_epsilon)
        self.add_mixin('cls', ClsMixin(args.hidden_size, args.num_classes))
        self.add_mixin('dec_forward', DecForward(args.hidden_size, args.num_layers, init_values=args.init_scale))
    @classmethod
    def add_model_specific_args(cls, parser):
        return super().add_model_specific_args(parser)

from sat.model import EncoderDecoderModel
import argparse

class CaiT(EncoderDecoderModel):
    def __init__(self, args, transformer=None, layernorm_epsilon=1e-6):
        encoder = CaiTEncoder(args, transformer=transformer, layernorm_epsilon=layernorm_epsilon)
        dec_args = argparse.Namespace(**vars(args))
        # dec_args.enc_hidden_size = dec_args.hidden_size  # used for cross attn
        override_attrs = ['num_layers', 'hidden_size', 'num_attention_heads', 'layernorm_order',
                            'max_sequence_length', 'inner_hidden_size', 'hidden_size_per_attention_head']
        for name in override_attrs:
            dec_attr = getattr(dec_args, 'dec_' + name, None)
            if dec_attr is not None:  # else use encoder-config
                setattr(dec_args, name, dec_attr)
        decoder = CaiTDecoder(dec_args, transformer=transformer, layernorm_epsilon=layernorm_epsilon)
        super().__init__(args, encoder=encoder, decoder=decoder)
        
    def forward(self, input_ids, enc_position_ids, dec_position_ids, *, enc_attention_mask=None, dec_attention_mask=None, cross_attention_mask=None, **kw_args):
        # Please use self.decoder for auto-regressive generation.
        if enc_attention_mask is None:
            enc_attention_mask = torch.ones(1, 1, dtype=self.encoder.transformer.word_embeddings.weight.dtype, device=input_ids.device)
        if cross_attention_mask is None:
            cross_attention_mask = enc_attention_mask
        encoder_outputs = self.encode(input_ids, enc_position_ids, enc_attention_mask, **kw_args)
        decoder_outputs, *mems = self.decode(input_ids, dec_position_ids, dec_attention_mask, encoder_outputs=encoder_outputs, cross_attention_mask=cross_attention_mask, **kw_args)
        return (encoder_outputs, decoder_outputs, *mems)