ChatGLM之modeling_glm.py代码理解

本文是对ChatGLM的model部分进行讲解，主要讲解时modeling_glm的代码，更接近ChaGLM计算的核心去了解，大模型chatGLM的运行逻辑！本文只是对代码内容进行个人理解的解释，如有问题欢迎批评改造。后续会逐步展开对chatGLM项目的讲解。让更多人了解chatGLM，从而更好理解大模型。

"""GPT-2 model."""
#导入了所需的Python库和自定义模块。
import torch
import torch.nn as nn
import torch.nn.functional as F

import mpu
from model.prompt import PromptSpell
from utils import print_rank_0


def init_method_normal(std=0.02):
    """基于正态分布的初始化方法。

    这仅用于嵌入层（embeddings）。Transformer模型有自己的初始化方法。
    """
    

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std) #初始化嵌入层（embeddings）

    return init_

#创建一个GLM（Generalized Language Model）语言模型
class GLMModel(torch.nn.Module):
    """一个继承自torch.nn.Module的类，代表了一个GLM语言模型。
    forward方法的输出是logits（即预测结果），这些logits可能是并行或串行的，
    这取决于parallel_output标志的设置。
    """

    def __init__(self,
                 num_layers,
                 vocab_size,
                 hidden_size,
                 num_attention_heads,
                 embedding_dropout_prob,
                 attention_dropout_prob,
                 output_dropout_prob,
                 max_sequence_length,
                 max_memory_length,
                 checkpoint_activations,
                 checkpoint_num_layers=1,
                 parallel_output=True,
                 relative_encoding=False,
                 block_position_encoding=False,
                 output_predict=True,
                 spell_length=None,
                 spell_func='lstm',
                 attention_scale=1.0,
                 ):

        super(GLMModel, self).__init__()

        self.parallel_output = parallel_output
        self.output_predict = output_predict
        self.hidden_size = hidden_size

        init_method = init_method_normal(std=0.02)
        
        '''
		创建了一个嵌入层（embeddings）对象，用于将输入的单词ID转换为对应的词向量。
	    mpu.VocabParallelEmbedding是一个特殊的并行嵌入层，它可以处理多个GPU并行工作。
	    它使用之前创建的init_method来初始化权重。
		'''
        # Word embeddings (parallel).
        self.word_embeddings = mpu.VocabParallelEmbedding(
            vocab_size, hidden_size, init_method=init_method)
        '''
		创建了一个Transformer模型对象，使用GPT-2的并行实现。
		这个Transformer模型是用来处理文本数据的核心部分。
		'''
        # Transformer
        self.transformer = mpu.GPT2ParallelTransformer(num_layers,
                                                       hidden_size,
                                                       num_attention_heads,
                                                       max_sequence_length,
                                                       max_memory_length,
                                                       embedding_dropout_prob,
                                                       attention_dropout_prob,
                                                       output_dropout_prob,
                                                       checkpoint_activations,
                                                       checkpoint_num_layers,
                                                       attention_scale=attention_scale,
                                                       relative_encoding=relative_encoding,
                                                       block_position_encoding=block_position_encoding)
        
        #                                               
        if spell_length is not None:
            self.prompt_spell = PromptSpell(spell_length, self.hidden_size, spell_func)

    #用于冻结Transformer模型的一部分或全部权重，以便在训练过程中固定它们
    #将词嵌入层（word_embeddings）的梯度计算设置为False，
    #这样在反向传播过程中，词嵌入层的权重不会更新，即冻结了这一层。
    def freeze_transformer(self, tune_prefix_layers=None):
        log_str = "Freeze transformer"
        self.word_embeddings.requires_grad_(False)
        self.transformer.requires_grad_(False)
        #这里检查是否传入了tune_prefix_layers参数。
        #如果tune_prefix_layers不为None，则会对Transformer的部分权重进行微调（解冻）。
        if tune_prefix_layers is not None:
            log_str += f" tune {tune_prefix_layers} prefix layers"
            
            #遍历从0到tune_prefix_layers-1的数字，根据tune_prefix_layers的值决定解冻哪些Transformer层。
            for i in range(tune_prefix_layers):
                #对于第i层Transformer，将其梯度计算设置为True，这样在反向传播时，第i层的权重会参与更新，
                #从而实现对前几层的微调
                self.transformer.layers[i].requires_grad_(True)
                
        #打印记录的日志信息，提示哪些部分的权重被冻结，哪些被解冻（微调）
        print_rank_0(log_str)
        
    #
    def forward(self, input_ids, position_ids, attention_mask, *mems, return_memory=False, detach_memory=True,
        prompt_pos=None):
        
        # Embeddings.
        batch_size = input_ids.size(0)
        #将输入的单词ID转换为对应的词向量（词嵌入），通过调用之前创建的self.word_embeddings对象。
        words_embeddings = self.word_embeddings(input_ids)
        embeddings = words_embeddings
        
		#检查是否传入了prompt_pos参数，用于指定插入Prompt的位置。
        if prompt_pos is not None:
            embeddings = embeddings.clone()
            
            #根据Prompt信息，生成与之相对应的嵌入向量（可能是通过调用self.prompt_spell对象实现的）。
            prompt_embeds = self.prompt_spell()
            #将生成的Prompt嵌入向量插入到对应位置的词嵌入中
            batch_index = torch.arange(batch_size, device=input_ids.device).unsqueeze(1)
            embeddings[batch_index, prompt_pos] = prompt_embeds
        
        #调用Transformer模型，传递嵌入向量和其他参数，得到Transformer的输出。    
        # Transformer.
        transformer_output = self.transformer(embeddings, position_ids, attention_mask, mems,
                                              return_memory=return_memory, detach_memory=detach_memory)
        #从Transformer输出中提取logits（预测结果）和隐藏层输出                                      
        logits, hidden_layers = transformer_output
        
        #将隐藏层输出保存到outputs变量中，以便在返回时使用。
        outputs = hidden_layers

        if self.output_predict:
            # Parallel logits.
            logits_parallel = mpu.copy_to_model_parallel_region(
                logits)
            logits_parallel = F.linear(logits_parallel, self.word_embeddings.weight)
             
            #根据是否采用多个GPU并行输出的结果
            if self.parallel_output:
                return (logits_parallel, *outputs)

            return (mpu.gather_from_model_parallel_region(logits_parallel), *outputs)
        else:
            return (logits, *outputs)

