理解 huggingface transformers 中 GPT2 模型

作为解码器，GPT2 在预测 token 时只同比自身位置更前的 token 计算相似度，且每一层 GPT2Block 都是遵循同样的原则。

一般来所，用 tokenizer 生成的 attention_mask 默认是形状为 (batch,seq) 的全 1 矩阵, 后文假设用户传入的 attention_mask 要么是全 1, 要么是 None;

回到最外层的 GPT2LMHeadModel 类，它的 forward 函数获得 attention_mask 参数后直接传递给 GPT2Model.foward , 中间没有对 attenton_mask 做额外操作。

在 GPT2Model.forward 中有以下代码：

# GPT2Attention mask.
if attention_mask is not None:
    if batch_size <= 0:
        raise ValueError("batch_size has to be defined and > 0")
    attention_mask = attention_mask.view(batch_size, -1)
    # We create a 3D attention mask from a 2D tensor mask.
    # Sizes are [batch_size, 1, 1, to_seq_length]
    # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
    # this attention mask is more simple than the triangular masking of causal attention
    # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
    attention_mask = attention_mask[:, None, None, :]

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and the dtype's smallest value for masked positions.
    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
    attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min

首先，根据注释可以进一步确认用户输入的 attention_mask 应该是一个 (batch, length) 的二维 tensor, 其中只包括 0 和 1 ； 此外 attention_mask 会被展开成形状为 (batch, 1, 1, length) 的四维矩阵；

最后一句中，原始矩阵中的 1 会变成 0, 0 则变成接近 -inf, 因此它的意图是用加法（广播）来对注意力得分做 mask 后再计算 softmax, 使得那些 -inf 的值会得到接近 0 的注意力权重。

注意本文用 -inf 表示 torch 能表示的最接近负无穷的数，它的真实值如下：

print(torch.finfo(torch.float32).min)
print(torch.finfo(torch.float16).min)
print(torch.finfo().min)

-3.4028234663852886e+38
-65504.0
-3.4028234663852886e+38

如果用 -inf 符号去计算，可能会出现数值问题，应该避免。

torch.tensor(float("-inf")) + 12.2

tensor(-inf)

经过以上计算，当前的 attention_mask 要么是就是形状为 (batch,1,1,length) 的全 0 矩阵，要么还是 None.

GPT2Model.forward 会调用循环调用 GPT2Block.forward 继而调用 GPT2Attention.forward ，这三个函数都没有对 attention_mask 做额外处理，因此直接进入到最核心的 GPT2Attention._attn 函数

该函数首先计算完整的 scaled Dot-Product attention weight 值，也就 \( \text{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V \) 中的 \( \frac{QK^{T}}{\sqrt{d_{k}}} \):

attn_weights = torch.matmul(query, key.transpose(-1, -2))

 if self.scale_attn_weights:
     attn_weights = attn_weights / torch.full(
         [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device
     )

假设输入序列长度是 6, batch 是 2, 采用默认 12 个注意力头，那么：

• attn_weights 的形状是 torch.Size([2, 12, 6, 6]): 表示有 2 个样本，每个样本 12 个头，当前有 6 个 token, 因此得到的权重是 6x6
• q,k,v 形状为 torch.Size([2, 12, 6, 64]): 表示每个 token 对应 12 个 heads, 每个 head 是 64 维度
• attention_mask 的形状为 torch.Size([2, 1, 1, 6]), 内容都是 0 。（或为 None）

接着下来涉及对 attention mask 操作的代码是：

if not self.is_cross_attention:
    # if only "normal" attention layer implements causal mask
    query_length, key_length = query.size(-2), key.size(-2)
    causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
    mask_value = torch.finfo(attn_weights.dtype).min
    # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
    # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
    mask_value = torch.full([], mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
    attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)

if attention_mask is not None:
    # Apply the attention mask
    attn_weights = attn_weights + attention_mask

attn_weights = nn.functional.softmax(attn_weights, dim=-1)

对于自注意力层，以上第一个分支内代码会被执行，前文中就已经知道 self.bias 是一个三角矩阵, 形状是 [1,1,max_position,max_position], max_positions 为 6 时如下：

tensor([[[[ True, False, False, False, False, False],
        [ True,  True, False, False, False, False],
        [ True,  True,  True, False, False, False],
        [ True,  True,  True,  True, False, False],
        [ True,  True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True,  True]]]])

但如果当前输入的序列长度只有 3 的话，只需要选择左上角 3x3 的矩阵，也就是 bias[:,:,:3,:3]. 因此第三第四维度最多只需要取前 key_length 个元素，这才有如下代码：

causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] 

不过在推理时，每次模型只是自回归式地解码出一个词汇作为下个时刻的输入，query_length 始终是 1, 因此每一步只需要对单个 token 计算注意力，也就是取出 key_length - 1 : key_length 这一行注意力 mask，形状是 [1,1,1,key_length], 它会在执行 torch.where 时从 (batch,heads, 1,key_length) 的 attn_weights 上 mask 掉无效值。

注意 torch.where 也会广播，但这个场景中不需要广播

torch.where(torch.tensor([[1,0]]).bool(), torch.tensor([[1,2],[4,5]]), 100)

tensor([[  1, 100],
        [  4, 100]])

执行完判断分支里的代码后会得到了一个与 attention_mask 无关的 attn_weights ，这说明：

• 如果 attention_mask 是 None 的话，以下语句不会触发， attention_mask 对注意力计算没有任何贡献。

  if attention_mask is not None:
      attn_weights = attn_weights + attention_mask
  

• 如果 attention_mask 不是 None, 那么它就是 (batch,1,1,length) 形状的，需要和形状为 (batch, heads, query_length , key_length) 的 attn_weights 进行广播求和，结果矩阵的形状和后者保持一致。

结论是： attention_mask 只是用来 mask 某个 token 的，如 pad_token;

假设训练时，输入两个长度为 6 的序列, 那么 causal_mask 的形状就是 (1, 1, 6, 6)。如果不想要关注第 2 个 token, 那么输入的 attention_mask 形状为 (2,6) 的矩阵，其中每一行可以是：

[1,0,1,1,1,1]

经过 GPT2Model.forward 的处理后，变成了形状为 [2,1,1,6] 的矩阵，其中最后一维变成了

[0,-inf,0,0,0,0]

一旦广播，最后两维度就变成：

[
    [0,-inf,0,0,0,0],
    [0,-inf,0,0,0,0],
    [0,-inf,0,0,0,0],
    [0,-inf,0,0,0,0],
    [0,-inf,0,0,0,0],
    [0,-inf,0,0,0,0],
 ]

再与 attn_weights 相加的话，每个 token （行）和第 2 个（列） token 的注意力权重都变成 0 了。

接着对 attention 做 softmax 变成概率权重，然后对 value 做加权平均：

attn_weights = nn.functional.softmax(attn_weights, dim=-1)
attn_weights = self.attn_dropout(attn_weights)
attn_output = torch.matmul(attn_weights, value)

本节给工程实践上带来的启示是：如果想要修改自注意力的下三角 mask 矩阵，只能修改 GPT2Attention._attn （要么继承覆写，要么 monkey patch）函数，模型没有留出对外的注意力定制接口。

首先要清楚的是，GPT2 默认是纯 decoder 结构，没有 encoder 提供外部上下文， 因此理论上是不需要实现交叉注意力的，然而工程上没必要把该接口删除，以保持随时可以把 GPT2 和其他 encoder 模型组合起来变成 encoder-decoder 结构的灵活性。

在 GPT2Model.forward 中，如果有设置 add_cross_attention 为 True 以及提供了 encoder 的隐藏层，那么会激活对交叉注意力的准备流程：

# If a 2D or 3D attention mask is provided for the cross-attention
# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
if self.config.add_cross_attention and encoder_hidden_states is not None:
    encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
    encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
    if encoder_attention_mask is None:
        encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
    encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
else:
    encoder_attention_mask = None

从这里可以看到，首先 encoder_attention_mask 是 forward 参数提供的，假设 encoder 是一个 bert 类的模型，那么即使提供该参数，它也是形状为 (batch, seq_len) 的全 1 矩阵（不考虑 padding token），因此它和以下语句生成的 tensor 是一样的（所以如果不想定制交叉注意力 mask, 可以不用传入该参数）

if encoder_attention_mask is None:
    encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)

