ChatGLM之modeling_glm.py代码理解_modeling_chatglm.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_

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

#创建一个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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

#构建一个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,)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19