本文是对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