#构建一个Seq2Seq Transformer模型。
#它包含了一个编码器（Encoder）和一个解码器（Decoder）部分，用于将源文本（source）转换为目标文本（target）。
class EncoderDecoder(torch.nn.Module):
    """Seq2Seq Transformer Model
    The output of the forward method are the logits (parallel or serial depending on the `parallel_output` flag).
    """

    def __init__(self,
                 num_layers,
                 vocab_size,
                 hidden_size,
                 num_attention_heads,
                 embedding_dropout_prob,
                 attention_dropout_prob,
                 output_dropout_prob,
                 max_sequence_length,
                 max_memory_length,
                 checkpoint_activations,
                 checkpoint_num_layers=1,
                 parallel_output=True,
                 output_predict=True
                 ):
        super(EncoderDecoder, self).__init__()

        self.parallel_output = parallel_output
        self.output_predict = output_predict

        init_method = init_method_normal(std=0.02)

        # Word embeddings (parallel).词嵌入层（Word embeddings）
        #mpu.VocabParallelEmbedding是一个用于处理词汇表并行嵌入的特殊层。
        #它接受词汇表大小（vocab_size）和隐藏层大小（hidden_size）作为参数，并使用init_method来初始化词嵌入的权重。词嵌入层用于将输入的单词ID转换为对应的词向量，用于后续的模型处理。
        self.word_embeddings = mpu.VocabParallelEmbedding(
            vocab_size, hidden_size, init_method=init_method)

        # Transformer
        '''
        编码器（self.encoder）和解码器（self.decoder）。它们都是使用GPT-2的并行实现的Transformer模型。
        这两个Transformer模型共享相同的参数设置，包括层数（num_layers）、隐藏层大小（hidden_size）、
        注意力头数（num_attention_heads）、最大序列长度（max_sequence_length）、最大记忆长度（max_memory_length）以及各种丢弃率参数等。
        解码器还使用了一个名为use_decoder_layer的参数，用于区分编码器和解码器的处理方式。'''
        self.encoder = mpu.GPT2ParallelTransformer(num_layers,
                                                   hidden_size,
                                                   num_attention_heads,
                                                   max_sequence_length,
                                                   max_memory_length,
                                                   embedding_dropout_prob,
                                                   attention_dropout_prob,
                                                   output_dropout_prob,
                                                   checkpoint_activations,
                                                   checkpoint_num_layers)
        self.decoder = mpu.GPT2ParallelTransformer(num_layers,
                                                   hidden_size,
                                                   num_attention_heads,
                                                   max_sequence_length,
                                                   max_memory_length,
                                                   embedding_dropout_prob,
                                                   attention_dropout_prob,
                                                   output_dropout_prob,
                                                   checkpoint_activations,
                                                   checkpoint_num_layers,
                                                   use_decoder_layer=True)
    '''
    前向传播过程，将输入数据（source和target）通过编码器（encoder）和解码器（decoder）进行处理，并生成模型的输出。
   '''
    def forward(self, source_ids, target_ids, source_position_ids, target_position_ids, source_mask, target_mask):
        # Embeddings. 通过词嵌入层（self.word_embeddings）将输入的源文本（source）和目标文本（target）的单词ID转换为对应的词向量。
        #源文本和目标文本都分别对应着嵌入后的向量source_embeddings和target_embeddings。
        source_embeddings = self.word_embeddings(source_ids)
        target_embeddings = self.word_embeddings(target_ids)

        # Transformer.
        '''
        将源文本的嵌入向量source_embeddings输入到编码器（self.encoder）中进行处理。
        调用self.encoder并传入嵌入向量、位置ID（source_position_ids）和掩码（source_mask），得到编码器的输出encoder_output。
        '''
        encoder_output, _ = self.encoder(source_embeddings, source_position_ids, source_mask)
        #调用self.decoder并传入嵌入向量、位置ID（target_position_ids）和掩码（target_mask），得到解码器的输出decoder_output。
        decoder_output, _ = self.decoder(target_embeddings, target_position_ids, target_mask)
        
        #根据模型的配置，这里判断是否进行输出预测（self.output_predict为True或False）
        if self.output_predict:
            # Parallel logits.
            '''将解码器的输出decoder_output复制到模型的并行区域（mpu.copy_to_model_parallel_region）。
            然后，通过线性变换F.linear将复制的输出与词嵌入的权重（self.word_embeddings.weight）相乘，
            得到并行的logits（预测结果），即logits_parallel。'''
            output_parallel = mpu.copy_to_model_parallel_region(decoder_output)
            logits_parallel = F.linear(output_parallel, self.word_embeddings.weight)
            
            #self.parallel_output为True，表示模型要返回并行的logits，则直接返回(logits_parallel,)，其中逗号表示返回一个元组。
            if self.parallel_output:
                return (logits_parallel,)
            '''self.parallel_output为False，表示模型要返回非并行的logits，因此需要将并行的logits从多个GPU上收集和合并成单个输出。
            通过mpu.gather_from_model_parallel_region函数实现，最终返回一个包含非并行logits的元组(logits_parallel,)
            '''
            return (mpu.gather_from_model_parallel_region(logits_parallel),)
        else:
            #self.output_predict为False，表示模型不进行输出预测，而是直接返回解码器的输出（decoder_output），因此返回(decoder_output,)。
            return (decoder_output,)

def glm_get_params_for_weight_decay_optimization(module):
    weight_decay_params = {'params': []}
    no_weight_decay_params = {'params': [], 'weight_decay': 0.0}
    for module_ in module.modules():
        if isinstance(module_, (mpu.LayerNorm, torch.nn.LayerNorm)):
            no_weight_decay_params['params'].extend(
                [p for p in list(module_._parameters.values())
                 if p is not None and p.requires_grad])
        else:
            weight_decay_params['params'].extend(
                [p for n, p in list(module_._parameters.items())
                 if p is not None and p.requires_grad and n != 'bias'])
            no_weight_decay_params['params'].extend(
                [p for n, p in list(module_._parameters.items())
                 if p is not None and p.requires_grad and n == 'bias'])

    return weight_decay_params, no_weight_decay_params