60 lines
3.6 KiB
Markdown
60 lines
3.6 KiB
Markdown
以下是针对Llama模型架构中各组件的功能解析,按模块分类说明其核心作用:
|
||
|
||
------
|
||
|
||
### **核心工具函数**
|
||
|
||
| 函数/常量 | 作用 | 与Qwen2的差异 |
|
||
| :----------------------------: | :-----------------------------------: | :--------------------: |
|
||
| `swiglu` | 实现SwiGLU激活函数:`x * silu(gate)` | Qwen2使用标准GLU |
|
||
| `rms_norm_fused` | 启用融合的RMSNorm计算(CUDA优化) | 实现相同但配置参数不同 |
|
||
| `__all__` | 定义模块的公开接口 | - |
|
||
| `_get_interleave` | 生成交错注意力头索引(用于长序列) | Llama特有 |
|
||
| `build_alibi_tensor` | 构建ALiBi位置偏置张量(相对位置编码) | Qwen2未使用 |
|
||
| `get_triangle_upper_mask` | 生成因果上三角掩码 | 实现逻辑相同 |
|
||
| `assign_kv_heads` | 分配KV头的索引(支持GQA/MQA) | - |
|
||
| `parallel_matmul` | 并行矩阵乘法(张量并行) | - |
|
||
| `scaled_dot_product_attention` | 核心注意力计算 | Llama支持更多掩码类型 |
|
||
| `_make_causal_mask` | 动态生成因果掩码(考虑padding) | Qwen2更简化 |
|
||
|
||
### **归一化层**
|
||
|
||
| 类/函数 | 作用 | 差异点 |
|
||
| :------------: | :-------------------: | :-------------: |
|
||
| `LlamaRMSNorm` | 带融合优化的RMS归一化 | 与Qwen2实现相同 |
|
||
|
||
### **位置编码(核心差异)**
|
||
|
||
| 类 | 作用 | 特性 |
|
||
| :-------------------------------------: | :------------------------: | :---------------: |
|
||
| `LlamaRotaryEmbedding` | 基础RoPE实现 | - |
|
||
| `LlamaLinearScalingRotaryEmbedding` | 线性缩放RoPE(扩展上下文) | Qwen2无此变体 |
|
||
| `LlamaNTKScalingRotaryEmbedding` | NTK-aware缩放RoPE | 动态调整高频/低频 |
|
||
| `LlamaDynamicNTKScalingRotaryEmbedding` | 动态NTK缩放(训练自适应) | Llama特有 |
|
||
| `Llama3RotaryEmbedding` | Llama3专用RoPE | 改进的旋转策略 |
|
||
|
||
### **前馈网络**
|
||
|
||
| 类 | 作用 | 差异 |
|
||
| :--------: | :-----------------: | :------------: |
|
||
| `LlamaMLP` | 使用SwiGLU的门控FFN | Qwen2用普通GLU |
|
||
|
||
### **注意力机制**
|
||
|
||
| 类 | 核心改进 | 说明 |
|
||
| :-----------------: | :----------------------------------------: | :------------: |
|
||
| `LlamaAttention` | - 多版本RoPE支持 - ALiBi融合 - 动态NTK缩放 | 比Qwen2更复杂 |
|
||
| `LlamaDecoderLayer` | 深度优化层实现 | 支持梯度检查点 |
|
||
|
||
### **预训练基础**
|
||
|
||
| 类 | 关键功能 | 扩展性 |
|
||
| :--------------------: | :--------------------------: | :-----------: |
|
||
| `LlamaPretrainedModel` | - 多设备加载 - FLOPs计算工具 | 比Qwen2更完善 |
|
||
|
||
### **任务模块**
|
||
|
||
| 类 | 用途 | 特色 |
|
||
| :----------------: | :------------: | :---------------: |
|
||
| `LlamaForCausalLM` | 语言建模 | 支持静态图导出 |
|
||
| `ConcatMaskedLoss` | 多任务损失合并 | 处理padding的梯度 | |