接着调用了 self.invert_attention_mask 函数，它来自 ModuleUtilsMixin 类（GPT2Model 继承到了它）其实现如下：

def invert_attention_mask(...):
    if encoder_attention_mask.dim() == 3:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
    if encoder_attention_mask.dim() == 2:
        encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
    encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
    encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min
    return encoder_extended_attention_mask

这些操作和自注意力的 attention_mask 的处理手段是一样的，都是先扩展维度，然后把 1 转为 0, 把 0 转为 -inf.

有一点不同的是，在第一个判断分支中表明 encoder_attention_mask 可以是三维结构，也就是说，用户可以通过函数参数来手动设置更复杂的交叉注意力 mask 。

接着调用 GPT2Block.forward ，它并没有对 mask 做额外处理。

再进入 GPT2Attention.forward, 在真正计算 attention 前，把 attenton_mask 赋值为 encoder_attention_mask, 也就是统一交叉注意力和自注意力的计算接口，区别只在于 key, value 的不同。

if encoder_hidden_states is not None:
    query = self.q_attn(hidden_states)
    key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
    attention_mask = encoder_attention_mask

之后就是调用 GPT2Attention._attn ，上文已经分析过该函数。交叉注意力场景下的区别在于，以下第一个条件分支是不会执行的， 不会构造 causal_mask, 对注意力的控制都是通过 attention_mask （也就是 encoder_attention_mask ）完成的。

if not self.is_cross_attention:
    # all skip

if attention_mask is not None:
    attn_weights = attn_weights + attention_mask

attn_weights = nn.functional.softmax(attn_weights, dim=-1)

本节给工程实践带来的启示是：如果要修改交叉注意力 mask 矩阵，可以直接在 GPT2Model.forward 参数里传入二维或三维的 encoder_attention_mask 矩阵：

• 二维矩阵可以使得编码器输入的某个 token 被 decoder 完全忽视
• 三维矩阵则可以定制 decoder 里各 token 对 encoder 中各 token 的关注程度。

在后文介绍 EncoderDecoderModel 时还会再次从编码器解码交互的角度梳理一遍交叉注意力的计算过程。

由于 past_key_values 的存在，在解码时，每次只需要传入最后一个 token_id, GPT2LMHeadModel 的 prepare_inputs_for_generation 函数对此进行了处理：

def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
    token_type_ids = kwargs.get("token_type_ids", None)
    # only last token for inputs_ids if past is defined in kwargs
    if past_key_values:
        input_ids = input_ids[:, -1].unsqueeze(-1)
        if token_type_ids is not None:
            token_type_ids = token_type_ids[:, -1].unsqueeze(-1)

    attention_mask = kwargs.get("attention_mask", None)
    position_ids = kwargs.get("position_ids", None)

    if attention_mask is not None and position_ids is None:
        # create position_ids on the fly for batch generation
        position_ids = attention_mask.long().cumsum(-1) - 1
        position_ids.masked_fill_(attention_mask == 0, 1)
        if past_key_values:
            position_ids = position_ids[:, -1].unsqueeze(-1)

    # skip
    model_inputs.update(
        {
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache"),
            "position_ids": position_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
        }
    )
    return model_inputs

当 past_key_values 存在的时候，下一次推理的输出只会取最后一个 token 的最后一项： input_ids[:, -1].unsqueeze(-1). 同样如果 position_ids 存在，也只取最后一个元素： position_ids[:, -1].unsqueeze(-1)

与之对应的， GPT2Attention._attn 函数 一节提到，在做自注意力的时候， causal_mask 矩阵不再是一个方阵，它从形状从 (1, 1, key_length, key_length) 变成了 (1, 1, 1, key_length), 后者在参与计算时核心是以下两段代码：

attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)
# skip...

if attention_mask is not None:
    attn_weights = attn_weights + attention_mask

causal_mask 的形状是 (1,1,1,key_length), attention_mask 是 None 则不会参与计算，如果不是 None 那么形状是 (batch,1,1,key_length), atten_weights 形状为 (batch, 12, 1, key_length) 广播出来还是 (batch, 12, 1, key_length), 这些形状和传入的 input_ids 和 posint_ids 都是匹配的。

prepare_inputs_for_generation 函数是在具体的解码策略函数里调用，比如以下是从 GenerationMixin.greedy_search 函数中提取出的相关部分，在循环解码的每一步，先调用输入准备函数，然后执行 self(*kwargs) 从而触发前向传播

while True:
    # ...
    # prepare model inputs
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

    # forward pass to get next token
    outputs = self(
        **model_inputs,
        return_dict=True,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
    )
    next_token_logits = outputs.logits[:, -1, :]
    next_tokens_scores = logits_processor(input_ids, next_token_logits)
    next_tokens = torch.argmax(next_tokens_scores, dim=-1)
    input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
    model_kwargs = self._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
            )

首先要清楚，解码时第一次的时候传入 prepare_inputs_for_generation 的 **model_kwargs 中是没有 past_key_values 的，也就是为默认值 None.

有了解码的结果 outputs 后通过以下两个函数对其进行后处理：

def _update_model_kwargs_for_generation(self, outputs: ModelOutput, model_kwargs: Dict[str, Any]) -> Dict[str, Any]:
    # update past_key_values
    model_kwargs["past_key_values"] = self._extract_past_from_model_output(
        outputs, standardize_cache_format=standardize_cache_format
    )
    #skip
def _extract_past_from_model_output(self, outputs: ModelOutput, standardize_cache_format: bool = False):
    past_key_values = None
    if "past_key_values" in outputs:
        past_key_values = outputs.past_key_values
    #skip
    return past_key_values

