PaddleNLP tie_weights 权重绑定能力设计与实现全解析（RFC No.103）
【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址: https://gitcode.com/gh_mirrors/pa/PaddleNLP
导读
权重绑定（Weight Tying）是预训练语言模型中最常用的参数共享技巧之一，它将输入层 Embedding 与输出层 Embedding 的权重绑定为同一份参数，既减少了网络参数量，又让 Embedding 层在训练中得到更充分的梯度更新。本文以 PaddleNLP 仓库中的 RFC 文档 20230304_api_design_for_tie_weight_task_103.md 为骨架，完整讲解该能力的背景动机、业内方案调研、API 设计思路、验收测试方案，并结合当前仓库源码（如 model_utils.py、configuration_utils.py 及 test_modeling_common.py）给出最终落地的实现细节与验证方法，帮助读者理解"为什么需要 tie_weights、它在 PaddleNLP 中如何设计与落地"。
一、背景：为什么预训练语言模型需要 tie_weights
权重绑定（tie weights）一般指将输入层 Embedding 与输出层 Embedding 的权重进行共享，其收益体现在两个方面：
- 减少网络参数量：以词汇表大小
vocab_size、隐藏维度hidden_size为例，不绑定时输入 Embedding 与输出投影层各占vocab_size * hidden_size个参数；绑定后输出层直接复用输入 Embedding 的权重，可省去一份规模巨大的参数矩阵（在词表动辄数万、数十万的 LLM 场景下，这一节省非常可观）。 - 让 Embedding 层参数训练更充分：共享权重后，输出侧的反向梯度同样会回传到 Embedding 权重上，使低频 token 的 Embedding 也能获得更多训练信号。
该技巧在《Attention Is All You Need》第 3.4 节中被明确提出（Transformer 将 encoder 输入 Embedding、decoder 输入 Embedding 与输出线性层权重共享），其有效性又在《Using the Output Embedding to Improve Language Models》中得到进一步验证。因此，预训练语言模型需要提供一个输入层 Embedding 与输出层 Embedding 权重共享的基础能力，方便使用者直接调用。
RFC 中对应的任务为No.103：新增 tie_weights 能力，目标是给预训练语言模型增加一个基础函数，实现输入层 Embedding 与输出层 Embedding 的权重共享绑定，并对齐 HuggingFace Transformers 中的PreTrainedModel.tie_weights功能。
二、飞桨现状：三种零散实现方式的调研
RFC 首先对飞桨框架当时的现状做了调研，结论是：PaddlePaddle 框架层面并没有对 tie weights 的统一实现，调用者需要自己写代码完成绑定。PaddleNLP 中已经存在三种典型的零散实现方式：
方式一：模型内定义tie_weights函数，通过赋值绑定
这种方式在modeling.py中定义tie_weights函数，模型同时实现get_input_embeddings()与get_output_embeddings()来获取输入、输出 Embedding 层权重，然后通过赋值完成绑定。RFC 引用的 Transformer-XL 示例（mem_transformer.py）展示了典型的按层赋值逻辑：
if tie_weight: for i in range(len(self.crit.out_layers_weight)): self.crit.out_layers_weight[i] = self.word_emb.emb_layers[i].weight if tie_projs: for i, tie_proj in enumerate(tie_projs): if tie_proj and div_val == 1 and d_model != d_embed: self.crit.out_projs[i] = self.word_emb.emb_projs[0] elif tie_proj and div_val != 1: self.crit.out_projs[i] = self.word_emb.emb_projs[i]而 RFC 引用的 Reformer 实现（paddlenlp/transformers/reformer/modeling.py）则更接近 HF 的结构，先判断配置开关再调用内部绑定逻辑：
def tie_weights(self): """ Tie the weights between the input embeddings and the output embeddings. """ tie_word_embeddings = ( self.tie_word_embeddings if hasattr(self, "tie_word_embeddings") else self.config.get("tie_word_embeddings", False) ) if hasattr(self, "get_output_embeddings") and hasattr(self, "get_input_embeddings") and tie_word_embeddings: output_embeddings = self.get_output_embeddings() if output_embeddings is not None: self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())方式二：在定义模型层时直接把输入 Embedding 权重传给输出层
这种方式在定义模型层时，直接将input_embeding的 weight 赋值给输出层 weight，即把 Embedding 的 weight 直接传给 head 来构建 Linear 输出层。RFC 引用的 ERNIE 实现（paddlenlp/transformers/ernie/modeling.py 中的ErnieLMPredictionHead）展示了这种模式：
class ErnieLMPredictionHead(nn.Layer): r""" Ernie Model with a `language modeling` head on top. """ def __init__( self, config: ErnieConfig, embedding_weights=None, weight_attr=None, ): super(ErnieLMPredictionHead, self).__init__() self.transform = nn.Linear(config.hidden_size, config.hidden_size, weight_attr=weight_attr) self.activation = getattr(nn.functional, config.hidden_act) self.layer_norm = nn.LayerNorm(config.hidden_size) self.decoder_weight = ( self.create_parameter( shape=[config.vocab_size, config.hidden_size], dtype=self.transform.weight.dtype, attr=weight_attr, is_bias=False, ) if embedding_weights is None else embedding_weights ) self.decoder_bias = self.create_parameter( shape=[config.vocab_size], dtype=self.decoder_weight.dtype, is_bias=True )调研结论
从调研可以看出，PaddleNLP 内大部分tie_weights实现是直接在模型 layer 定义层面实现的，而不是像 HuggingFace Transformers 一样在模型基类外统一实现。这种分散式实现的缺点是：每个模型都要自己写一遍绑定逻辑，代码重复且容易出错。RFC 的核心命题由此提出——能否在模型基类（如PretrainedModel）层面统一实现 tie_weights，让调用者不再为每个模型重复开发，最佳落点被建议放在模型基类model_utils.py中统一实现。
三、业内方案调研：三大主流实现对照
RFC 调研了业内主流深度学习框架与库的实现方式，作为设计参考。
1. HuggingFace Transformers（PyTorch 动态图）
Transformers 库在PreTrainedModel基类中统一实现了tie_weights方法（RFC 引用 v4.26.1 的实现），其核心逻辑如下：
def tie_weights(self): """ Tie the weights between the input embeddings and the output embeddings. If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning the weights instead. """ if getattr(self.config, "tie_word_embeddings", True): output_embeddings = self.get_output_embeddings() if output_embeddings is not None: self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings()) if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False): if hasattr(self, self.base_model_prefix): self = getattr(self, self.base_model_prefix) self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix) for module in self.modules(): if hasattr(module, "_tie_weights"): module._tie_weights()可以看到 HF 的实现覆盖了三类绑定场景：词嵌入绑定（tie_word_embeddings）、encoder-decoder 权重绑定（tie_encoder_decoder）、以及通过遍历子模块调用各模块自定义的_tie_weights钩子。
2. Tensor2Tensor（TensorFlow）
Tensor2Tensor 通过在命名空间层面复用变量实现共享（shared_embedding_and_softmax_weights开关），当开启共享时使用"shared"的variable_scope配合tf.AUTO_REUSE，使 softmax 权重与 Embedding 权重指向同一份变量：
def symbol_top(body_output, targets, model_hparams, vocab_size): del targets # unused arg if model_hparams.shared_embedding_and_softmax_weights: scope_name = "shared" reuse = tf.AUTO_REUSE else: scope_name = "softmax" reuse = False with tf.variable_scope(scope_name, reuse=reuse): body_output_shape = common_layers.shape_list(body_output) var = get_weights(model_hparams, vocab_size, body_output_shape[-1]) if (model_hparams.factored_logits and model_hparams.mode == tf_estimator.ModeKeys.TRAIN): # insert channels dimension body_output = tf.expand_dims(body_output, 3) return common_layers.FactoredTensor(body_output, var) else: body_output = tf.reshape(body_output, [-1, body_output_shape[-1]]) logits = tf.matmul(body_output, var, transpose_b=True) return tf.reshape(logits, body_output_shape[:-1] + [1, vocab_size])3. Fairseq（PyTorch）
Fairseq 的 FConv 模型通过直接把输出线性层fc3的 weight 指向embed_tokens.weight完成共享，同时用断言保证维度一致（共享要求out_embed_dim == embed_dim）：
self.fc2 = Linear(in_channels, out_embed_dim) if share_embed: assert out_embed_dim == embed_dim, ( "Shared embed weights implies same dimensions " " out_embed_dim={} vs embed_dim={}".format(out_embed_dim, embed_dim) ) self.fc3 = nn.Linear(out_embed_dim, num_embeddings) self.fc3.weight = self.embed_tokens.weight else: self.fc3 = Linear(out_embed_dim, num_embeddings, dropout=dropout)对比结论
由于 PaddlePaddle 与 HuggingFace Transformers 都基于动态图开发，RFC 明确选择参照 HuggingFace Transformers 的tie_weights函数思路在 PaddleNLP 中实现：以配置开关控制、通过get_input_embeddings()/get_output_embeddings()获取权重对象、最终以赋值方式让二者指向同一份参数。
四、设计思路与实现方案
RFC 将tie_weights的实现收敛为清晰的三步：
- 获取模型 input embedding 权重对象 A；
- 获取模型 output embedding 权重对象 B；
- 让 A 和 B 都指向同一个权重值（即
B.weight = A.weight，二者共享同一份 Parameter）。
命名与参数设计
命名与参数设计遵循飞桨 API 设计及命名规范（RFC 中引用了飞桨官方 API 设计指南）。关键的开关参数为tie_word_embeddings：
- 类型：
bool； - 默认值：在 PaddleNLP 的
PretrainedConfig基类中默认取True（见下文源码解读），即默认开启词嵌入绑定； - 语义：是否将输入与输出词 Embedding 权重绑定（"Whether input and output word embeddings should be tied for all MLM, LM and Seq2Seq models."）。
底层 OP 设计
从实现角度看，tie_weights不需要新增底层 OP：它只是参数层面的引用共享（Parameter 赋值/别名），不引入新的计算逻辑，所有前向/反向计算仍然复用已有的Embedding、Linear等算子。这正是"在模型外统一实现"能够成立的根本原因。
五、仓库源码级解读：tie_weights 在 PaddleNLP 的最终落地
RFC 提出的设计最终在 PaddleNLP 的PretrainedModel基类中统一落地。以下是当前仓库中的实际实现（paddlenlp/transformers/model_utils.py 中的tie_weights方法）：
def tie_weights(self): """ Tie the weights between the input embeddings and the output embeddings. """ if self.config.tie_word_embeddings: output_embeddings = self.get_output_embeddings() input_embeddings = self.get_input_embeddings() if output_embeddings is not None and input_embeddings is not None: if input_embeddings.weight.shape != output_embeddings.weight.shape: logger.warning( f"The shape of input embeddings is {input_embeddings.weight.shape} and the shape of output embeddings is {output_embeddings.weight.shape}. " "This is only expected if you are calling the `resize_token_embeddings` method" ) output_embeddings.weight = input_embeddings.weight if getattr(output_embeddings, "bias", None) is not None: # need to pad if output_embeddings.weight.shape[0] > output_embeddings.bias.shape[0]: old_bias = output_embeddings.bias pad_length = output_embeddings.weight.shape[0] - old_bias.shape[0] output_embeddings.bias = output_embeddings.create_parameter( shape=[output_embeddings.weight.shape[0]], attr=output_embeddings._bias_attr, dtype=output_embeddings._dtype, is_bias=True, ) new_bias = paddle.concat( [old_bias, paddle.zeros([pad_length], dtype=output_embeddings.bias.dtype)] ) output_embeddings.bias.set_value(new_bias) # need to trim elif output_embeddings.weight.shape[0] < output_embeddings.bias.shape[0]: new_bias = output_embeddings.bias[: output_embeddings.weight.shape[0]] output_embeddings.bias = output_embeddings.create_parameter( shape=[output_embeddings.weight.shape[0]], attr=output_embeddings._bias_attr, dtype=output_embeddings._dtype, is_bias=True, ) output_embeddings.bias.set_value(new_bias)该实现相比 RFC 的最初三步设计，落地时补充了两个重要细节：
- 形状不匹配告警：当输入、输出 Embedding 形状不一致时（典型场景是调用
resize_token_embeddings调整词表后），打印 warning 提示而非静默失败； - bias 对齐处理：当输出层带 bias 且词表大小变化时，自动对 bias 做pad（补零）或 trim（截断），保证共享权重后输出层仍然可用。
与基类配套的三个关键钩子方法定义在同一文件中：
get_input_embeddings()：返回模型输入 Embedding（nn.Embedding），支持通过base_model_prefix转发到 base model，未实现时抛出NotImplementedError；set_input_embeddings(value)：设置新的输入 Embedding；get_output_embeddings()：返回模型输出 Embedding，默认返回None，需要带输出头（LM Head）的模型覆写。
此外，基类的resize_token_embeddings在调整词表后会主动调用self.tie_weights()（见 model_utils.py 中resize_token_embeddings方法），保证词表伸缩后绑定关系依然成立。
配置开关：tie_word_embeddings
配置项tie_word_embeddings在 paddlenlp/transformers/configuration_utils.py 的PretrainedConfig基类中被统一解析：
self.tie_word_embeddings = kwargs.pop( "tie_word_embeddings", True ) # Whether input and output word embeddings should be tied for all MLM, LM and Seq2Seq models.即默认开启绑定；各模型可以在自己的配置类中覆写默认值，例如当前仓库中：
- llama/configuration.py：默认
tie_word_embeddings=False； - gemma/configuration.py：默认
tie_word_embeddings=True； - codegen/configuration.py 与 gptj/configuration.py：默认
False； - deepseek_v2/configuration.py：默认
False。
模型侧对接示例：Llama
以当前仓库中的 Llama 为例（paddlenlp/transformers/llama/modeling.py），LlamaForCausalLM在初始化时根据配置决定是否绑定，并将绑定后的 head 显式构造为转置共享形式：
self.llama = LlamaModel(config) if config.tie_word_embeddings: self.lm_head = LlamaLMHead(config, embedding_weights=self.llama.embed_tokens.weight, transpose_y=True) self.tie_weights() else: self.lm_head = LlamaLMHead(config)结合LlamaModel中覆写的get_input_embeddings()/set_input_embeddings()钩子（返回/设置self.embed_tokens），tie_weights()即可在基类层面完成输入输出权重绑定。
兼容旧实现的过渡：ConvBERT 的模型内实现
除基类统一实现外，个别模型（如 ConvBERT，paddlenlp/transformers/convbert/modeling.py）仍保留了模型内的tie_weights与_tie_or_clone_weights实现，用于处理输出权重需要转置复制（transpose）的特殊情况：
def _tie_or_clone_weights(self, output_embeddings, input_embeddings): """Tie or clone layer weights""" if output_embeddings.weight.shape == input_embeddings.weight.shape: output_embeddings.weight = input_embeddings.weight elif output_embeddings.weight.shape == input_embeddings.weight.t().shape: output_embeddings.weight.set_value(input_embeddings.weight.t()) else: ...这印证了 RFC 调研中提到的"现有实现分散在模型内部"的历史背景，也说明统一基类实现与模型内特殊实现可以共存。
六、测试与验收方案：以权重 id 一致性为核心
RFC 在测试章节提出了两种 tie_weights 的验收办法：
- 直接判断 id 一致性：判断输出层 weight 与输入层 weight 的对象 id 是否一致，一致即通过，否则 Failed；
- 训练若干 step 后验证：经过几个前反向之后，检查输出层 weight 与输入层 weight 是否仍保持一致，一致即通过。
RFC 最终选定第一种方式：用 id 的一致性判断绑定是否成功，简单高效。具体做法是构建单元测试，断言get_input_embeddings()得到的权重 id 与get_output_embeddings()得到的权重 id 一致。
该验收方案在仓库测试框架中得到了落实。在 tests/transformers/test_modeling_common.py 的ModelTesterMixin.test_tie_weight中：
def test_tie_weight(self): # test whether id of input_embeding equal id of output_embeding ? if not self.test_tie_weights: return config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() for model_class in self.all_model_classes: if "CausalLM" not in model_class.__name__ and "MaskedLM" not in model_class.__name__: continue model = self._make_model_instance(config, model_class) if not model.config.tie_word_embeddings: continue if hasattr(model, "get_input_embeddings") and hasattr(model, "get_output_embeddings"): try: input_embeddings = model.get_input_embeddings() except NotImplementedError: continue try: output_embeddings = model.get_output_embeddings() except NotImplementedError: continue if input_embeddings is not None and output_embeddings is not None: if hasattr(output_embeddings, "weight"): output_embeddings_weight = output_embeddings.weight else: output_embeddings_weight = output_embeddings if hasattr(input_embeddings, "weight"): input_embeddings_weight = input_embeddings.weight else: input_embeddings_weight = input_embeddings print( "model name :{},id is{},{}".format( model_class, id(output_embeddings_weight), id(input_embeddings_weight) ) ) self.assertEqual(id(output_embeddings_weight), id(input_embeddings_weight))测试逻辑与 RFC 的验收设计完全一致：只对CausalLM/MaskedLM类模型生效、要求config.tie_word_embeddings为真、通过id(...)断言输入输出权重为同一对象。各模型测试类通过test_tie_weights = True/False控制是否执行该用例，例如 tests/transformers/bert/test_modeling.py、tests/transformers/albert/test_modeling.py、tests/transformers/convbert/test_modeling.py、tests/transformers/electra/test_modeling.py 等均开启，而 tests/transformers/bloom/test_modeling.py 中test_tie_weights = False。
七、可行性验证脚本与排期规划
RFC 在可行性分析章节提供了一个可直接运行的验证脚本，用两个独立paddle.nn.Embedding对象验证"赋值绑定"是否真正共享参数，并验证修改其中一个是否会同步影响另一个：
import numpy as np from paddle.nn import Embedding """step1 定义两个不同的embedding 对象 AA 和 BB""" print('------------step1') AA = Embedding(1,2) BB = Embedding(1,2) AA.weight = BB.weight # 进行权重的绑定 """ step2 测试一下绑定结果""" print('------------step2') print('检测 AA 和 BB 的id是否一致:', AA is BB,id(AA), id(BB)) # AA 和 BB 的id 不一致 print('检测 AA.weight 和 BB.weight 的id是否一致:',AA.weight is BB.weight,id(AA.weight), id(BB.weight)) # 但是AA.weight 和 BB.weight 的id是一致的 print("AA.weight: ",AA.weight) print("BB.weight: ",BB.weight) """ step3 尝试修改一下AA的weight的值 BB的weight的值是否也跟着会一起修改""" # 修改一下其中一个AA 的权重值, 看一下 BB的权重值会不会变化 print('------------step3') AA.weight.set_value(np.array([[4.0,6.0]],dtype=np.float32)) print('检测 修改后的 AA.weight 和 BB.weight 的id是否一致:',AA.weight is BB.weight,id(AA.weight), id(BB.weight)) # AA.weight 和 BB.weight 的id是一致的 print("AA.weight 修改后的值: ",AA.weight) print("BB.weight:",BB.weight)脚本的预期结论是：
AA与BB两个 Layer 对象本身 id 不同（AA is BB为False）；- 但
AA.weight与BB.weight的 id 一致（AA.weight is BB.weight为True），证明二者指向同一 Parameter； - 通过
AA.weight.set_value(...)修改后，BB.weight的值同步变化，证明参数真正共享。
该脚本从实验层面验证了"在动态图中通过 Parameter 赋值即可实现权重共享"这一核心前提，为在基类统一实现提供了可行性依据。
RFC 给出的时间与开发排期规划（主要 milestone）为：
| 时间节点 | 里程碑 |
|---|---|
| 3.10 | 与官方确认开发思路 |
| 3.17 | 提交实现代码 |
八、总结
围绕 RFC No.103，PaddleNLP 的tie_weights能力经历了从"问题提出 → 现状与业内调研 → 三步设计 → 基类统一实现 → id 一致性验收"的完整闭环：
- 统一入口：
tie_weights最终在PretrainedModel基类（paddlenlp/transformers/model_utils.py）统一实现，调用者无需为每个模型重复编写绑定逻辑； - 配置驱动：通过
PretrainedConfig.tie_word_embeddings（默认True，configuration_utils.py）控制开关，各模型可覆写默认值； - 钩子协作：模型只需实现/覆写
get_input_embeddings()与get_output_embeddings()，基类自动完成绑定，并在resize_token_embeddings词表伸缩后自动重绑； - 可靠验收：以权重对象
id一致性为核心判据的test_tie_weight用例已进入通用模型测试框架（tests/transformers/test_modeling_common.py），被 BERT、ALBERT、ConvBERT、ELECTRA 等多个模型测试继承启用。
对于希望在自己的预训练模型中启用权重绑定的开发者，只需三步：在配置中开启tie_word_embeddings、实现get_input_embeddings/get_output_embeddings两个钩子、并在模型初始化末尾调用self.tie_weights()（或复用基类逻辑），即可获得参数共享带来的参数量缩减与训练收益。
【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址: https://gitcode.com/gh_mirrors/pa/PaddleNLP
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
