多模态学习笔记 - 3
参考repo:WatchTower-Liu/VLM-learning; url: VLLM-BASE
吐槽
今天接着昨天的源码继续看,黑神话:悟空正好今天发售,希望广大coder能玩的开心~
学习心得
前情提要
详情请看多模态学习笔记 - 2
上次我们讲到利用view()函数对token_type_ids、position_ids进行重新塑形,确保这些张量的最后一个维度和input_shape(输入序列数据)的最后一个维度相等。重构的代码中默认启用缓存键值对(显然use_cache的bool值有点可有可无了QAQ),如果past_key_values的值为空,代表处于推理或者训练的第一步,此时我们初始化past_length为0,初始化past_key_values为长度为Qwen模型层数量的元组,self.h是Qwen模型的成员变量,我们无需太过关心(因为我们只是继承Qwen模型的成员变量,并重构了forward方法)。
如果我们当前不处于训练或推理的第一步,past_key_values显然就不为空(因为我们默认启用缓存键值对,ps:科研级代码是这样的),不管缓存量化(use_cache_quantization)启用与否,我们将past_length更新为第一个注意力头键张量的第二个或倒数第二个维度。这里唯一的区别只是元组的维数和维度不太一样。
如果position_ids为None,我们需要初始化一个position_ids,起始位置为past_length,终止位置为psst_lenght + input_shape[=1],确保我们的position_ids长度与input_shape的最后一个维度相等,随后重新塑形,同样是为了确保position_ids为二维张量,且最后一个维度与input_shape对齐,代码如下:
if token_type_ids is not None:
token_type_ids = token_type_ids.view(-1, input_shape[-1])
if position_ids is not None:
position_ids = position_ids.view(-1, input_shape[-1])
if past_key_values is None:
past_length = 0
past_key_values = tuple([None] * len(self.h))
else:
if self.use_cache_quantization:
past_length = past_key_values[0][0][0].size(2)
else:
past_length = past_key_values[0][0].size(-2)
if position_ids is None:
position_ids = torch.arange(
past_length,
input_shape[-1] + past_length,
dtype=torch.long,
device=device,
)
position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
新的记忆
代码块1
接着上面的代码,继续看MQwen.py中MQwenModel中重构的forward方法,代码如下:
if attention_mask is not None:
# image_feaute_length = self.otherConfig["image_context_length"]*self.otherConfig["image_feature_hidden_size"]
# attention_mask_length = attention_mask.shape[-1] - image_feaute_length + self.otherConfig["image_context_length"]
# attention_mask = torch.ones((batch_size, attention_mask_length), dtype=torch.long, device=device)
if batch_size <= 0:
raise ValueError("batch_size has to be defined and > 0")
attention_mask = attention_mask.view(batch_size, -1)
attention_mask = attention_mask[:, None, None, :]
attention_mask = attention_mask.to(dtype=self.dtype)
attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
如果没有传入attention_mask,我们需要根据batch_size重塑一个注意力掩码,注意力掩码用于告诉模型应该关注和忽略序列数据中的哪些部分,并且防止信息泄露,在处理序列到序列任务时利用未来信息生成当前输出。
首先检测传入的batch_size是否小于等于0,这很显然,对于空数据,是无法初始化一个合法的注意力掩码的。
如果batch_size合法,我们重塑attention_mask的第一个维度为batch_size。并且将attention_mask扩展为一个四维张量,其中第二第三维度为1,attention_mask的维度大致为(batch_size,1,1,未知),扩展为四维是为了适用于多头机制,对每一个头的输出进行操作。而后将attention_mask的数据类型变更为self.dtype这一题继承而来的成员变量。对于attention_mask的值进行翻转,将原先的1变为0,0变为1,然后让attention_mask乘以一个极大的负数。这样做的目的是让应该被忽略的地方变为一个极大的负数,而被注意的地方仍为0,考虑到softmax函数如下:
S
o
f
t
m
a
x
(
x
1
)
=
e
x
i
∑
j
e
x
j
Softmax(x_1) = \frac{e^{x_i}}{\sum_{j}e^{x_j}}
Softmax(x1)=∑jexjexi
其中
x
i
x_i
xi是输入序列中当前元素的掩码值,
x
j
x_j
xj代表任意元素的掩码值。如果掩码值为0,
e
x
i
e^{x_i}
exi的值为1,如果掩码只为极大负数,值趋近于0而不为0。
这样做的目的是为了让模型完全忽略本不应该关注的部分。如果按照原先的mask,我们将应当被忽略的地方置0,在softmax操作时,幂0的e值为1,仍然会对输出有贡献,如果将其变为一个极大的负数,那么它就能真正的趋于0,被完全忽略。
代码块2
encoder_attention_mask = None
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
我们将encoder_attention_mask置为None,在多模态场景中,Qwen作为解码器使用,不需要encoder_attention_mask,后续也没有给它赋值。head_mask则使用继承成员方法self.get_mask获取,传入两个参数,一个是head_mask,默认为None,一个是隐藏层层数,头掩码用来选择性地忽略部分头的输出,效果与attention_mask类似
代码块3
if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
hidden_states = inputs_embeds
if images is not None and first_step:
new_hidden_states = []
for b_idx, img_idx in enumerate(image_index):
new_hidden_states.append(torch.cat([hidden_states[b_idx][:img_idx], images[b_idx], hidden_states[b_idx][img_idx:]], dim = 0)) ############# concat image and text
hidden_states = torch.stack(new_hidden_states, dim = 0).to(hidden_states)
如果没有传入input_embeds,将传入的input_ids利用成员方法生成inputs_embeds,并将其作为初始的隐藏状态。
如果我们当前处于推理或者训练的第一步,并且传入了图像数据,就对图像数据和文本数据进行融合,具体来说,我们先新初始化一个列表new_hidden_states用于存储每个批次的合并数据。image_index在多模态大模型学习笔记 - 1中说明过,用来判断每个输入序列数据中图像信息的起始位置。利用torch.cat方法,将每一批次的图像信息插入到word_embeds中,最后再用torch,stack堆叠为一个新的批次,至此,图像数据和文本数据的融合完毕。
代码块4
kv_seq_len = hidden_states.size()[1]
if past_key_values[0] is not None:
# past key values[0][0] shape: bs * seq_len * head_num * dim
if self.use_cache_quantization:
kv_seq_len += past_key_values[0][0][0].shape[2]
else:
kv_seq_len += past_key_values[0][0].shape[1]
hidden_states的size大致为(batch_size, new_seq_len, 未知),new_seq_len是原始的文本数据序列长度加上图上数据序列长度,kv_seq_len获取不同模态数据合并后的序列长度
如果发现有过去缓存的键值对信息,我们就对kv_seq_len进行累加,这里的shape有点抽象,我们只用知道这些都是以缓存键值对的序列长度即可~
代码块5(ntk,选看)
if self.training or not self.use_dynamic_ntk:
ntk_alpha_list = [1.0]
elif kv_seq_len != hidden_states.size()[1]:
ntk_alpha_list = self.rotary_emb._ntk_alpha_cached_list
else:
ntk_alpha_list = []
if attention_mask is not None and kv_seq_len > self.seq_length:
true_seq_lens = attention_mask.squeeze(1).squeeze(1).eq(0).sum(dim=-1, dtype=torch.int32)
for i in range(hidden_states.size()[0]):
true_seq_len = true_seq_lens[i].item()
ntk_alpha = self.get_ntk_alpha(true_seq_len)
ntk_alpha_list.append(ntk_alpha)
else:
ntk_alpha = self.get_ntk_alpha(kv_seq_len)
ntk_alpha_list.append(ntk_alpha)
NTK比较复杂,作用也很多,这里不展开说,它的主要目的是加速收敛,线性化训练动态,提高模型解释性等(ps:我也不知道干啥用的,但感觉是用来分析模型的训练和决策过程,增强可解释性的)。
首先我们检查当前是否处于训练状态,并且不使用动态NTK,如果是,我们初始化NTK系数为1.0。
反之,我们进一步判断kv_seq_len是否和hidden_states的seq_len长度相等,假如我们先前更新了kv_seq_len的长度,即我们有以缓存的键值对,那么这里必然是不相等的,我们初始化一个ntk_alpha_list,这里调用的是继承的成员变量。
其他情况,我们初始化一个空的ntk_alpha_list,如果存在attention_mask且kv_seq_len大于继承的成员变量self.seq_len,我们用attenrion_mask计算序列的实际长度,这里去除掉四维张量attenrion_mask的中间两个维度,计算seq_len维度中指为0的元素数量(由于之前翻转了attention_mask,所以值为0代表我们需要关注的元素)。我们获取每个批次的true_seq_len,并利用成员方法获取ntk_alpha值,添加到之前初始化的ntk_alpha_list中。
如果没有提供注意力掩码或键值序列长度不大于设定的序列长度,直接为整个键值序列长度计算一个NTK缩放因子,并添加到列表中。
ps:最一头雾水的代码块。
代码块6
self.rotary_emb._ntk_alpha_cached_list = ntk_alpha_list
rotary_pos_emb_list = [
self.rotary_emb(kv_seq_len, ntk_alpha=ntk_alpha) for ntk_alpha in ntk_alpha_list
]
hidden_states = self.drop(hidden_states)
将初始化好的ntk_alpha_list缓存到_ntk_alpha_cached_list中,以便重复利用,调用self.rotary_emb方法生成旋转嵌入,传递参数皆在之前初始化完成,生成的旋转嵌入都存储于旋转嵌入列表中。
最后启用dropout随即将一些激活值置为0,提高泛化能力,防止过拟合。
代码块7
output_shape = input_shape + (hidden_states.size(-1),)
if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False
回顾一下inpui_shape的size为(batch_size,text_seq_len + image_seq_len)是一个二维张量,这里再加上hidden_stete的最后一个维度,结合为三维张量,其中hidden_state的最后一个维度就是多模态数据融合后的embed_size(参考之前代码块3中的融合过程)。
如果启用了梯度累积,并且当前处于训练状态,我们检查是否启用了缓存,由于梯度累积和缓存冲突,将use_cache置为False。梯度累积是一个内存优化技术,可以模拟大batch_size的训练,多次小批量训练后将梯度累积,并一次性用于优化器更新权重,这样能够让小批量训练类似于使用大批量训练,提高训练的稳定性和性能。
代码块8
presents = () if use_cache else None
all_self_attentions = () if output_attentions else None
all_hidden_states = () if output_hidden_states else None
for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if self.gradient_checkpointing and self.training:
def create_custom_forward(module):
def custom_forward(*inputs):
# None for past_key_value
return module(*inputs, use_cache, output_attentions)
return custom_forward
outputs = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
rotary_pos_emb_list,
None,
attention_mask,
head_mask[i],
encoder_hidden_states,
encoder_attention_mask,
)
else:
outputs = block(
hidden_states,
layer_past=layer_past,
rotary_pos_emb_list=rotary_pos_emb_list,
attention_mask=attention_mask,
head_mask=head_mask[i],
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
use_cache=use_cache,
output_attentions=output_attentions,
)
hidden_states = outputs[0]
if use_cache is True:
presents = presents + (outputs[1],)
if output_attentions:
all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
presents用于缓存键值对信息,如果启用了use_cache。
all_self_attentions用于输出每一层的注意力分数
all_hidden_states则存储每一层的隐藏层状态。
遍历模型的每一层,如果我们要输出每一层的隐藏状态,就添加当前层的隐藏状态进入元祖all_hidden_states。
如果启用了梯度累积技术,并且当前处于训练状态,我们就新建一个工厂函数,这个函数接受一个模块,并且返回一个新的函数customer_forward,这个函数可以在调用原始函数前向传播的同时,传递入新的参数。
使用pytorch的checkpoint函数执行前向传播,传递参数含义如下:
create_custom_forward(block):当前层的自定义前向传播函数,
hidden_states:当前层的隐藏状态
rotary_pos_emb_list:旋转位置嵌入列表,在之前初始化完成
各种mask:用于控制模型注意和忽略的部分
encoder_hidden_states:编码器的隐藏状态
反之,直接调用当前层的网络块进行前向传播计算,参数含义前文中都有说明,不再赘述。
从outputs中提取到当前层的隐藏状态。
判断是否启动缓存,如果启用,将当前层计算得到的键值对存储到prensents元组中。
如果output_attentions为True,将自注意力力权重存入all_self_attentions,这里根据是否启动缓存,索引有所不同。
代码块9
hidden_states = self.ln_f(hidden_states)
hidden_states = hidden_states.view(output_shape)
# Add last hidden state
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
首先对hidden_states执行层归一化操作,提高训练过程的稳定性。然后将hidden_states塑形为out_put_shape,在之前有提及,就是Input_shape + 图像和文本embed合并后的最后一个维度,如果out_put_hidden_states为True,将当前层归一化后的hidden_states添加入all_hidden_states。
代码块10
if not return_dict:
return tuple(
v for v in [hidden_states, presents, all_hidden_states] if v is not None
)
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=presents,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
)
这段代码主要处理的是返回值类型。如果不要求返回值为字典类型,则返回一个元祖,对于hidden_states等元组,依次遍历里面不为None的元素并返回。
如果返回字典,我们创建一个自定义类BaseModelOutputWithPast的实例,将各种元组传递进去,最后的返回值应该是一个字典类型的数据。
至此,MQwenModel类的forward源码看完,下面要看的就是MQwenLMHeadModel的源码。
QwenModel是基座模型,包含了Qwen的主要架构,QwenLMHeadModel 是在 QwenModel 的基础上增加了一个或多个特定的下游任务头,可以用于特定的下游任务。