这个时候，outputs 中已经有了非 None 的 past_key_values, 它从 outputs 中提取出，通过 _update_model_kwargs_for_generation 函数又放回到 model_kwargs 中。

past_key_values 是在执行以下命令后返回到 outputs 里的

outputs = self(
    **model_inputs,
    return_dict=True,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
)

这里 self 是 GPT2LMHeadModel, 但 GPT2LMHeadModel 的 forward 里没有对 past_key_values 做处理，直接就传入到 self.transformer, 也就是 GPT2Model.forward 中

transformer_outputs = self.transformer(
    input_ids,
    past_key_values=past_key_values,
    #skip
    )

GPT2Model.forward 中相关代码如下：

if past_key_values is None:
    past_length = 0
    past_key_values = tuple([None] * len(self.h))
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])

这里如果 past_key_value 为 None, 那么会创建一个长度等于 len(self.h) 的全为 None 的 tuple,

接着检查如果没有传入 position_ids （大部分时候都不会手动传入）的话，会根据 past_key_values 维度创建 position_ids ， past_key_values[0][0].size(-2) 是已经生成的序列的长度。考虑两种情况：

• 训练的时候，那么每次前向传播只是手动输入 input_ids 和 labels, 其他参数默认, 因此每次 past_key_values 都是 None, past_length 始终 0, 于是 position_ids 计算如下

  print(torch.arange(0, 3))
  print(torch.arange(0, 3).unsqueeze(0).view(-1, 3))
  

  tensor([0, 1, 2])
  tensor([[0, 1, 2]])
  

  也就是覆盖了整个输入序列。

• 解码的时候，假设 past_key_values 中 key 和 value 已经存了 2 个 token 的缓存，那么 past_length 等于 2. 当前为了预测第三个 token, 传入的 input_id 实际是一个 (batch,1) 形状的矩阵（后文会提到是哪里对输入进行切分的），因此 input_shape[-1] 等于 1, position_ids 计算为：

  print(torch.arange(2, 1 + 2))
  print(torch.arange(2, 1 + 2).unsqueeze(0).view(-1, 1))
  

  tensor([2])
  tensor([[2]])
  

  计算出来的 position_ids 是最后这个 token 的，形状为 (batch, 1) 。

回到 past_key_values, 在 GPT2Model.forward 中，它的以 layer_past 之名分别输入到了各层 GPT2Block 的 forward 中：

for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
    #skip
    outputs = block(layer_past, ...)

在 GPT2Block 中，layer_past 传递给了自注意力层 GPT2Attention 的 forward 里：

attn_outputs = self.attn(
            hidden_states,
            layer_past=layer_past,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )

而在计算交叉注意力时，没有传入 layer_past

outputs = attn_outputs[1:] # 获得 present (attentions)

if encoder_hidden_states is not None:
    # skip
    residual = hidden_states
    hidden_states = self.ln_cross_attn(hidden_states)
    cross_attn_outputs = self.crossattention(
        hidden_states,
        attention_mask=attention_mask,
        head_mask=head_mask,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_attention_mask,
        output_attentions=output_attentions,
    )
    attn_output = cross_attn_outputs[0]

由于 GPT2Attention.forward 中 layer_past 的默认值为 None, 这大致说明交叉注意力是不需要 key value 缓存的，具体细节见下节。

GPT2Attention.forward 形参如下：

# GPT2Attention
def forward(
     self,
     hidden_states: Optional[Tuple[torch.FloatTensor]],
     layer_past: Optional[Tuple[torch.Tensor]] = None,
     # skip

接着是 layer_past 相关代码：

query = self._split_heads(query, self.num_heads, self.head_dim)
key = self._split_heads(key, self.num_heads, self.head_dim)
value = self._split_heads(value, self.num_heads, self.head_dim)

if layer_past is not None:
    past_key, past_value = layer_past
    key = torch.cat((past_key, key), dim=-2)
    value = torch.cat((past_value, value), dim=-2)

if use_cache is True:
    present = (key, value)
else:
    present = None

# skip
attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

在切分出 query,key,value 后，如果 layer_past 不是 None, 就取出 past_key 和 past_value, 继续把当前计算出的 key value 在 dim=-2 维度拼接，该维度的意义是 decoder_sequence_length, 拼接完维度就是 (10,12,6,64) 。

解码时 use_cache 默认为 True, 因此将 key,value 放到 present 变量继而放到 outputs[1] 中保存下来， present 的形状可以看作是 (2,10,12,6,64), 尽管它是个 tuple.

layer_past 会为 None 的情况再次回顾:

• 训练时，我们直接调用 GPT2LMHeadModel 的 forward 函数，且根本不会手动传入 past_key_values 参数，又上文可知 past_key_values 就是一个全 None 的 tuple, layer_past 也就是 None
• 交叉注意力层，不会传入 layer_past, 因此用默认值 None
• 解码第一个 token 的时候

GPT2Attention 的 foward 中计算完自注意力后， (key, value) 会放到 outputs 的第二个元素里返回到 GPT2Block 中

#skip merge_heads, c_proj and  dropout
outputs = (attn_output, present)
if output_attentions:
    outputs += (attn_weights,)
return outputs  # a, present, (attentions)

GPT2Attention.forward 返回的 outputs 会在 GPT2Block 中以 attn_outputs 变量会接收，然后其第二个元素之后的内容放到 outputs 中，注意 attn_outputs[1] 就是 present, 也就是本层的 (key, value) 缓存对。

attn_outputs = self.attn(layer_past=layer_past,)
outputs = attn_outputs[1:]

然后执行以下 use_cache 为 True 的分支

if use_cache:
    outputs = (hidden_states,) + outputs
else:
    outputs = (hidden_states,) + outputs[1:]

return outputs  # hidden_states, present, (attentions, cross_attentions)

也就是说 GPT2Block 返回的结果的第二个是本曾的 (key, value) 缓存对

GPT2Model.forward 中所有与 past_key_values 相关代码如下：

if past_key_values is None:
    past_length = 0
    past_key_values = tuple([None] * len(self.h))
# skip
presents = () if use_cache else None

for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
    # skip
    outputs = block(layer_past, ...)
    if use_cache is True:
        presents = presents + (outputs[1],)
# skip
# 当前得到 n_layers 个 (key,value) 对, 放在结果第二个位置被返回
return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

GPT2Model.forward 拿到每一层 GPT2Block 的返回值后取出了 outputs[1] 加入到 presents 变量，构成了一个长度为 len(self.h) 的 tuple, 然后包装在 BaseModelOutputWithPastAndCrossAttentions 的第二个变量 past_key_values 中。

在 GPT2LMHeadModel.forward 中的 GPT2Model 是以 self.transformer 对象执行的，其返回的 transformer_outputs 变量的各个参数又被包装到 CausalLMOutputWithCrossAttentions 里然后再次返回

transformer_outputs = self.transformer(
    input_ids,
    past_key_values=past_key_values,
    #skip
)

#skip
return CausalLMOutputWithCrossAttentions(
    loss=loss,
    logits=lm_logits,
    past_key_values=transformer_outputs.past_key_values,
    hidden_states=transformer_outputs.hidden_states,
    attentions=transformer_outputs.attentions,
    cross_attentions=transformer_outputs.cross_attentions,
)

我们再回到第一节调用 Model.generate 的场景，代码如下：

while True:
    # ...
    # prepare model inputs
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

    # forward pass to get next token
    outputs = self(
        **model_inputs,
        return_dict=True,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
    )
    next_token_logits = outputs.logits[:, -1, :]
    next_tokens_scores = logits_processor(input_ids, next_token_logits)
    next_tokens = torch.argmax(next_tokens_scores, dim=-1)
    input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
    model_kwargs = self._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
            )

