ChatGLM3 源码解析（五）_chatglm prefix decoder

PrefixEncoder

# 根据前缀 ID 获取前缀嵌入
# 前缀嵌入将连接到分头之后的 K 和 V 上
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: ChatGLMConfig):
        super().__init__()
        # 控制是否开启前缀投影，即用两层 MLP 处理前缀嵌入
        self.prefix_projection = config.prefix_projection
        if self.prefix_projection:
            # KVSize = NLayer * 2 * NGroup * HeadSize
            kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
            # 将 ID 变为嵌入的嵌入层，[PreSeqLen, KVSize]
            self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
            # 处理嵌入的 MLP
            # 映射到 HidSize, 计算 tanh，在映射到 KVSize
            self.trans = torch.nn.Sequential(
                torch.nn.Linear(kv_size, config.hidden_size),
                torch.nn.Tanh(),
                torch.nn.Linear(config.hidden_size, kv_size)
            )
        else:
            # 将 ID 变为嵌入的嵌入层
            self.embedding = torch.nn.Embedding(config.pre_seq_len,
                                                config.num_layers * config.kv_channels * config.multi_query_group_num * 2)

    def forward(self, prefix: torch.Tensor):
        # 前缀 ID 尺寸为 [BatchSize, PreSeqLen]
        # 根据前缀 ID 获取嵌入，尺寸为 [BatchSize, PreSeqLen, KVSize]
        # 如果设定了需要投影，就用两层 MLP 处理嵌入
        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


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

ChatGLMPreTrainedModel

class ChatGLMPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and
    a simple interface for downloading and loading pretrained models.
    """

    is_parallelizable = False
    supports_gradient_checkpointing = True
    config_class = ChatGLMConfig
    base_model_prefix = "transformer"
    _no_split_modules = ["GLMBlock"]

    def _init_weights(self, module: nn.Module):
        """Initialize the weights."""
        return

    # 从输入单词 ID，KVCache生成默认的（上三角）掩码矩阵
    def get_masks(self, input_ids, past_key_values, padding_mask=None):
        # 单词 ID 尺寸为 [BatchSize, SeqLen]
        batch_size, seq_length = input_ids.shape
        # 掩码矩阵初始化为全 1，形状为 [BatchSize, SeqLen, SeqLen]，每个输入序列一个
        full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
        # 保留其下三角元素，其余设为 9
        full_attention_mask.tril_()
        # CacheLen：KVCache 中序列长度
        # 如果没有提供则设为 0，如果提供了，从中获取长度
        past_length = 0
        if past_key_values:
            past_length = past_key_values[0][0].shape[0]
        # 如果提供了 KVCache，在每个掩码矩阵的上方填充 1，形状为 [BatchSize, SeqLen, CacheSeqLen]
        if past_length:
            full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
                                                        device=input_ids.device), full_attention_mask), dim=-1)

        # 如果提供了掩码数组（[BatchSize, (Cache)SeqLen]）
        # 将其变形为 [BatchSize, 1, (Cache)SeqLen]
        # 然后与掩码矩阵相乘
        # 将掩码数组为0的列设为0
        if padding_mask is not None:
            full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
        # 如果提供了掩码数组，并且没有提供 KVCache
        # 将其变形为 [BatchSize, SeqLen, 1]
        # 然后将掩码数组为 0 的行设为 1
        if not past_length and padding_mask is not None:
            full_attention_mask -= padding_mask.unsqueeze(-1) - 1
        # 小于 0.5 变成 true，大于 0.5 变成 false，相当于将其翻转，上三角不为 0
        full_attention_mask = (full_attention_mask < 0.5).bool()
        # 分头，变形为 [BatchSize, 1, SeqLen, SeqLen]
        full_attention_mask.unsqueeze_(1)
        return full_attention_mask

    # 从输入单词 ID 生成默认的（从零开始的）序列 ID
    def get_position_ids(self, input_ids, device):
        # 单词 ID 尺寸为 [BatchSize, SeqLen]
        batch_size, seq_length = input_ids.shape
        # 序列 ID 创建为 0~(SeqLen-1)的一维数组
        # 变形为 [1, SeqLen]，之后重复第一维 BatchSize 次，得到 [BatchSize, SeqLen]
        position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
        return position_ids

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, GLMTransformer):
            module.gradient_checkpointing = value

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

ChatGLMForConditionalGeneration.stream_generate()

    @torch.inference_mode()
    def stream_generate(
            self,
            input_ids,
            generation_config: Optional[GenerationConfig] = None,
            logits_processor: Optional[LogitsProcessorList] = None,
            stopping_criteria: Optional[StoppingCriteriaList] = None,
            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
            return_past_key_values=False,
            **kwargs,
    ):
        batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]

        if generation_config is None:
            generation_config = self.generation_config
        generation_config = copy.deepcopy(generation_config)
        model_kwargs = generation_config.update(**kwargs)
        model_kwargs["use_cache"] = generation_config.use_cache
        bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id

        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None

        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        if has_default_max_length and generation_config.max_new_tokens is None:
            warnings.warn(
                f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
                "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
                " recommend using `max_new_tokens` to control the maximum length of the generation.",
                UserWarning,
            )
        elif generation_config.max_new_tokens is not None:
            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
            if not has_default_max_length:
                logger.warn(
                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                    "Please refer to the documentation for more information. "
                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
                    UserWarning,
                )

        # 如果 SeqLen 大于等于配置里设定的 MaxSeqLen，发出警告
        if input_ids_seq_length >= generation_config.max_length:
            input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
            logger.warning(
                f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
                " increasing `max_new_tokens`."
            )

        # 如果没有提供 logits 处理器，初始化为空列表
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        # 没有提供停止标准，初始化为空列表
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()

        # 根据生成配置等对象获取 logits 处理器
        logits_processor = self._get_logits_processor(
            generation_config=generation_config,
            input_ids_seq_length=input_ids_seq_length,
            encoder_input_ids=input_ids,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            logits_processor=logits_processor,
        )
        # 根据生成配置等对象获取停止标准
        stopping_criteria = self._get_stopping_criteria(
            generation_config=generation_config, stopping_criteria=stopping_criteria
        )
        # 根据生成配置获取 logits 包装器
        logits_warper = self._get_logits_warper(generation_config)

        # 未完成标志，表示每个序列是否生成完毕的数组
        # 初始化为 [BatchSize] 尺寸的全 1 数组
        unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
        scores = None
        while True:
            # 根据传入参数组装成字典，请见该方法定义
            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
            # 将单词 ID 传入模型，得到（所有前缀）下一个单词的 logits
            # [BatchSize, SeqLen, VocabSize]
            outputs = self(
                **model_inputs,
                return_dict=True,
                output_attentions=False,
                output_hidden_states=False,
            )
            # 截取 SeqLen 维度的最后一维，得到整句话下一个单词的 logits
            # [BatchSize, VocabSize]
            next_token_logits = outputs.logits[:, -1, :]

            # 传入 logits 处理器和包装器，修正 logits
            next_token_scores = logits_processor(input_ids, next_token_logits)
            next_token_scores = logits_warper(input_ids, next_token_scores)

            # 计算 softmax 得到概率值
            probs = nn.functional.softmax(next_token_scores, dim=-1)
            # 如果设定了需要采样，对其进行多项式采样，样本容量为 1
            # 否则直接取最大的
            # 得到下个单词 ID，尺寸为　[BatchSize]
            if generation_config.do_sample:
                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
            else:
                next_tokens = torch.argmax(probs, dim=-1)
            # 下个单词 ID 变形为 [BatchSize, 1]，然后和输入单词 ID 拼接
            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
            # 根据当前输出更新KVCache、注意力掩码和位置ID
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
            )
            # `next_tokens` 变形为 [1, BatchSize]，再将第一维重复 NEOS 次，[NEOS, BatchSize]
            # `eos_token_id_tensor` 变形为 [NEOS, 1]，将广播第二维变成 [NEOS, BatchSize]
            # 之后二者逐元素比较是否不相等，形成一个比较结果，尺寸为 [NEOS, BatchSize]
            # 之后按照 BatchSize 维度计算乘积，得到未完成标志，[BatchSize]
            # 如果某个序列等于终止符集合里面的任意一个，那么比较结果就会出现一个 0，未完成标志将会是 0。
            unfinished_sequences = unfinished_sequences.mul(
                next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
            )
            # 如果指定了返回 KVCache
            # 产生输入ID和已生成的输出ID
            # 和 KVCache
            # 否则只产生第一个
            if return_past_key_values:
                yield input_ids, outputs.past_key_values
            else:
                yield input_ids
            # 如果未完成标志全为零（表示序列都已生成完毕），或者达到了停止标准，就停止生成
            if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
                break

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