嵌入层优化：从查找到计算
嵌入层：被忽视的显存大户 在LLM中，嵌入层（Embedding Layer）往往是被忽视的显存消耗者。以Llama-3为例，词表大小128K、隐藏维度4096、FP16精度下，嵌入层参数量为 128K × 4096 × 2 = 1GB。对于更大的词表（如Qwen的152K），嵌入层可达1.5GB以上。 在70B模型中，嵌入层占总参数量的约2%，但在推理时它还需要一个等大的输出投影层（LM Head），两者合计占4%——这还不算小。 标准嵌入层 查找表 标准嵌入层是一个简单的查找表——词表中的每个token对应一个可学习的向量： class StandardEmbedding(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.weight = nn.Parameter(torch.randn(vocab_size, d_model)) def forward(self, input_ids): # 简单的索引查找 return self.weight[input_ids] # [batch, seq_len, d_model] 权重共享 许多LLM共享嵌入层和输出层的权重（Weight Tying）： class SharedEmbeddingModel(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) # 输出层直接使用嵌入矩阵的转置 # 不需要额外的参数 def forward(self, input_ids): x = self.embedding(input_ids) x = self.transformer(x) # 输出logits = x @ embedding.weight.T logits = F.linear(x, self.embedding.weight) return logits 权重共享的好处： 减少参数量（节省1GB+显存） 输入和输出在同一个语义空间中 训练时梯度信号更强 但也有研究指出，权重共享可能导致冲突——输入嵌入需要捕捉"词的语义"，输出层需要捕捉"词的预测概率分布"，两者的最优表示可能不同。 嵌入层压缩 矩阵分解 将大的嵌入矩阵分解为两个小矩阵的乘积： class FactorizedEmbedding(nn.Module): def __init__(self, vocab_size, d_model, d_factor): """ d_factor: 分解后的中间维度 """ super().__init__() # V × D → V × d_factor × d_factor × D self.factor1 = nn.Parameter(torch.randn(vocab_size, d_factor)) self.factor2 = nn.Parameter(torch.randn(d_factor, d_model)) def forward(self, input_ids): # 查找低维因子，然后投影到高维 low_dim = self.factor1[input_ids] # [batch, seq, d_factor] return low_dim @ self.factor2 # [batch, seq, d_model] 参数量从 V×D 降低到 V×d + d×D。当 D=4096、d=256、V=128K 时，参数量减少约16倍。 ...