outputs 就是一个 CausalLMOutputWithCrossAttentions 类型的对象。其中包括 past_key_values 变量，它在 _update_model_kwargs_for_generation 函数里被取出又放回到 model_kwargs 之中，进入下一次循环采样（如不熟悉，可以回顾本章第一节 input_ids 的准备 的内容来重新看清整个流程。）

以 greedy_search 为例，以上循环结束最后返回结果是：

if return_dict_in_generate:
    if self.config.is_encoder_decoder:
        return GreedySearchEncoderDecoderOutput(
            sequences=input_ids,
            scores=scores,
            encoder_attentions=encoder_attentions,
            encoder_hidden_states=encoder_hidden_states,
            decoder_attentions=decoder_attentions,
            cross_attentions=cross_attentions,
            decoder_hidden_states=decoder_hidden_states,
        )
    else:
        return GreedySearchDecoderOnlyOutput(
            sequences=input_ids,
            scores=scores,
            attentions=decoder_attentions,
            hidden_states=decoder_hidden_states,
        )
else:
    return input_ids

两个 if 分支里不会再包含 past_key_values, 因此最后缓存是不会被 generate 函数返回的，在循环采样结束后会自动释放掉。

另外，以上的 input_ids 实际是解码出来的结果 token id 序列，之所以叫做 input_ids, 是因为循环中它不断作为下一次解码的输入，因此直到停止也叫做 input_ids。而我们最后通过 outputs.sequences 就可以取出最终的结果

在 generate 函数中，如果有 num_beams 参数（以及没有其他采样策略选项），那么会调用 beam_search 函数，具体的触发条件如下：

is_beam_gen_mode = (
    (generation_config.num_beams > 1)
    and (generation_config.num_beam_groups == 1)
    and generation_config.do_sample is False
    and not is_constraint_gen_mode
    and not is_contrastive_search_gen_mode
)

为了能够维持路径得分，先实例化一个 BeamSearchScorer 类的对象，它继承自 BeamScorer 抽象类，用于跟踪搜索过程中的状态：

beam_scorer = BeamSearchScorer(
    batch_size=batch_size,
    num_beams=generation_config.num_beams,
    device=inputs_tensor.device,
    length_penalty=generation_config.length_penalty,
    do_early_stopping=generation_config.early_stopping,
    num_beam_hyps_to_keep=generation_config.num_return_sequences,
    max_length=generation_config.max_length,
)

该类对 batch 里每个样本构建一个 BeamHypotheses 对象来追踪 beam 的状态，这些对象在属性 _beam_hyps 中：

self._beam_hyps = [
            BeamHypotheses(
                num_beams=self.num_beams,
                length_penalty=self.length_penalty,
                early_stopping=self.do_early_stopping,
                max_length=max_length,
            )
            for _ in range(batch_size)
        ]

注意每个样本搜索时要在 num_beams 个"平行世界" 里同时进行探索, 因此接下来调用 _expand_inputs_for_generation 对 input_ids 、attention_mask 等 tensor 类型的模型输入做扩充：

