以下是针对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的梯度 |
wehub-resource-sync