直接赋值 lm_head.weight = embedding_layer.weight 合法有效,因二者 shape 均为 [V, d] 且 lm_head 无 bias,PyTorch 中 nn.Linear 权重存储为 [out_features, in_features],与 nn.Embedding 的 [vocab_size, embedding_dim] 天然对齐,赋值后共享内存与梯度。

直接赋值 lm_head.weight = embedding_layer.weight 就行,前提是二者 shape 完全一致([V, d]),且 lm_head 不带 bias。
为什么 lm_head.weight = embedding_layer.weight 是合法且有效的?
PyTorch 的 nn.Embedding 和 nn.Linear 在权重 shape 上天然对齐:
-
embedding_layer.weight形状是[vocab_size, embedding_dim](即[V, d]) -
lm_head = nn.Linear(embedding_dim, vocab_size, bias=False)的weight形状也是[vocab_size, embedding_dim](PyTorch 中Linear存储为[out_features, in_features]) - 所以赋值后,两者指向同一块内存;反向传播时梯度会自动累加到同一个
.grad - 调用
lm_head(hidden_states)实际执行的是hidden_states @ embedding_layer.weight.T,语义上就是“计算隐藏状态与各 token embedding 的点积”
必须避开的三个典型错误
实操中容易因细节不一致导致共享失效或报错:
- 给
lm_head加了bias=True:权重 shape 变成[V, d]+[V],无法和embedding_layer.weight对齐;务必设bias=False - 初始化顺序搞反:先
self.lm_head = ...再self.lm_head.weight = self.wte.weight;不能反过来(wte还没创建就去引用) - 在
forward里动态赋值:比如写成lm_head.weight.data = embedding_layer.weight.data—— 这只是拷贝数值,不共享梯度;必须在__init__中做对象级引用
transformers 库加载时怎么维持权重共享?
HF 的 transformers 模型(如 GPT2LMHeadModel)默认启用权重共享,但加载逻辑有隐含约定:
立即学习“Python免费学习笔记(深入)”;
- 保存时只存一份 embedding 权重(
model.state_dict()["transformer.wte.weight"]),不单独存lm_head.weight - 加载时跳过
lm_head.weight的加载步骤,然后在post_init或_tie_weights中执行self.lm_head.weight = self.transformer.wte.weight - 如果你手动修改 state dict,漏掉这步绑定,模型会退化为两套独立权重——参数量翻倍、收敛变慢、logits 偏移
真正关键的不是“怎么写那行赋值”,而是确保从初始化、训练到保存加载,整个生命周期里 embedding_layer.weight 和 lm_head.weight 始终是同一个 Python 对象。任何中间环节的深拷贝、.data 赋值、或分步加载都可能悄悄破坏这个引用关系。


















