1
(transformer): GPT2Model(
2
  (wte): Embedding(100, 64)
3
  (wpe): Embedding(64, 64)
4
)

1
(drop): Dropout(p=0.1)
2
(h): ModuleList(
3
  (0-1): 2 x GPT2Block(...)
4
)
5
(ln_f): LayerNorm((64,), eps=1e-05)

1
# GPT2Model = TokenEmbedding + PositionEmbedding (+ TokenTypeEmbedding)
2
#           → Dropout
3
#           → num_layers × GPT2Block
4
#           → Final LayerNorm
5
#           → 输出 last_hidden_state（以及可选的 past_key_values / attentions / hidden_states）
6
class GPT2Model(GPT2PreTrainedModel):
7
    _supports_param_buffer_assignment = False
8

9
    def __init__(self, config):
10
        super().__init__(config)
11
        self.embed_dim = config.hidden_size # 隐藏层维度
12
        self.wte = nn.Embedding(config.vocab_size, self.embed_dim) # 词嵌入层 Word Token Embedding
13
        self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) # 位置嵌入层 Word Position Embedding
14

15
        self.drop = nn.Dropout(config.embd_pdrop) # dropout 层
16
        self.h = nn.ModuleList([GPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)]) # Transformer 堆叠
17
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) # 最终输出前的 LayerNorm
18
        # Initialize weights and apply final processing
19
        self.post_init()
20

21
    def forward(
22
        self,
23
        input_ids: Optional[torch.LongTensor] = None,
24
        past_key_values: Optional[Cache] = None,
25
        cache_position: Optional[torch.LongTensor] = None,
26
        attention_mask: Optional[torch.FloatTensor] = None,
27
        token_type_ids: Optional[torch.LongTensor] = None,
28
        position_ids: Optional[torch.LongTensor] = None,
29
        head_mask: Optional[torch.FloatTensor] = None,
30
        inputs_embeds: Optional[torch.FloatTensor] = None,
31
        encoder_hidden_states: Optional[torch.Tensor] = None,
32
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
33
        use_cache: Optional[bool] = None,
34
        output_attentions: Optional[bool] = None,
35
        output_hidden_states: Optional[bool] = None,
36
        return_dict: Optional[bool] = None,
37
        **kwargs,
38
    ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:
39
        r"""
40
        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
41
            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
42
            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input
43
            sequence tokens in the vocabulary.
44

45
            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
46
            `input_ids`.
47

48
            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
49
            [`PreTrainedTokenizer.__call__`] for details.
50

51
            [What are input IDs?](../glossary#input-ids)
52
        """
53
    # 选项初始化
54
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
55
        output_hidden_states = (
56
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
57
        )
58
        use_cache = use_cache if use_cache is not None else self.config.use_cache
59
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
60
    # 检查输入
61
        if input_ids is not None and inputs_embeds is not None:
62
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
63
        elif input_ids is not None:
64
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
65
      # 把多维输入压平，转换成标准的批次输入格式
66
      # [batch, segments, seq_len]
67
      # [batch, n_sentences, seq_len]
68
      # [batch, n_variants, seq_len]
69
      # 统一变成 [N, seq_len]，Embedding 期望的输入
70
            input_shape = input_ids.size()
71
            input_ids = input_ids.view(-1, input_shape[-1])
72
            batch_size = input_ids.shape[0]
73
        elif inputs_embeds is not None:
74
            input_shape = inputs_embeds.size()[:-1]
75
            batch_size = inputs_embeds.shape[0]
76
        else:
77
            raise ValueError("You have to specify either input_ids or inputs_embeds")
78

79
        device = input_ids.device if input_ids is not None else inputs_embeds.device
80

81
        if token_type_ids is not None:
82
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
83
    # 梯度检查点与缓存冲突
84
        if self.gradient_checkpointing and self.training:
85
            if use_cache:
86
                logger.warning_once(
87
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
88
                )
89
                use_cache = False
90

91
        # based on pattern from src/transformers/models/whisper/modeling_whisper.py::WhisperDecoder
92
    # 初始化 kv cache
93
        if use_cache:
94
            if past_key_values is None:
95
                past_key_values = DynamicCache(config=self.config)
96
            if self.config.add_cross_attention and not isinstance(past_key_values, EncoderDecoderCache):
97
                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config))
98
    # 词嵌入
99
        if inputs_embeds is None:
100
            inputs_embeds = self.wte(input_ids) # (N, seq_len) -> (N, seq_len, hidden_size)
101
    # 生成绝对位置编码
102
        if cache_position is None:
103
      # 先前token数
