PaddleNLP tie_weights 权重绑定能力设计与实现全解析（RFC No.103）

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_weightsNone, weight_attrNone, ): super(ErnieLMPredictionHead, self).__init__() self.transform nn.Linear(config.hidden_size, config.hidden_size, weight_attrweight_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], dtypeself.transform.weight.dtype, attrweight_attr, is_biasFalse, ) if embedding_weights is None else embedding_weights ) self.decoder_bias self.create_parameter( shape[config.vocab_size], dtypeself.decoder_weight.dtype, is_biasTrue )调研结论从调研可以看出PaddleNLP 内大部分tie_weights实现是直接在模型 layer 定义层面实现的而不是像 HuggingFace Transformers 一样在模型基类外统一实现。这种分散式实现的缺点是每个模型都要自己写一遍绑定逻辑代码重复且容易出错。RFC 的核心命题由此提出——能否在模型基类如PretrainedModel层面统一实现 tie_weights让调用者不再为每个模型重复开发最佳落点被建议放在模型基类model_utils.py中统一实现。三、业内方案调研三大主流实现对照RFC 调研了业内主流深度学习框架与库的实现方式作为设计参考。1. HuggingFace TransformersPyTorch 动态图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, cant 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. Tensor2TensorTensorFlowTensor2Tensor 通过在命名空间层面复用变量实现共享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, reusereuse): 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_bTrue) return tf.reshape(logits, body_output_shape[:-1] [1, vocab_size])3. FairseqPyTorchFairseq 的 FConv 模型通过直接把输出线性层fc3的 weight 指向embed_tokens.weight完成共享同时用断言保证维度一致共享要求out_embed_dim embed_dimself.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, dropoutdropout)对比结论由于 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( fThe 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]], attroutput_embeddings._bias_attr, dtypeoutput_embeddings._dtype, is_biasTrue, ) new_bias paddle.concat( [old_bias, paddle.zeros([pad_length], dtypeoutput_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]], attroutput_embeddings._bias_attr, dtypeoutput_embeddings._dtype, is_biasTrue, ) output_embeddings.bias.set_value(new_bias)该实现相比 RFC 的最初三步设计落地时补充了两个重要细节形状不匹配告警当输入、输出 Embedding 形状不一致时典型场景是调用resize_token_embeddings调整词表后打印 warning 提示而非静默失败bias 对齐处理当输出层带 bias 且词表大小变化时自动对 bias 做pad补零或 trim截断保证共享权重后输出层仍然可用。与基类配套的三个关键钩子方法定义在同一文件中get_input_embeddings()返回模型输入 Embeddingnn.Embedding支持通过base_model_prefix转发到 base model未实现时抛出NotImplementedErrorset_input_embeddings(value)设置新的输入 Embeddingget_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_embeddingsFalsegemma/configuration.py默认tie_word_embeddingsTruecodegen/configuration.py 与 gptj/configuration.py默认Falsedeepseek_v2/configuration.py默认False。模型侧对接示例Llama以当前仓库中的 Llama 为例paddlenlp/transformers/llama/modeling.pyLlamaForCausalLM在初始化时根据配置决定是否绑定并将绑定后的 head 显式构造为转置共享形式self.llama LlamaModel(config) if config.tie_word_embeddings: self.lm_head LlamaLMHead(config, embedding_weightsself.llama.embed_tokens.weight, transpose_yTrue) self.tie_weights() else: self.lm_head LlamaLMHead(config)结合LlamaModel中覆写的get_input_embeddings()/set_input_embeddings()钩子返回/设置self.embed_tokenstie_weights()即可在基类层面完成输入输出权重绑定。兼容旧实现的过渡ConvBERT 的模型内实现除基类统一实现外个别模型如 ConvBERTpaddlenlp/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]],dtypenp.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.103PaddleNLP 的tie_weights能力经历了从问题提出 → 现状与业内调研 → 三步设计 → 基类统一实现 → id 一致性验收的完整闭环统一入口tie_weights最终在PretrainedModel基类paddlenlp/transformers/model_utils.py统一实现调用者无需为每个模型重复编写绑定逻辑配置驱动通过PretrainedConfig.tie_word_embeddings默认Trueconfiguration_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），仅供参考