聊聊ChatGLM中P-tuning v2的应用
论文PDF地址：
https://arxiv.org/pdf/2110.07602.pdf
P-Tuning v2
摘录自第三部分
桔色块指代
可训练
的prompt embedding；蓝色块是由固定(冻结)的
预训练
语言模型 存储或计算的embedding。
Deep Prompt Tuning
continuous prompts(连续提示) 仅仅能够插入到input embedding序列层。如此，有两个问题：首先由于序列长度的约束限制，可调参数的数量有限。其次，输入的embedding对模型预测有间接的影响。
为了解决这些问题，P-Tuning v2使用deep prompt tuning的方案。正如上图的b部分，prompt作为prefix token插入到不同的层中。一方面，p-tuning v2有更多可调的特定任务参数(从 0.01% 到 0.1%~3%)，扩大了任务的容量也提高了参数效率；另一方面，添加到更深层的prompt对模型的预测会有更多直接的影响。
理解
在P-tuning v2的方案中，从图直观来看，有两个关键的点：
- prompts会加在序列的前端，而不仅仅是插入到input embedding
- 每一层都会插入prompts
v2版本主要基于p-tuning和prefix-tuning技术。prompt 向量是在模型的 embedding 层与其他输入 token 的 embedding 相拼接的，且通过在预训练模型的每一层引入可训练的 prompt 向量来提高模型对特定任务的适应性。
p-tuning主要是利用一个prompt encoder，将prompt先encoder再与input embedding进行拼接。
prefix-tuning是在Transformer的Encoder和Decoder的网络中都加了一些特定的前缀。
而基于这两种技术的v2版本，则是将两者结合。在embedding与transformer模块都做了prompt向量的插入。
ChatGLM中，首先要对prompt做encode，作为前缀prefix拼接插入到input embedding与transformer模型中。
# 转载请备注出处：https://www.cnblogs.com/zhiyong-ITNote/
class PrefixEncoder(torch.nn.Module):
"""
The torch.nn model to encode the prefix
Input shape: (batch-size, prefix-length)
Output shape: (batch-size, prefix-length, 2*layers*hidden)
"""
def __init__(self, config):
super().__init__()
self.prefix_projection = config.prefix_projection
if self.prefix_projection:
# Use a two-layer MLP to encode the prefix
self.embedding = torch.nn.Embedding(config.pre_seq_len, config.hidden_size)
self.trans = torch.nn.Sequential(
torch.nn.Linear(config.hidden_size, config.hidden_size),
torch.nn.Tanh(),
torch.nn.Linear(config.hidden_size, config.num_layers * config.hidden_size * 2)
)
else:
self.embedding = torch.nn.Embedding(config.pre_seq_len, config.num_layers * config.hidden_size * 2)
def forward(self, prefix: torch.Tensor):
if self.prefix_projection:
prefix_tokens = self.embedding(prefix)
past_key_values = self.trans(prefix_tokens)
else:
past_key_values = self.embedding(prefix)
return past_key_values
在ChatGLMModel中调用并插入到每一个transformer模型层中。
class ChatGLMModel(ChatGLMPreTrainedModel):
'''
省略其它....
'''
def __init__(self, config: ChatGLMConfig, empty_init=True):
if self.pre_seq_len is not None:
for param in self.parameters():
param.requires_grad = False
self.prefix_tokens = torch.arange(self.pre_seq_len).long()
# encode prompt
self.prefix_encoder = PrefixEncoder(config)
self.dropout = torch.nn.Dropout(0.1)
# 调用prompt
def get_prompt(self, batch_size, device, dtype=torch.half):
prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
# 调用prompt并返回
past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
past_key_values = past_key_values.view(
batch_size,
self.pre_seq_len,
self.num_layers * 2,
self.num_attention_heads,
self.hidden_size // self.num_attention_heads
)
# seq_len, b, nh, hidden_size
past_key_values = self.dropout(past_key_values)
past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
# past_key_values = [(v[0], v[1]) for v in past_key_values]
return past_key_values
# 返回transformer模型
def get_layer(layer_id):
return GLMBlock(
self.hidden_size,
self.num_attention_heads,
self.layernorm_epsilon,
layer_id,
inner_hidden_size=self.inner_hidden_size,
hidden_size_per_attention_head=self.hidden_size_per_attention_head,
layernorm=LayerNorm,
use_bias=True,
params_dtype=self.params_dtype,
position_encoding_2d=self.position_encoding_2d,
empty_init=empty_init
)
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
inputs_embeds: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPast]:
# 其它代码
if past_key_values is None:
if self.pre_seq_len is not None:
# 调用prompt
past_key_values = self.get_prompt(batch_size=input_ids.shape[0], device=input_ids.device,
dtype=inputs_embeds.dtype)
else:
past_key_values = tuple([None] * len(self.layers))
if attention_mask is None:
attention_mask = self.get_masks(
input_ids,
device=input_ids.device
)
# 其它代码
for i, layer in enumerate(self.layers):
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
# 准备参数传递到layer
layer_past = past_key_values[i]
# 每个layer 是一个GLMBlock即transformer模型层
if self.gradient_checkpointing and self.training:
# 将prompt传递到每个层中
layer_ret = torch.utils.checkpoint.checkpoint(
layer,
hidden_states,
position_ids,
attention_mask,
torch.tensor(i),
layer_past,
use_cache,
output_attentions
)
else:
layer_ret = layer(
hidden_states,
position_ids=position_ids,
attention_mask=attention_mask,
layer_id=torch.tensor(i),
layer_past=layer_past,
use_cache=use_cache,
output_attentions=output_attentions
)
# 其它代码