
是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 |
为梯度检查点创建自定义前向 |