SwissArmyTransformer

Introdução

sat ( SwissArmyTransformer ) é uma biblioteca flexível e poderosa para desenvolver suas próprias variantes de transformador.

sat recebeu o nome de "faca do exército suíço", o que significa que todos os modelos (por exemplo, Bert, GPT, T5, GLM, Cogview, Vit ...) compartilham o mesmo código de backbone e atendem a usos versáteis com algumas mixins extras leves.

sat é alimentado pelo paralelismo deepspeed-ZeRO e Model, com o objetivo de fornecer as melhores práticas para modelos grandes de pré-treinamento e fino (parâmetros de 100m ~ 20b).

Instalar

    pip install SwissArmyTransformer

Características

• Adicione os componentes agnósticos do modelo , por exemplo, ajuste de prefixo, em apenas uma linha!

  • O ajuste do prefixo (ou tunning P) melhora o Finetuning por meio da adição de parâmetros treináveis ​​em cada camada de atenção. Para aplicá -lo a um modelo de classificação GLM (ou qualquer outro), é fácil com a nossa biblioteca.

      class ClassificationModel ( GLMModel ): # can also be BertModel, RobertaModel, etc. 
          def __init__ ( self , args , transformer = None , ** kwargs ):
              super (). __init__ ( args , transformer = transformer , ** kwargs )
              self . add_mixin ( 'classification_head' , MLPHeadMixin ( args . hidden_size , 2048 , 1 ))
              # Arm an arbitrary model with Prefix-tuning with this line!
              self . add_mixin ( 'prefix-tuning' , PrefixTuningMixin ( args . num_layers , args . hidden_size // args . num_attention_heads , args . num_attention_heads , args . prefix_len ))

  • O GPT e outros modelos auto-regressivos agem de maneira diferente durante o treinamento e a inferência. Durante a inferência, o texto é gerado por token e precisamos armazenar em cache os estados anteriores para obter eficiência. Com a nossa LIB, você só precisa considerar o comportamento durante o treinamento (forçando o professor) e transformá-lo em um modelo auto-regressivo em cache através da adição de um mixin:

      model , args = AutoModel . from_pretrained ( 'glm-10b-chinese' , args )
      model . add_mixin ( 'auto-regressive' , CachedAutoregressiveMixin ())
      # Generate a sequence with beam search
      from sat . generation . autoregressive_sampling import filling_sequence
      from sat . generation . sampling_strategies import BeamSearchStrategy
      output , * mems = filling_sequence ( model , input_seq ,
                      batch_size = args . batch_size ,
                      strategy = BeamSearchStrategy ( args . batch_size ))

• Crie seu modelo baseado em transformador com códigos mínimos . Mencionamos o GLM, que difere apenas do transformador padrão (chamado BasEmodel) na incorporação de posição (e perdas de treinamento). Precisamos apenas nos concentrar na parte relacionada ao codificar.

  Estender toda a definição:

   class BlockPositionEmbeddingMixin ( BaseMixin ):
      # Here define parameters for the mixin
      def __init__ ( self , max_sequence_length , hidden_size , init_method_std = 0.02 ):
          super ( BlockPositionEmbeddingMixin , self ). __init__ ()
          self . max_sequence_length = max_sequence_length
          self . hidden_size = hidden_size
          self . block_position_embeddings = torch . nn . Embedding ( max_sequence_length , hidden_size )
          torch . nn . init . normal_ ( self . block_position_embeddings . weight , mean = 0.0 , std = init_method_std )
      
      # Here define the method for the mixin
      def position_embedding_forward ( self , position_ids , ** kwargs ):
          position_ids , block_position_ids = position_ids [:, 0 ], position_ids [:, 1 ]
          position_embeddings = self . transformer . position_embeddings ( position_ids )
          block_position_embeddings = self . block_position_embeddings ( block_position_ids )
          return position_embeddings + block_position_embeddings
  
  class GLMModel ( BaseModel ):
      def __init__ ( self , args , transformer = None ):
          super (). __init__ ( args , transformer = transformer )
          self . add_mixin ( 'block_position_embedding' , 
              BlockPositionEmbeddingMixin ( args . max_sequence_length , args . hidden_size )
          ) # Add the mixin for GLM

• Apoios abrangentes para o treinamento . sat tem como objetivo fornecer as melhores práticas para pré -treinamento e finetuning, onde você só precisa finalizar forward_step e create_dataset_function , mas com hiperparâmetros para alterar as configurações de treinamento úteis.

  • Estenda o treinamento a várias GPUs ou nós especificando --num_nodes , --num_gpus e um simples hostfile .
  • A velocidade profunda e o paralelismo do modelo.
  • Melhor integração de zero-2 e verificação de ativação.
  • Dados automáticos de treinamento e mudança de treinamento e memmap .
  • Apoie com sucesso o treinamento de CogView2 e Cogvideo.
  • A única base de código de código aberto que suporta o Finetuning T5-10B nas GPUs atualmente.

Tour rápido

O arquivo python mais típico para usar Bert no SAT (para inferência) é o seguinte:

 # @File: inference_bert.py
from sat import get_args , get_tokenizer , AutoModel
# Parse args, initialize the environment. This is necessary.
args = get_args () 
# Automatically download and load model. Will also dump model-related hyperparameters to args.
model , args = AutoModel . from_pretrained ( 'bert-base-uncased' , args ) 
# Get the BertTokenizer according to args.tokenizer_type (automatically set).
tokenizer = get_tokenizer ( args ) 
# Here to use bert as you want!
# ...

Então podemos executar o código via

    SAT_HOME=/path/to/download python inference_bert.py --mode inference

Todos os nomes de modelos oficialmente suportados estão em urls.py.

Para Finetune ou Pré -TRAIN, um transformador também é extremamente fácil!

 # @File: finetune_bert.py
from sat import get_args , get_tokenizer , AutoModel
from sat . model . mixins import MLPHeadMixin

def create_dataset_function ( path , args ):
    # Here to load the dataset
    # ...
    assert isinstance ( dataset , torch . utils . data . Dataset )
    return dataset

def forward_step ( data_iterator , model , args , timers ):
    inputs = next ( data_iterator ) # from the dataset of create_dataset_function.
    loss , * others = model ( inputs )
    return loss
    
# Parse args, initialize the environment. This is necessary.
args = get_args () 
model , args = AutoModel . from_pretrained ( 'bert-base-uncased' , args ) 
tokenizer = get_tokenizer ( args ) 
# Here to use bert as you want!
model . del_mixin ( 'bert-final' )
model . add_mixin ( 'classification_head' , MLPHeadMixin ( args . hidden_size , 2048 , 1 ))
# ONE LINE to train! 
# args already includes hyperparams such as lr, train-iters, zero-stage ...
training_main ( args , 
    model_cls = model , 
    forward_step_function = forward_step , # user define
    create_dataset_function = create_dataset_function # user define
)

Então podemos executar o código via

deepspeed --include localhost:0,1 finetune_bert.py 
    --experiment-name ftbert 
    --mode finetune --train-iters 1000 --save /path/to/save 
    --train-data /path/to/train --valid-data /path/to/valid 
    --lr 0.00002 --batch-size 8 --zero-stage 1 --fp16

Aqui usamos o paralela de dados nas GPUs 0,1. Também podemos iniciar o treinamento em muitas máquinas interconectadas via --hostfile /path/to/hostfile . Veja o tutorial para obter mais detalhes.

Para escrever seu próprio modelo, você só precisa considerar a diferença entre o transformador padrão. Por exemplo, se você tiver uma idéia para melhorar a operação de atenção:

 from sat . model import BaseMixin
class MyAttention ( BaseMixin ):
    def __init__ ( self , hidden_size ):
        super ( MyAttention , self ). __init__ ()
        # MyAttention may needs some new params, e.g. a learnable alpha.
        self . learnable_alpha = torch . nn . Parameter ( torch . ones ( hidden_size ))
    
    # This is a hook function, the name `attention_fn` is special.
    def attention_fn ( q , k , v , mask , dropout = None , ** kwargs ):
        # Code for my attention.
        # ...
        return attention_results

Aqui, attention_fn é uma função de gancho, substituindo a ação padrão pela nova função. Todos os ganchos disponíveis estão em transformador_defaults.py. Agora, podemos usar add_mixin para aplicar nossa alteração a todos os transformadores, como Bert, Vit e Cogview. Veja o tutorial para obter mais detalhes.

Tutoriais

• Como usar modelos pré -terem coletados no SAT?
• Por que e como treinar modelos em SAT?

Citação

Atualmente, não temos um artigo, então você não precisa citar formalmente! ~

Se este projeto ajudar sua pesquisa ou engenharia, use footnote{https://github.com/THUDM/SwissArmyTransformer} para nos mencionar e recomendar SwissArmyTransformer a outros.

O tutorial para contribuir com SAT está a caminho!

O projeto é baseado em (um usuário de) DeepSpeed, Megatron-LM e Huggingface Transformers. Obrigado pelo seu trabalho incrível.