【大模型训练】sglang 权重绑定和roll HF Meg相互转化
Answer
您询问的是 tie_embeddings_and_output_weights，这是指将输入词嵌入层（embedding）和输出层（lm_head）的权重绑定在一起的机制。
功能概述
在许多语言模型中，输入的词嵌入矩阵和输出的语言模型头（lm_head）可以共享相同的权重矩阵，这样可以减少模型参数量。这个特性通过配置文件中的 tie_word_embeddings 参数控制。
实现方式
1. ParallelLMHead 的 tie_weights 方法
ParallelLMHead 类提供了 tie_weights() 方法来实现权重绑定： 1
这个方法的逻辑是：
- 对于 GGUF 量化模型，返回原始的
embed_tokens对象 2 - 对于其他情况，直接将
lm_head.weight指向embed_tokens.weight，实现权重共享 3
2. 在模型初始化中的使用
多个模型在初始化时会根据 config.tie_word_embeddings 决定是否绑定权重：
Granite 模型示例： 4
Olmo 模型示例： 5
Llama Eagle 模型示例： 6
3. 在权重加载时的处理
当 tie_word_embeddings=True 时，模型在加载权重时会跳过 lm_head.weight，因为它已经与 embed_tokens.weight 共享：
Granite 模型的处理： 7
Olmo 模型的处理： 8
MiniCPM3 模型的处理： 9
4. GGUF 模型的特殊处理
对于 GGUF 量化模型，加载器会检查是否存在 lm_head.weight，如果不存在则自动设置 tie_word_embeddings=True： 10
配置示例
在模型配置中，tie_word_embeddings 通常在 config.json 中定义：
{"tie_word_embeddings": true // 或 false
}
例如 InternVL 配置默认设置为 False： 11
优势
权重绑定的主要优势：
- 减少参数量：避免存储两份相同大小的权重矩阵
- 减少内存占用：特别是对于大词汇表的模型
- 可能提高泛化能力：输入和输出共享表示
Notes
不同模型对 tie_word_embeddings 的默认设置不同。例如 Llama 3.2 1B Instruct 设置为 True，而 Llama 3.1 8B Instruct 设置为 False。 12 在 BailingMoE 模型中，当启用 DP attention 时，lm_head 使用 use_attn_tp_group 参数来控制张量并行组。 13
Wiki pages you might want to explore:
- Model Execution (sgl-project/sglang)
- DeepSeek Model Optimizations (sgl-project/sglang)
Wiki pages you might want to explore:
- DeepSeek Model Optimizations (sgl-project/sglang)
Citations
File: python/sglang/srt/layers/vocab_parallel_embedding.py (L560-567)
def tie_weights(self, embed_tokens: VocabParallelEmbedding):"""Tie the weights with word embeddings."""# GGUF quantized embed_tokens.if self.quant_config and self.quant_config.get_name() == "gguf":return embed_tokenselse:self.weight = embed_tokens.weightreturn self
File: python/sglang/srt/models/granite.py (L322-329)
self.lm_head = ParallelLMHead(config.vocab_size,config.hidden_size,quant_config=quant_config,prefix=add_prefix("lm_head", prefix),)if self.config.tie_word_embeddings:self.lm_head.tie_weights(self.model.embed_tokens)
File: python/sglang/srt/models/granite.py (L400-404)
if "lm_head.weight" in name and self.config.tie_word_embeddings:# Input and output embeddings are tied, so the output embeddings# may not be present in the checkpoint. We assume that the input# embeddings are always present in the checkpoint.continue
File: python/sglang/srt/models/olmo.py (L308-318)
if config.tie_word_embeddings:self.lm_head = self.model.embed_tokenselse:self.unpadded_vocab_size = config.vocab_sizeself.lm_head = ParallelLMHead(self.unpadded_vocab_size,config.hidden_size,org_num_embeddings=config.vocab_size,quant_config=quant_config,prefix=add_prefix("lm_head", prefix),)
File: python/sglang/srt/models/olmo.py (L355-356)
if self.config.tie_word_embeddings and "lm_head.weight" in name:continue
File: python/sglang/srt/models/llama_eagle.py (L127-137)
# Llama 3.2 1B Instruct set tie_word_embeddings to True# Llama 3.1 8B Instruct set tie_word_embeddings to Falseif self.config.tie_word_embeddings:self.lm_head = self.model.embed_tokenselse:self.lm_head = ParallelLMHead(getattr(config, "hot_vocab_size", config.vocab_size),config.hidden_size,quant_config=quant_config,prefix=add_prefix("lm_head", prefix),)
File: python/sglang/srt/models/minicpm3.py (L470-471)
if self.config.tie_word_embeddings and "lm_head.weight" in name:continue
File: python/sglang/srt/model_loader/loader.py (L1453-1457)
# we can only know if tie word embeddings after mapping weightsif "lm_head.weight" in get_gguf_extra_tensor_names(local_model_path, gguf_weights_map):model_config.hf_config.update({"tie_word_embeddings": True})
File: python/sglang/srt/configs/internvl.py (L347-349)
# By default, we use tie_word_embeddings=False for models of all sizes.self.tie_word_embeddings = Falseself.llm_config.tie_word_embeddings = self.tie_word_embeddings
File: python/sglang/srt/models/bailing_moe.py (L816-827)
# tie_word_embeddings为true，复用tie_word_embeddings，反之是独立的if config.tie_word_embeddings:self.lm_head = self.model.word_embeddingselse:# TODO something wrong with ParallelLMHead with DP attention enabledself.lm_head = ParallelLMHead(config.vocab_size,config.hidden_size,quant_config=quant_config,prefix=add_prefix("lm_head", prefix),use_attn_tp_group=get_global_server_args().enable_dp_lm_head,)
Answer
您问的是 mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py:621-622 这段代码的含义 1 。
代码含义
这段代码的逻辑是：
if self.revert and name == MCORE_LM_HEAD:return None # don't need a duplicate lm head
“revert 模式” 指的是将分布式并行的模型权重转换回单机 HuggingFace 格式的过程 2 。这个转换器有两个方向：
- 正向转换（
revert=False）：从 HF 格式转换为 Megatron-Core 分布式格式 2 - 反向转换（
revert=True）：从 Megatron-Core 分布式格式转换回 HF 格式 2
为什么跳过 lm_head
当满足以下条件时，会跳过 lm_head 权重：
- 使用
tie_embeddings_and_output_weights=True：这个配置表示lm_head和word_embeddings共享权重 3 - 使用 pipeline parallelism：模型被切分到多个 pipeline stage 4
- 在 pipeline 最后一个 stage：
lm_head和word_embeddings都在最后一个 stage 5 - 正在进行 revert 转换：从分布式格式转回 HF 格式 6
在这种情况下，返回 None 意味着不需要单独保存 lm_head 权重，因为它和 word_embeddings 是同一份权重 7 。
完整的权重共享逻辑
在同一个函数中，您可以看到完整的处理逻辑 8 ：
- 正向转换时（
not self.revert）：将word_embeddings的权重加载到lm_head9 - 反向转换时（
self.revert）：跳过lm_head，避免重复保存 1
与您的问题的关系
这段代码是在模型保存/转换时使用的，而不是在权重更新时使用的 10 。您遇到的 lm_head 权重未更新问题，更可能是因为：
- 权重更新时序问题：在
dummy模式下，lm_head初始化为随机权重，但权重更新可能没有覆盖到它 - 参数名称不匹配：训练端发送的参数名称和推理端期望的参数名称不一致
- 权重共享配置不一致：训练端和推理端对
tie_embeddings_and_output_weights的配置不同
Notes
MCORE_LM_HEAD和MCORE_WORD_EMBEDDING是常量，分别定义为"output_layer.weight"和"embedding.word_embeddings.weight"11- 这个转换器主要用于 checkpoint 的保存和加载，例如在
convert_checkpoint_to_hf函数中 12 - 权重共享是 Qwen 等模型的常见配置，可以减少参数量 3
- 如果您想调试权重更新问题，建议在
update_parameter_in_bucket方法中添加日志，查看哪些参数被成功加载
Citations
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L33-34)
MCORE_WORD_EMBEDDING = "embedding.word_embeddings.weight"
MCORE_LM_HEAD = "output_layer.weight"
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L214-214)
revert: bool = False,
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L616-624)
if (self.mca_config.tie_embeddings_and_output_weightsand self.mca_config.pipeline_model_parallel_size > 1and self.is_pipeline_last_stage()):if self.revert and name == MCORE_LM_HEAD:return None # don't need a duplicate lm headelif not self.revert and name == MCORE_WORD_EMBEDDING:name = MCORE_LM_HEAD # load word embedding weight to lm head
File: mcore_adapter/src/mcore_adapter/models/converter/post_converter.py (L55-55)
mca_named_weights = dist_reverter(mca_name, weights)
File: mcore_adapter/src/mcore_adapter/models/converter/post_converter.py (L79-81)
def convert_checkpoint_to_hf(model_name_or_path: str, save_directory: str, torch_dtype: Optional["torch.dtype"] = None, verbose: bool = True
):
Answer
您问的是当 Pipeline Parallelism (PP) size 为 1 且 tie_embeddings_and_output_weights=True 时的情况。
关键区别
当 PP size 为 1 时，权重转换逻辑会有所不同。让我们看看 dist_convert 方法中的条件判断 1 ：
if (self.mca_config.tie_embeddings_and_output_weightsand self.mca_config.pipeline_model_parallel_size > 1 # 注意这个条件and self.is_pipeline_last_stage()
):if self.revert and name == MCORE_LM_HEAD:return None # don't need a duplicate lm headelif not self.revert and name == MCORE_WORD_EMBEDDING:name = MCORE_LM_HEAD # load word embedding weight to lm head
关键点：这个特殊处理逻辑只在 pipeline_model_parallel_size > 1 时才会触发 2 。
PP size = 1 的行为
当 PP size 为 1 时：
- 不会跳过
lm_head权重：因为条件pipeline_model_parallel_size > 1不满足，所以不会执行return None的逻辑 2 lm_head和word_embeddings仍然共享权重：虽然不会跳过转换，但在模型配置层面，它们仍然是同一份权重 3- 权重更新应该正常工作：因为没有特殊的跳过逻辑，
lm_head权重应该能够正常接收更新
为什么 PP > 1 时需要特殊处理
在 Pipeline Parallelism 中，word_embeddings 通常在第一个 stage，而 lm_head 在最后一个 stage。但当 tie_embeddings_and_output_weights=True 时，word_embeddings 也需要在最后一个 stage 存在（用作 lm_head） 4 ：
if self.name_match(weight_name, self.config.pre_process_weights):# mtp and tie_embeddings_and_output_weights use embedding weights in last stageif weight_name == MCORE_WORD_EMBEDDING and (self.mca_config.mtp_num_layers or self.mca_config.tie_embeddings_and_output_weights):if self.is_pipeline_last_stage():return True
这就是为什么在 PP > 1 时需要特殊处理：避免在保存 checkpoint 时重复保存同一份权重 5 。
对您问题的影响
如果您的配置是 PP size = 1 且 tie_embeddings_and_output_weights=True：
lm_head权重应该能正常更新：不会触发跳过逻辑- 如果仍然出现乱码：问题可能不在于这个转换逻辑，而在于：
- 权重更新的时序问题（第一次推理在权重更新之前）
- SGLang 在
dummy模式下的初始化问题 - 参数名称映射问题
建议
对于 PP size = 1 的情况，建议：
- 检查权重更新日志：确认
lm_head是否在load_weights调用中 - 对比参数名称：确认训练端和推理端的参数名称一致
- 如果问题仍然存在：考虑使用
load_format="auto"而不是dummy
Notes
- PP size = 1 是最简单的情况，不涉及跨 stage 的权重共享问题
- 在 PP size = 1 时，所有层都在同一个 stage，
is_pipeline_first_stage()和is_pipeline_last_stage()都返回True6 - 如果您使用的是 Qwen 模型，
tie_embeddings_and_output_weights默认为True3 - 这个转换逻辑主要影响 checkpoint 的保存和加载，对于运行时的权重更新，应该通过
update_parameter_in_bucket方法处理
Citations
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L266-272)
# mtp and tie_embeddings_and_output_weights use embedding weights in last stageif weight_name == MCORE_WORD_EMBEDDING and (self.mca_config.mtp_num_layers or self.mca_config.tie_embeddings_and_output_weights):if self.is_pipeline_last_stage():return Truereturn self.is_pipeline_first_stage()
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L295-303)
def is_pipeline_last_stage(self):return self.pipeline_model_parallel_rank == (self.mca_config.pipeline_model_parallel_size - 1) and self.virtual_pipeline_model_parallel_rank == ((self.mca_config.virtual_pipeline_model_parallel_size or 1) - 1)def is_pipeline_first_stage(self):return self.pipeline_model_parallel_rank == 0 and self.virtual_pipeline_model_parallel_rank == 0
File: mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py (L616-624)
if (self.mca_config.tie_embeddings_and_output_weightsand self.mca_config.pipeline_model_parallel_size > 1and self.is_pipeline_last_stage()):if self.revert and name == MCORE_LM_HEAD:return None # don't need a duplicate lm headelif not self.revert and name == MCORE_WORD_EMBEDDING:name = MCORE_LM_HEAD # load word embedding weight to lm head
File: mcore_adapter/src/mcore_adapter/models/model_config.py (L177-180)
tie_embeddings_and_output_weights: bool = field(default=False,metadata={"help": "Untie embeddings and output weights."},)