![img](https://i-blog.csdnimg.cn/blog_migrate/4757ad5c98d4344547227fc52684bac1.png)

是Qwen2模型中各组件和函数的详细作用说明，按模块分类整理：

### **核心工具函数**

|           函数/常量            |                           作用                           |
| :----------------------------: | :------------------------------------------------------: |
|           `__all__`            | 定义模块的公开接口，控制`from module import *`时的可见性 |
|   `get_triangle_upper_mask`    |       生成上三角因果注意力掩码（防止未来信息泄露）       |
|       `assign_kv_heads`        |           分配Key/Value头的索引（用于GQA/MQA）           |
|       `parallel_matmul`        |               并行矩阵乘法（支持张量并行）               |
| `scaled_dot_product_attention` |                实现缩放点积注意力核心计算                |
|         `masked_fill`          |       按掩码填充张量（如将padding位置设为负无穷）        |
|        `is_casual_mask`        |                 判断是否为因果注意力掩码                 |
|      `_make_causal_mask`       |            创建因果注意力掩码（考虑padding）             |
|       `_expand_2d_mask`        |            将2D掩码扩展为4D（适配多头注意力）            |
|          `repeat_kv`           |              重复Key/Value头（用于GQA/MQA）              |

### **归一化层**

|    类/函数     |               作用               |
| :------------: | :------------------------------: |
| `Qwen2RMSNorm` | **RMS归一化层**（替代LayerNorm） |

### **位置编码**

|         类/函数          |                作用                |
| :----------------------: | :--------------------------------: |
|  `Qwen2RotaryEmbedding`  |       **旋转位置编码(RoPE)**       |
|     - `rotate_half`      | 旋转向量的后半部分（RoPE核心操作） |
| - `apply_rotary_pos_emb` |   将旋转位置编码应用到注意力分数   |

### **前馈网络**

|  类/函数   |             作用              |
| :--------: | :---------------------------: |
| `Qwen2MLP` | **门控线性单元(GLU)前馈网络** |

### **注意力机制**

|       类/函数       |                        作用                        |
| :-----------------: | :------------------------------------------------: |
|  `Qwen2Attention`   |                 **多头注意力机制**                 |
|    - `__init__`     |          初始化Q/K/V投影层、输出层和RoPE           |
|     - `forward`     |      处理输入序列，计算注意力分数并聚合值向量      |
| `Qwen2DecoderLayer` |               **Transformer解码层**                |
|    - `__init__`     |              组合自注意力层和前馈网络              |
|     - `forward`     | 执行：`LN -> Attention -> Add -> LN -> MLP -> Add` |

### **预训练基础**

|        类/函数         |                作用                |
| :--------------------: | :--------------------------------: |
| `Qwen2PretrainedModel` |         **预训练模型基类**         |
|    - `config_class`    |    关联的配置类（Qwen2Config）     |
| - `_get_name_mappings` | 定义参数名称映射（用于加载检查点） |
|   - `_init_weights`    |           参数初始化策略           |
|  - `_get_model_flops`  |           计算模型FLOPs            |

### **主干模型**

|               类/函数               |           作用            |
| :---------------------------------: | :-----------------------: |
|            `Qwen2Model`             |     **模型主干架构**      |
| - `_prepare_decoder_attention_mask` |      生成解码器掩码       |
|             - `forward`             | 执行完整的Transformer堆栈 |
|         `Qwen2ForCausalLM`          |     **因果语言模型**      |
|  - `prepare_inputs_for_generation`  |   处理生成时的输入格式    |
|             - `forward`             |     计算语言建模损失      |

### **任务特定头部**

|                类                |           作用           |
| :------------------------------: | :----------------------: |
|          `Qwen2LMHead`           | 语言模型头部（词表投影） |
| `Qwen2ForSequenceClassification` |     序列分类任务适配     |
|  `Qwen2ForTokenClassification`   |     标记分类任务适配     |
|     `Qwen2SentenceEmbedding`     |       句子向量提取       |

### **训练相关**

|           类/函数           |            作用            |
| :-------------------------: | :------------------------: |
| `Qwen2PretrainingCriterion` |       预训练损失计算       |
|  `recompute_training_full`  |       激活重计算策略       |
|   `create_custom_forward`   | 为梯度检查点创建自定义前向 |