"""权重初始化：LLM 模型权重初始化策略""" import torch import torch.nn as nn import math def init_weights(module, std=0.02): """ 初始化模型权重（适用于 LLM，参考 GPT/LLaMA） 参数: module: PyTorch 模块 std: 正态分布的标准差（默认: 0.02） 初始化策略: - nn.Embedding: 正态分布 N(0, std) - nn.Linear: 正态分布 N(0, std)，偏置初始化为 0 - RMSNorm: 可学习参数（scale）初始化为 1.0（RMSNorm 会自动处理） """ if isinstance(module, nn.Embedding): # 词嵌入层：正态分布初始化 nn.init.normal_(module.weight, mean=0.0, std=std) elif isinstance(module, nn.Linear): # 线性层：权重正态分布初始化，偏置初始化为 0 nn.init.normal_(module.weight, mean=0.0, std=std) if module.bias is not None: nn.init.zeros_(module.bias) def init_weights_with_scaling(module, hidden_size=None, std=0.02): """ 初始化模型权重（带缩放，适用于输出层） 参数: module: PyTorch 模块 hidden_size: 隐藏层维度（用于输出层缩放） std: 基础标准差（默认: 0.02） 初始化策略: - nn.Embedding: 正态分布 N(0, std) - nn.Linear: - 输出层（如果 hidden_size 提供）: N(0, std / sqrt(hidden_size)) - 其他层: N(0, std) - 偏置: 初始化为 0 """ if isinstance(module, nn.Embedding): # 词嵌入层：正态分布初始化 nn.init.normal_(module.weight, mean=0.0, std=std) elif isinstance(module, nn.Linear): # 线性层 if hidden_size is not None: # 输出层：使用缩放的标准差 output_std = std / math.sqrt(hidden_size) nn.init.normal_(module.weight, mean=0.0, std=output_std) else: # 普通线性层：标准正态分布初始化 nn.init.normal_(module.weight, mean=0.0, std=std) # 偏置初始化为 0 if module.bias is not None: nn.init.zeros_(module.bias) def apply_llm_init(model, std=0.02, init_output_layer=True): """ 对整个模型应用 LLM 权重初始化 参数: model: Transformer 模型 std: 正态分布的标准差（默认: 0.02） init_output_layer: 是否对输出层使用特殊初始化（默认: True） 返回: 初始化后的模型 """ # 获取 hidden_size（用于输出层初始化） hidden_size = None if hasattr(model, "config"): hidden_size = model.config.get("hidden_size") elif hasattr(model, "hidden_size"): hidden_size = model.hidden_size # 遍历所有模块并初始化 for module in model.modules(): if isinstance(module, (nn.Embedding, nn.Linear)): if init_output_layer and isinstance(module, nn.Linear): # 检查是否是输出层（lm_head） # 如果 tie_word_embeddings=True，lm_head 可能为 None，使用 embedding 的权重 if hasattr(model, "lm_head") and module is model.lm_head: # 输出层：使用缩放初始化 init_weights_with_scaling(module, hidden_size=hidden_size, std=std) else: # 普通线性层 init_weights(module, std=std) else: # 标准初始化 init_weights(module, std=std) return model if __name__ == "__main__": import sys import io # 设置输出编码为 UTF-8（Windows 兼容） if sys.platform == "win32": sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8") print("=" * 60) print("权重初始化测试") print("=" * 60) # 测试参数 vocab_size = 100 hidden_size = 320 intermediate_size = 960 std = 0.02 print(f"\n测试参数:") print(f" vocab_size: {vocab_size}") print(f" hidden_size: {hidden_size}") print(f" intermediate_size: {intermediate_size}") print(f" std: {std}") # 测试 Embedding 初始化 print("\n1. 测试 Embedding 初始化") embedding = nn.Embedding(vocab_size, hidden_size) init_weights(embedding, std=std) weight_mean = embedding.weight.mean().item() weight_std = embedding.weight.std().item() print(f" Embedding 权重均值: {weight_mean:.6f} (应该接近 0)") print(f" Embedding 权重标准差: {weight_std:.6f} (应该接近 {std})") # 测试 Linear 初始化 print("\n2. 测试 Linear 初始化") linear = nn.Linear(hidden_size, intermediate_size) init_weights(linear, std=std) weight_mean = linear.weight.mean().item() weight_std = linear.weight.std().item() bias_mean = linear.bias.mean().item() if linear.bias is not None else 0.0 print(f" Linear 权重均值: {weight_mean:.6f} (应该接近 0)") print(f" Linear 权重标准差: {weight_std:.6f} (应该接近 {std})") print(f" Linear 偏置均值: {bias_mean:.6f} (应该为 0)") # 测试输出层初始化（带缩放） print("\n3. 测试输出层初始化（带缩放）") output_layer = nn.Linear(hidden_size, vocab_size, bias=False) init_weights_with_scaling(output_layer, hidden_size=hidden_size, std=std) weight_mean = output_layer.weight.mean().item() weight_std = output_layer.weight.std().item() expected_std = std / math.sqrt(hidden_size) print(f" 输出层权重均值: {weight_mean:.6f} (应该接近 0)") print(f" 输出层权重标准差: {weight_std:.6f}") print(f" 期望标准差: {expected_std:.6f}") # 测试完整模型初始化 print("\n4. 测试完整模型初始化") from llm.model.transformer import Transformer config = { "vocab_size": vocab_size, "hidden_size": hidden_size, "num_hidden_layers": 2, # 使用较小的层数用于测试 "num_attention_heads": 10, "num_key_value_heads": 2, "intermediate_size": intermediate_size, "rms_norm_eps": 1e-5, "max_position_embeddings": 1024, "rope_theta": 10000.0, "sliding_window": 256, "sliding_window_overlap": True, "tie_word_embeddings": True, } model = Transformer(config) # 记录初始化前的权重统计 print(" 初始化前的权重统计:") for name, param in model.named_parameters(): if "embedding" in name or "weight" in name: print(f" {name}: mean={param.data.mean().item():.6f}, std={param.data.std().item():.6f}") # 应用初始化 apply_llm_init(model, std=std, init_output_layer=True) # 记录初始化后的权重统计 print("\n 初始化后的权重统计:") for name, param in model.named_parameters(): if "embedding" in name or ("weight" in name and "norm" not in name): print(f" {name}: mean={param.data.mean().item():.6f}, std={param.data.std().item():.6f}") # 验证初始化效果 print("\n5. 验证初始化效果") embedding_weight = model.embedding.embedding.weight print(f" Embedding 权重均值: {embedding_weight.mean().item():.6f}") print(f" Embedding 权重标准差: {embedding_weight.std().item():.6f}") # 检查第一个 Transformer Block 的线性层 first_block = model.layers[0] attn_q_proj = first_block.attn.q_proj print(f" Attention Q 投影权重均值: {attn_q_proj.weight.mean().item():.6f}") print(f" Attention Q 投影权重标准差: {attn_q_proj.weight.std().item():.6f}") ffn_gate_proj = first_block.ffn.gate_proj print(f" FFN Gate 投影权重均值: {ffn_gate_proj.weight.mean().item():.6f}") print(f" FFN Gate 投影权重标准差: {ffn_gate_proj.weight.std().item():.6f}") print("\n" + "=" * 60) print("所有测试完成！") print("=" * 60)