def _expand_inputs_for_generation(
    expand_size: int = 1,
    is_encoder_decoder: bool = False,
    input_ids: Optional[torch.LongTensor] = None,
    **model_kwargs,
) -> Tuple[torch.LongTensor, Dict[str, Any]]:
    """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""

    def _expand_dict_for_generation(dict_to_expand):
        for key in dict_to_expand:
            if dict_to_expand[key] is not None and isinstance(dict_to_expand[key], torch.Tensor):
                dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
        return dict_to_expand

    if input_ids is not None:
        input_ids = input_ids.repeat_interleave(expand_size, dim=0)

    model_kwargs = _expand_dict_for_generation(model_kwargs)

    if is_encoder_decoder:
        # Skip some code
        model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])

    return input_ids, model_kwargs

例如 beam_num = 2 的时候，需要对输入中所有 tensor 都复制出额外的 1 份从而扩大 batch 为 2，如果原本是 [1,2,3] 当前会变成 [1,1,2,2,3,3]. 这样就有 beam_num 个单独的向前传播路径。 当然，输入的一般是形状为 (batch, seq_len) ，并且第一个 token_id 一般是 101, 对应 CLS token (注：CLS 一般是在 BERT 类编码器模型中用来做句子分类的 token 的代号, 在 GPT 类生成模型中，一般用 bos, begin of sentence 来表示，但本文还是沿用 CLS 这个标记了, 另外，基于 gpt-2 的模型有很多，不同模型可能使用不同的分词方式和词表，因此 bos 对应的 id 也不同，gpt2 官方实现里该 token 编码是 50256，但本文用的是另一个以 101 作为 bos id 的模型进行示例的) 。

num_beams = 2
print(torch.tensor([1,2,3]).repeat_interleave(num_beams, dim=0))
print(torch.tensor([[101, 23]]).repeat_interleave(num_beams, dim=0))

tensor([1, 1, 2, 2, 3, 3])
tensor([[101,  23],
        [101,  23]])

被扩充后的 input_ids 的第一个维度被称为 batch_beam_size, 它等于 batch_size*num_beams

扩充 input_ids 之后，则真正调用 beam_search 函数：

# 13. run beam search
return self.beam_search(
    input_ids,
    beam_scorer,
    logits_processor=logits_processor,
    stopping_criteria=stopping_criteria,
    pad_token_id=generation_config.pad_token_id,
    eos_token_id=generation_config.eos_token_id,
    output_scores=generation_config.output_scores,
    return_dict_in_generate=generation_config.return_dict_in_generate,
    synced_gpus=synced_gpus,
    **model_kwargs,
)

beam_search 函数的核心代码为可以分成三部分来看：

1. 初始化得分矩阵记录每个路径得分，进入循环，将输入传到模型 forward 中计算出下一时刻各个 token 的得分
2. 用下一时刻各 token 得分更新当前的得分矩阵
3. 修剪当前的得分矩阵, 只保留得分最高的 id 对应的那些路径

以下是第一个阶段：

• 首先注释里解释了把被复制的样本的得分设置为 -1e9 使得第一次计算的时候它们的得分不会被考虑进去，这是解决上一节中提到的问题的代码实现。
• prepare_inputs_for_generation 已经在上文提到过，主要逻辑是，如果 past_key_values 存在，则只取 input_ids 、position_ids 里最后一个 token 进行推理。

# initialise score of first beam with 0 and the rest with -1e9. This makes sure that only
# tokens of the first beam are considered to avoid sampling the exact same tokens across all beams.
beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
beam_scores[:, 1:] = -1e9
beam_scores = beam_scores.view((batch_size * num_beams,))
while True:
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

    outputs = self(
        **model_inputs,
        return_dict=True,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
    )

    # skip part2: score calculate
    # skip part3: reshape

用一个简单例子来说明（batch=3, beam_size = 2, 词库大小为 5，初始输入为 CLS ）

batch_size, num_beams = 3, 2
beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float)
beam_scores[:, 1:] = -1e9
beam_scores = beam_scores.view((batch_size * num_beams,))
print(beam_scores)

tensor([ 0.0000e+00, -1.0000e+09,  0.0000e+00, -1.0000e+09,  0.0000e+00,
        -1.0000e+09])

这里可以看到，每个样本的第二 beam 得分都非常小。

有了模型的 outputs 输出，关键就是获取下一个时刻各个 token 的得分、累加到历史得分里并找出最高的 num_beam 个得分

while True:
    # skip 
    outputs = self(**model_inputs,...)# model forward

    next_token_logits = outputs.logits[:, -1, :]
    next_token_scores = nn.functional.log_softmax(
        next_token_logits, dim=-1
    )  # (batch_size * num_beams, vocab_size)

    next_token_scores_processed = logits_processor(input_ids, next_token_scores)
    next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)

    vocab_size = next_token_scores.shape[-1]
    next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)

    # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
    next_token_scores, next_tokens = torch.topk(
        next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
    )

    next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
    next_tokens = next_tokens % vocab_size

    # skip part3: reshape

用上文中的例子来解释每一行代码：

• 第一次循环时，outputs.logits 是一个形状为 [6, 1, 5] 的 tensor, next_token_logits 的形状为 [6, 5] , 计算 softmax 并求 log 之后得到的 next_token_scores 继续保持 [6, 5] 的形状。 假设是如下得分（注意由于第一次循环时 input_ids 都是 CLS, 因此组内各 beam 得分是一样的）：

  torch.manual_seed(1)
  input_ids = torch.tensor([[101]]*6)
  next_token_scores = torch.randn((3, 5)).repeat_interleave(num_beams, dim=0)
  print(input_ids)
  print(next_token_scores)
  

  tensor([[101],
          [101],
          [101],
          [101],
          [101],
          [101]])
  tensor([[ 0.6614,  0.2669,  0.0617,  0.6213, -0.4519],
          [ 0.6614,  0.2669,  0.0617,  0.6213, -0.4519],
          [-0.1661, -1.5228,  0.3817, -1.0276, -0.5631],
          [-0.1661, -1.5228,  0.3817, -1.0276, -0.5631],
          [-0.8923, -0.0583, -0.1955, -0.9656,  0.4224],
          [-0.8923, -0.0583, -0.1955, -0.9656,  0.4224]])
  

  肉眼上看，各样本最高得分的两个 token 下标分别是： (0, 3), (2, 0), (4, 1);

• logits_processor 用于对 logit 做修改，可以用来限制采样，比如假设有一些敏感词汇要过滤，那么可以通过 processor 来把这些词的得分变得很低。这里假设过滤掉最后一个词，那么 logits_processor 可以看作以下函数（尽管它是一个类，但用 __call__ 方法转成了带参数的函数）

  def logits_processor(input_ids, next_token_scores):
      next_token_scores[:,-1] = -1e9
      return next_token_scores
  
  next_token_scores_processed = logits_processor(input_ids, next_token_scores)
  print(next_token_scores_processed)
  

  tensor([[ 6.6135e-01,  2.6692e-01,  6.1677e-02,  6.2132e-01, -1.0000e+09],
          [ 6.6135e-01,  2.6692e-01,  6.1677e-02,  6.2132e-01, -1.0000e+09],
          [-1.6613e-01, -1.5228e+00,  3.8168e-01, -1.0276e+00, -1.0000e+09],
          [-1.6613e-01, -1.5228e+00,  3.8168e-01, -1.0276e+00, -1.0000e+09],
          [-8.9229e-01, -5.8250e-02, -1.9551e-01, -9.6564e-01, -1.0000e+09],
          [-8.9229e-01, -5.8250e-02, -1.9551e-01, -9.6564e-01, -1.0000e+09]])
  

  next_token_scores_processed 的形状仍然是 [6, 5], 不过由于最后一个元素被过滤掉了，此时各样本最高得分的两个 token 下标分别是： (0, 3), (2, 0), (1, 2);

• 接着把以上得分广播并累加到路径得分 beam_scores 上：

  next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)
  

  tensor([[ 6.6135e-01,  2.6692e-01,  6.1677e-02,  6.2132e-01, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-1.6613e-01, -1.5228e+00,  3.8168e-01, -1.0276e+00, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-8.9229e-01, -5.8250e-02, -1.9551e-01, -9.6564e-01, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09]])
  

  这实际可以简写成：

  next_token_scores_processed + beam_scores[:, None]
  

  tensor([[ 6.6135e-01,  2.6692e-01,  6.1677e-02,  6.2132e-01, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-1.6613e-01, -1.5228e+00,  3.8168e-01, -1.0276e+00, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-8.9229e-01, -5.8250e-02, -1.9551e-01, -9.6564e-01, -1.0000e+09],
          [-1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09]])
  

  还是保留原始代码：

  next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(next_token_scores)
  

  此时 next_token_scores 还是形状为 [6,5] 的 tensor, 第一次循环中，每个样本中额外复制出来的 beam 的得分都是无穷小，只有第一个 beam 保存了各个 token 的得分

• 为了对所有路径中的节点统计得分，要把 [6,5] 维度变成 [3, 10], 也就是每个样本的所有 beam 里的 token 得分都放在一个维度：

  vocab_size = next_token_scores.shape[-1] # 5
  next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
  print(next_token_scores)
  

  tensor([[ 6.6135e-01,  2.6692e-01,  6.1677e-02,  6.2132e-01, -1.0000e+09,
           -1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-1.6613e-01, -1.5228e+00,  3.8168e-01, -1.0276e+00, -1.0000e+09,
           -1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09],
          [-8.9229e-01, -5.8250e-02, -1.9551e-01, -9.6564e-01, -1.0000e+09,
           -1.0000e+09, -1.0000e+09, -1.0000e+09, -1.0000e+09, -2.0000e+09]])
  

• 接着对每个样本取出得分最高的两个元素的得分以及对应下标，正常思路就是在第二个维度上进行排序并取前 2 个，。 不过这里取了 2*num_beams 个最高得分，有一倍的冗余。

  next_token_scores, next_tokens = torch.topk(
        next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
    )
  print(next_token_scores)
  print(next_tokens)
  

  tensor([[ 0.6614,  0.6213,  0.2669,  0.0617],
          [ 0.3817, -0.1661, -1.0276, -1.5228],
          [-0.0583, -0.1955, -0.8923, -0.9656]])
  tensor([[0, 3, 1, 2],
          [2, 0, 3, 1],
          [1, 2, 0, 3]])
  

  next_token_scores 维度变成了 [3, 4] 也就是 [batch_size, 2*num_beams], 另外next_tokens 每行的前 2 个和上文中肉眼识别的结果 (0, 3), (2, 0), (1, 2) 是一样的。

• 为了确定最高得分来自于哪个 beam 束，用 torch.div 和求模的方式来获得 beam_idx 和 token_id. (不过在以上例子中，由于是第一次循环，其他 beam 的值都接近负无穷，因此 topk 都来自于 beam_idx=0 的路径)

  next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
  next_tokens = next_tokens % vocab_size
  print(next_indices)
  print(next_tokens)
  

  tensor([[0, 0, 0, 0],
          [0, 0, 0, 0],
          [0, 0, 0, 0]])
  tensor([[0, 3, 1, 2],
          [2, 0, 3, 1],
          [1, 2, 0, 3]])
  

  这两个 tensor 形状都是 [batch_size, 2*num_beams]

上一节内容已经展示了 beam search 算法核心的 80% 了，但还有很多工程上的工作，例如要根据以上得到的各种得分以及排序结果来更新 beam_scores 、input_ids 等进入下一次循环所需要的变量，还要判断各个 beam 中是否已经采样出了结束符号（EOS）从而要停止某些 beam 的采样。

以下代码中先调用 beam_scorer.process 来更新 beam_scores ，源代码里有 "#stateless" 注释表明该函数的执行是无状态的，意味着函数的返回值 beam_outputs 完全是由函数输入所决定的，不需要 beam_scorer 对象里额外的其他信息。 因为各个 beam 的得分已经都广播到 next_token_scores 中了，这里要做的是从该变量里”提取“出新的 beam_scores 。（最初看代码的时候我以为是 beam_scorer 维持了各个路径的历史得分，然后通过 process 函数去去更新路径得分，但这种做法就不是 stateless, 而是 stateful 了。 beam_scorer 实际的主要作用是跟踪 beam 是否停止，后文将具体介绍）

while True:
    # skip 

    # stateless
    beam_outputs = beam_scorer.process(
        input_ids, #(batch_size, seq_len)
        next_token_scores, #(batch_size, 2*num_beams)
        next_tokens, #(batch_size, 2*num_beams)
        next_indices, #(batch_size, 2*num_beams)
        pad_token_id=pad_token_id,
        eos_token_id=eos_token_id,
        beam_indices=beam_indices, # None
    )
    beam_scores = beam_outputs["next_beam_scores"]
    beam_next_tokens = beam_outputs["next_beam_tokens"]
    beam_idx = beam_outputs["next_beam_indices"]

beam_scorer.process 中先初始化 beam 的 scores, tokens 和 indices, 如下，这里 self.group_size 默认等于 num_beams

cur_len = input_ids.shape[-1] + 1  # add up to the length which the next_scores is calculated on
batch_size = len(self._beam_hyps)
next_beam_scores = torch.zeros((batch_size, self.group_size)) # (batch_size, num_beam)
next_beam_tokens = torch.zeros((batch_size, self.group_size))
next_beam_indices = torch.zeros((batch_size, self.group_size))

接着对 batch 里每个样本进行处理：

for batch_idx, beam_hyp in enumerate(self._beam_hyps): # #batch_size
    if self._done[batch_idx]:
        next_beam_scores[batch_idx, :] = 0
        next_beam_tokens[batch_idx, :] = pad_token_id
        next_beam_indices[batch_idx, :] = 0
        continue

以上先检查当前样本是否可以停止采样，判断标准是 self._done[batch_idx], 该变量在 init 里初始化为：

self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool)

接着对每个样本的每个搜索路径 beam 进行检查：

for batch_idx, beam_hyp in enumerate(self._beam_hyps):
    #skip
    # next tokens for this sentence
    beam_idx = 0
    for beam_token_rank, (next_token, next_score, next_index) in enumerate(
        zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx])
    ):
        batch_beam_idx = batch_idx * self.group_size + next_index
        # add to generated hypotheses if end of sentence
        if (eos_token_id is not None) and (next_token.item() in eos_token_id):
            #skip
        else:
            # add next predicted token since it is not eos_token
            next_beam_scores[batch_idx, beam_idx] = next_score
            next_beam_tokens[batch_idx, beam_idx] = next_token
            next_beam_indices[batch_idx, beam_idx] = batch_beam_idx
            beam_idx += 1

        # once the beam for next step is full, don't add more tokens to it.
        if beam_idx == self.group_size:
            break

以上核心是 if (eos_token_id 判断的两个分支，本节先分析 if 判断失败的情况，该分支表明当前采样出的并不是 eos, 因此采样继续循环，所以这里做的就是把 socre, token 和 beam_idx 重新添加回到对应的 next_beam_* tensor 中，而 if beam_idx == self.group_size: 判断意味着，只有前 num_beam 个元素会被加入到这些 tensor 中（这与 next_beam_scores 等变量形状是一样的），而 process 函数返回类型如下：

return UserDict(
     {
         "next_beam_scores": next_beam_scores.view(-1),
         "next_beam_tokens": next_beam_tokens.view(-1),
         "next_beam_indices": next_beam_indices.view(-1),
     }
 )

beam_scores = beam_outputs["next_beam_scores"]
beam_next_tokens = beam_outputs["next_beam_tokens"]
beam_idx = beam_outputs["next_beam_indices"]

这意味着经过 process 函数后，如果每个搜索路径都没有遇到终点, 那么以上三个变量经历的操作类似于：

beam_scores = next_token_scores[:, :num_beams].reshape(-1)
beam_next_tokens = next_tokens[:, :num_beams].reshape(-1)
beam_idx = next_indices[:, :num_beams].reshape(-1)
print(beam_scores)
print(beam_next_tokens)
print(beam_idx)

tensor([ 0.6614,  0.6213,  0.3817, -0.1661, -0.0583, -0.1955])
tensor([0, 3, 2, 0, 1, 2])
tensor([0, 0, 0, 0, 0, 0])

以上结果是合理的，比如所有 tensor 都回到了 (batch_size * num_beams,) 形状，它们表示的都是当前 step 下的各类状态的切片，因此第二个时间维度形状为 1.

while True:
    # skip part1 model forward
    # skip part2 score calculate

    input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)

    model_kwargs = self._update_model_kwargs_for_generation(
        outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
    )
    if model_kwargs["past_key_values"] is not None:
        model_kwargs["past_key_values"] = self._reorder_cache(model_kwargs["past_key_values"], beam_idx)

本节的主角是 beam_idx, 当前的例子中，由于是第一次循环，各个样本上最高分的 beam 都来自原始的第一个 beam, 因此是一个全 0 的 tensor, 每个位置的意义表示的是当前的"胜利者" beam 的下标，而只有胜利者路径会参与下一次采样。 在后续的采样中，各个 beam 都有自己的得分，那么其他 beam 可能会有更高的得分，这个时候 beam_idx 就不是全 0, 而可以是 0 到 num_beam-1 中任何一个整数值。

拼接完的 input_ids

input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
print(input_ids)

tensor([[101,   0],
        [101,   3],
        [101,   2],
        [101,   0],
        [101,   1],
        [101,   2]])

接着调用 _update_model_kwargs_for_generation ，该函数上文提到过，后文也继续会提到，这里要知道的它的关键作用就是给 attention_mask 、 token_type_ids 拼接一个 tensor([1]) ，使得长度变长。

最后，和 input_ids 类似，要对 past_key_values 也根据 beam_idx 进行重新选择，只留下"胜利者"所属的历史:

@staticmethod
def _reorder_cache(
    past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
) -> Tuple[Tuple[torch.Tensor]]:
    return tuple(
        tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
        for layer_past in past_key_values
    )

该函数是在模型类里具体实现的，比如在 GPT2LMHeadModel 和 GPT2DoubleHeadsModel 中都有。

前文都是理解 beam search 是如何不断循环采样的，本节则理解算法如何停止，这比如何继续采样更复杂一点，所有逻辑都在 BeamSearchScorer 中，要考虑：

• 有些 beam 如果先遇到了结束符号，如何和其他 beam 一起继续采样还是自己先停止？

  回到上文提到的 BeamSearchScorer.proces 里遇到 eos 的 if 分支，代码如下：

  # if beam_token does not belong to top num_beams tokens, it should not be added
  is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size
  if is_beam_token_worse_than_top_num_beams:
      continue
  if beam_indices is not None:
      beam_index = beam_indices[batch_beam_idx]
      beam_index = beam_index + (batch_beam_idx,)
  else:
      beam_index = None
  
  beam_hyp.add(
      input_ids[batch_beam_idx].clone(),
      next_score.item(),
      beam_indices=beam_index,
  )
  

  第一行表示如果这个 token 不是在前 num_beam 里，那么不考虑（因为前文把 2*num_beams 个最高分 token 都保留了）。而 beam_indices 默认就是 None, 因此真正执行的是 beam_hyp.add 函数：

  def add(self, hyp: torch.LongTensor, sum_logprobs: float, beam_indices: Optional[torch.LongTensor] = None):
      score = sum_logprobs / (hyp.shape[-1] ** self.length_penalty)
      if len(self) < self.num_beams or score > self.worst_score:
          self.beams.append((score, hyp, beam_indices))
          if len(self) > self.num_beams:
              sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)])
              del self.beams[sorted_next_scores[0][1]]
              self.worst_score = sorted_next_scores[1][0]
          else:
              self.worst_score = min(score, self.worst_score)
  

  这里首先计算加了长度惩罚后的得分，length_penalty 越大，对长句得分惩罚越大，越小的话得分鼓励长句（可以小于 0），默认是 1. 接着之后的代码实际是一个最大堆的更新算法（只不过每次更新是重新调用一次排序，而不是进行 heapfy，不过 beam size 一般不大，效率瓶颈不在这里而是在神经网络，因此可能没必要过度优化），self.beams 中最多只能有 num_beams 个元素，并且只保留得分最高的元素，维持 worst_score 是为了快速过滤低分的 beam.

  这里看到 self.beams 保存的是某个样本中已经结束的最高得分的 beam. 在这种情况下，next_beam_* 不会更新，也就是得分已经定格了，不过上一轮的 token 还是会拼接到 input_ids 上 比如，有样本的某个 beam 的 input_ids 是 [1,2,3], 接着它预测出了 102, 这是一个结束符号（eos），那么这条路径会被加入到 beam_hyp.beams 中封存起来，但由于其他 beam 还没有结束，为了保持 tensor 形状，input_ids 变成 [1,2,3,3], 仍然继续采样下去。但最后我们从 beam_hyp 里取最终返回的序列。

• 某个样本里所有 beam 都结束了，但 batch 里其他样本没结束，如何继续采样？

  相关函数：BeamSearchScorer.is_done

• 所有样本都结束了，如何返回结果

  相关函数：BeamSearchScorer.finalize

由于训练场景遵循和解码一样的输入要求，因此先考虑解码场景。

GenerationMixin.generate 函数有一个 inputs 参数，它会根据是纯解码器模型调用还是编码解码器调用来决定它是给编码器还是解码器的，纯解码器结构，它对应的就是解码器的 input_ids, 对于编码解码结构则默认是给编码器的输入序列。

不过为了防止混乱，可以不使用该参数，而是手动指定 input_ids 和 decoder_input_ids 明确表示是分别给编码器和解码器的，如下：

output_ids = model.generate(
            input_ids=encoder_input_ids
            pad_token_id=decoder_tokenizer.pad_token_id,
            eos_token_id=decoder_tokenizer.eos_token_id,
            decoder_input_ids=decode_input_ids
            logits_processor=processors,)

相关函数执行顺序：

1. _prepare_model_inputs: 它是处理编码器的输入的函数， 其核心代码如下：

   inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
               inputs, generation_config.bos_token_id, model_kwargs)
   

   注意 _prepare_model_inputs 输入的是 generate 的 inputs 参数，但是在 _prepare_model_inputs 中有一段代码如下：

   model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None or k != input_name}
   inputs_kwarg = model_kwargs.pop(input_name, None)
   if inputs_kwarg is not None and inputs is not None:
       raise ValueError(
           f"`inputs`: {inputs}` were passed alongside {input_name} which is not allowed."
           f"Make sure to either pass {inputs} or {input_name}=..."
       )
   elif inputs_kwarg is not None:
       inputs = inputs_kwarg
   

   这里由于 input_ids 是在 model_kwargs 中的，而编码器的 input_name 是 input_ids, 因此 inputs 就被赋值成了 input_ids

   接着，_prepare_model_inputs 会调用 _maybe_initialize_input_ids_for_generation, 其核心函数如下：

   if inputs is not None:
       return inputs
   
   encoder_outputs = model_kwargs.get("encoder_outputs")
   if self.config.is_encoder_decoder and encoder_outputs is not None:
       # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding
       shape = encoder_outputs.last_hidden_state.size()[:-1]
       return torch.ones(shape, dtype=torch.long, device=self.device) * -100
   
   if bos_token_id is None:
       raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")
   
   #skip some codes
   return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id
   

   可以看到，如果用户输入了 inputs 或 (inputs_ids) 那么直接返回，但是如果用户手动传入了编码器输出（而不是自动执行 encoder 得到的），那么会生成一个全是 -100 的矩阵，形状是 encoder 的 (batch, seq_len), 这里完全是作为一个 placeholder, 因为既然有了 encoder_outputs, input_ids 是不需要执行的。

   更重要的是 encoder_outputs 没有给定的情况下，会返回一个 (batch, 1) 形状的 bos_token_id 矩阵，这里说明，编码器第一个元素默认需要一个初始 token, generation_config.bos_token_id 默认是 101 也就是 [CLS], 而只要用户给定了 input_ids，那么不会对其进行修改。

   另外，如果当前不是 EncoderDecoderModel 而是 GPT2LMHeadModel 调用 generate, 那么在没有输入的情况下，会自动使用 CLS 作为输入, 这是合理的，因为要预测第一个元素必须要有一个启动指示 token. 但如果 inputs 不是空，那么不会做处理。

2. 接着是 _prepare_decoder_input_ids_for_generation 来对 decoder 的输入进行处理，调用接口如下：

   if self.config.is_encoder_decoder:
      input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
                      batch_size=batch_size,
                      model_input_name=model_input_name,
                      model_kwargs=model_kwargs,
                      decoder_start_token_id=generation_config.decoder_start_token_id,
                      bos_token_id=generation_config.bos_token_id,
                      device=inputs_tensor.device,
                  )
   

   首先，要注意，只有 encoder decoder 结构才会触发以上函数，如果直接是 GPT2LMHeadModel 调用 generate 并不会做这个处理。

   另外 generation_config.decoder_start_token_id 和 generation_config.bos_token_id 一般都是 CLS, 在构建模型的时候，还可以用以下语句来设置 token

   decoder.config.bos_token_id = decoder_tokenizer.cls_token_id
   

   进入该函数后，首先有一个参数优先级判断：

   if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
       decoder_input_ids = model_kwargs.pop("decoder_input_ids")
   elif "input_ids" in model_kwargs and model_input_name != "input_ids":
       decoder_input_ids = model_kwargs.pop("input_ids")
   else:
       decoder_input_ids = None
   

   这里， decoder_input_ids 是比 input_ids 高的，选用 decoder_input_ids 作为解码输入。 然后以下代码给现有的输入都加上了 CLS token

    #skip
    decoder_input_ids_start = torch.ones((batch_size, 1), dtype=torch.long, device=device) * decoder_start_token_id
   if decoder_input_ids is None:
      decoder_input_ids = decoder_input_ids_start
   elif (decoder_input_ids[:, 0] != decoder_start_token_id).all().item():
      decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
      #skip
   

   总的来说， generate 中同时输入 input_ids 和 decoder_input_ids 的话， input_ids 是给编码器的，如果不是 None, 那么不会对这个输入有任何修改，decoder_input_ids 则是给解码器的，如果没有手动添加 CLS, 那么 generate 会自动添加 CLS, 等进入到 beam_search 或 greey_search 之后，decoder 输入会自动添加 101.

   但是对于纯 decoder 模型，generate 中输入的 input_ids 就是给解码器的，如果它是 None ，那么会被 _maybe_initialize_input_ids_for_generation 修正为 CLS. 当然一般通过 tokenizer 的 encode 或 encoder_plus 对序列进行编码时会自动在开头加上 CLS 结尾加上 SEP.

   在 generate 函数中，不会对 position_ids 进行处理，它是在具体的模型里准备的，因此放到训练阶段来说明

考虑训练的场景，首先对于编码器，EncoderDecoderModel 不会对 encoder 的输入进行任何进行额外处理，就是把 input_ids 传给 encoder 的 forward.

对于解码器，由于 GPT2 采用自监督的学习方式，因此训练的时候，理论上只给定 labels_ids， 例如假设在训练中只给定了 lable_ids 为 [1,2,3], 那么模型应该可以自动计算出 input_is 为 [101, 1, 2]

在 EncoderDecoderModel 的 forward 中计算完 encoder 后有以下代码 :

if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
    decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)

    # Decode
decoder_outputs = self.decoder(
    input_ids=decoder_input_ids,
    attention_mask=decoder_attention_mask,
    encoder_hidden_states=encoder_hidden_states,
    encoder_attention_mask=attention_mask,
    inputs_embeds=decoder_inputs_embeds,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
    use_cache=use_cache,
    past_key_values=past_key_values,
    return_dict=return_dict,
    **kwargs_decoder,
)

# Compute loss independent from decoder (as some shift the logits inside them)
loss = None
if labels is not None:
    warnings.warn(DEPRECATION_WARNING, FutureWarning)
    logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
    loss_fct = CrossEntropyLoss()
    loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))

总的来看有三步：构造出 decoder_input_ids （如果没有传入的话）, decoder.forward, 计算损失。

此处的注意点是：

• 不使用 decoder_outputs 里的损失，而是重新计算一遍。

接着来看第一步：在调用 decoder.forward 前，如果只给定 labels, 那么先调用 shift_tokens_right 函数构造出解码器输入，其实现如下

def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
    """
    Shift input ids one token to the right.
    """
    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
    if decoder_start_token_id is None:
        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
    shifted_input_ids[:, 0] = decoder_start_token_id

    if pad_token_id is None:
        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
    # replace possible -100 values in labels by `pad_token_id`
    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

    return shifted_input_ids

该函数先构造一个形状同 labels 的全 0 矩阵 shifted_input_ids ，假设 labels 长度都是 n，那么 shifted_input_ids 后 n-1 个元素会被赋予 labels 前 n-1 个元素的值，而第一个值赋值为 CLS 编码（101），还是以之前的例子来看的话，得到的是 [101,1,2], 最后一句还对 shifted_input_ids 中值为 -100 的 token_id 重新赋值为 pad_token_id, label 中的 -100 是告诉损失函数不计算该 token 处的损失的。因此如果 labels 是 [1,2,3,-100,4], 同时没有给定 decoder_input_ids, 那么会得到 [101, 1,2,3,0]

最后进入到 decoder 前向传播部分，首先 decoder_input_ids 被赋值给了 input_ids, 而 GPT2LMHeadModel 的 forward 中没有对 input_ids 的处理，但在执行完 self.transformer 和 self.lm_head(hidden_states) 后会对 得到的 logits 做以下处理并计算损失：

if labels is not None:
    # move labels to correct device to enable model parallelism
    labels = labels.to(lm_logits.device)
    # Shift so that tokens < n predict n
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    # Flatten the tokens
    loss_fct = CrossEntropyLoss()
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

从这里可以看到 labels 会被截断，所以如果它是 [1,2,3,-100,4] 形式，shift_labels 就是 [2,3,-100, 4] shift_logits 对应的 id 就是 [1, 2, 3 0]. 可以看到这里无法计算到预测第一个 token 的损失。好在上文知道 EncoderDecoderModel 不会用此处计算的损失。

因此这里的注意点是：

• 如果直接基于 GPT2LMHeadModel 训练，那么输入的 labels 应该手动加上 CLS, 并且 input_ids 应该是和 labels 一样的，这样损失才是合理的。
• 如果 GPT2LMHeadModel 只是作为 EncoderDecoderModel 中的解码器，那么 labels 应该是不带 CLS 的， 如果不希望 labels 为 -100 位置对应的 decoder_input_ids 变成 0, 可以手动设置 decoder_input_ids, 其第一个元素得是 CLS, 剩余为 labels[:-1]

最后，对于 position_ids, 会在 GPT2Model.forward 中计算（也就是输入 embedding 层之前），如下：

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])