104
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
105
            cache_position = torch.arange(
106
                past_seen_tokens,
107
        past_seen_tokens + inputs_embeds.shape[1], # inputs_embeds.shape[1] 为序列长度
108
        device=inputs_embeds.device
109
            ) # (seq_len), eg: tensor([ 0,  1,  2,  3,  4,  5])
110
        if position_ids is None:
111
      # 为什么要 unsqueeze(0) ？
112
      # 因为模型需要：(batch_size, seq_len)
113
      # 而 cache_position 是 (seq_len)
114
            position_ids = cache_position.unsqueeze(0) # 增加维度 (1, seq_len)
115
    # 位置嵌入
116
        position_embeds = self.wpe(position_ids) # (1, seq_len, hidden_size)
117
    # (N, seq_len, hidden_size) + (1, seq_len, hidden_size)，广播
118
        hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)
119

120
        # Attention mask.
121
        # ._update_causal_mask() and ._prepare_4d_causal_attention_mask_with_cache_position() copied from LlamaModel
122
        if attention_mask is not None and attention_mask.ndim < 4:
123
            attention_mask = attention_mask.view(batch_size, -1)
124
    # 构造最终的因果注意力 mask
125
        causal_mask = create_causal_mask(
126
            config=self.config,
127
            input_embeds=inputs_embeds,
128
            attention_mask=attention_mask,
129
            cache_position=cache_position,
130
            past_key_values=past_key_values,
131
            position_ids=position_ids,
132
        )
133
    # 交叉注意力处理
134
        # If a 2D or 3D attention mask is provided for the cross-attention
135
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
136
        _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
137
        if self.config.add_cross_attention and encoder_hidden_states is not None:
138
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
139
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
140
            if encoder_attention_mask is None:
141
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
142
            if _use_sdpa:
143
                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
144
                    mask=encoder_attention_mask, dtype=inputs_embeds.dtype, tgt_len=input_shape[-1]
145
                )
146
            elif self._attn_implementation != "flash_attention_2":
147
                encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
148
        else:
149
            encoder_attention_mask = None
150
    # 把传入的 head 层级屏蔽参数整理成统一形状，用于在 block 内屏蔽某些注意力头。
151
        # Prepare head mask if needed
152
        # 1.0 in head_mask indicate we keep the head
153
        # attention_probs has shape bsz x n_heads x N x N
154
        # head_mask has shape n_layer x batch x n_heads x N x N
155
        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
156
    # token 类型嵌入
157
        if token_type_ids is not None:
158
            token_type_embeds = self.wte(token_type_ids) # (batch_size, seq_len) -> (batch_size, seq_len, hidden_size)
159
            hidden_states = hidden_states + token_type_embeds
160
    # dropout
161
        hidden_states = self.drop(hidden_states)
162
    # 输出形状记录，主要用于最后 view 把 batch/seq 形状恢复
163
        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),) # (-1, sample_shape, hidden_size)
164
    # 如果需要输出 attentions，则初始化为 ()，否则 None。其他同理
165
        all_self_attentions = () if output_attentions else None
166
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
167
        all_hidden_states = () if output_hidden_states else None
168
    # 模型主体：依次通过每一层 GPT2Block
169
        for i, block in enumerate(self.h):
170
            # Model parallel
171
            if self.model_parallel:
172
                torch.cuda.set_device(hidden_states.device)
173
                if isinstance(head_mask, torch.Tensor):
174
                    head_mask = head_mask.to(hidden_states.device)
175
            if output_hidden_states:
176
                all_hidden_states = all_hidden_states + (hidden_states,) # 收集中间隐状态
177
      # 调用单个 block 的 forward
178
            outputs = block(
179
                hidden_states, # (N, seq_len, hidden_size)
180
                past_key_values if not (self.gradient_checkpointing and self.training) else None,
181
                cache_position,
182
                causal_mask,
183
                head_mask[i],
184
                encoder_hidden_states,  # as a positional argument for gradient checkpointing
185
                encoder_attention_mask=encoder_attention_mask,
186
                use_cache=use_cache,
187
                output_attentions=output_attentions,
188
                **kwargs,
189
            )
190
            hidden_states = outputs[0] # 取隐藏层状态
191
      # 收集中间隐状态
192
            if output_attentions:
193
                all_self_attentions = all_self_attentions + (outputs[1],)
194
                if self.config.add_cross_attention:
195
                    all_cross_attentions = all_cross_attentions + (outputs[2],)
196
            # 若启用模型并行，根据 self.device_map 判断当前层是否是该设备上的最后一层，如果是，就把 hidden_states 迁移到下一个设备上。
197
            if self.model_parallel:
198
                for k, v in self.device_map.items():
199
                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
200
                        hidden_states = hidden_states.to("cuda:" + str(k + 1))
201
    # 最终输出前的 LayerNorm
202
        hidden_states = self.ln_f(hidden_states)
203
    # 恢复输出形状 (-1, sample_shape, hidden_size)
204
        hidden_states = hidden_states.view(output_shape)
205
        # 收集最后一层隐状态
206
        if output_hidden_states:
207
            all_hidden_states = all_hidden_states + (hidden_states,)
208
    # 缓存输出
209
        past_key_values = past_key_values if use_cache else None
210
        if not return_dict:
211
      # [hidden_states, past_key_values, all_hidden_states, all_self_attentions, all_cross_attentions] 中非空的按顺序返回
212
            return tuple(
213
                v
214
                for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions, all_cross_attentions]
215
                if v is not None
216
            )
217

218
        return BaseModelOutputWithPastAndCrossAttentions(
219
            last_hidden_state=hidden_states,
220
            past_key_values=past_key_values,
221
            hidden_states=all_hidden_states,
222
            attentions=all_self_attentions,
223
            cross_attentions=all_cross_attentions,
224
        )

1
GPT2Block(
2
  (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
3
  (attn): GPT2Attention(
4
    (c_attn): Conv1D(nf=192, nx=64)
5
    (c_proj): Conv1D(nf=64, nx=64)
6
    (attn_dropout): Dropout(p=0.1)
7
    (resid_dropout): Dropout(p=0.1)
8
  )
9
  (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
10
  (mlp): GPT2MLP(
11
    (c_fc): Conv1D(nf=128, nx=64)
12
    (c_proj): Conv1D(nf=64, nx=128)
13
    (act): GELUActivation()
14
    (dropout): Dropout(p=0.1)
15
  )
16
)

1
(lm_head): Linear(in_features=64, out_features=100, bias=False)

1
# GPT2LMHeadModel = GPT2Model → 取 hidden_states → Linear(lm_head) → logits（+ 可选 loss）
2
class GPT2LMHeadModel(GPT2PreTrainedModel, GenerationMixin):
3
    _tied_weights_keys = ["lm_head.weight"]
4

5
    def __init__(self, config):
6
        super().__init__(config)
7
        self.transformer = GPT2Model(config)  # GPT2 模型部分
8
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) # 任务头
9

10
        # Initialize weights and apply final processing
11
        self.post_init()
12

13
    def forward(self,...,**kwargs,) -> Union[tuple, CausalLMOutputWithCrossAttentions]:
14
        r"""
15
        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
16
            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
17
            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input
18
            sequence tokens in the vocabulary.
19

20
            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
21
            `input_ids`.
22

23
            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
24
            [`PreTrainedTokenizer.__call__`] for details.
25

26
            [What are input IDs?](../glossary#input-ids)
27
        labels (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
28
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
29
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
30
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
31
        """
32
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
33

34
        transformer_outputs = self.transformer( # GPT‑2 本体
35
            input_ids,  # (batch_size, input_ids_length)
36
            past_key_values=past_key_values,
37
            attention_mask=attention_mask,
38
            cache_position=cache_position,
39
            token_type_ids=token_type_ids,
40
            position_ids=position_ids,
41
            head_mask=head_mask,
42
            inputs_embeds=inputs_embeds,
43
            encoder_hidden_states=encoder_hidden_states,
44
            encoder_attention_mask=encoder_attention_mask,
45
            use_cache=use_cache,
46
            output_attentions=output_attentions,
47
            output_hidden_states=output_hidden_states,
48
            return_dict=return_dict,
49
        )
50
        hidden_states = transformer_outputs[0]  # (batch_size, seq_len, hidden_size)，获取隐藏层，使用[0]取，因为返回的可能不是dict
51
    # 对隐藏层切片送入 lm_head
52
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
53
        logits = self.lm_head(hidden_states[:, slice_indices, :]) # (batch_size, kept_seq_len, vocab_size)
54

55
        loss = None  # 表示默认情况下（没有标签）是不计算损失的
56
        if labels is not None:
57
            # 把 token 维度拉平
58
            loss = self.loss_function(
59
                logits,
60
                labels,
61
                vocab_size=self.config.vocab_size,
62
                **kwargs,
63
            )
64

65
        if not return_dict:
66
            output = (logits,) + transformer_outputs[1:]
67
            return ((loss,) + output) if loss is not None else output
68

69
        return CausalLMOutputWithCrossAttentions(
70
            loss=loss,
71
            logits=logits,
72
            past_key_values=transformer_outputs.past_key_values,
73
            hidden_states=transformer_outputs.hidden_states,
74
            attentions=transformer_outputs.attentions,
75
            cross_attentions=transformer_outputs.cross_attentions,
76
        )