【大模型训练】sglang 权重绑定和roll HF Meg相互转化
在许多语言模型中，输入的词嵌入矩阵和输出的语言模型头（lm_head）可以共享相同的权重矩阵，这样可以减少模型参数量。这个特性通过配置文件中的参数控制。
Answer
您询问的是 tie_embeddings_and_output_weights，这是指将输入词嵌入层（embedding）和输出层（lm_head）的权重绑定在一起的机制。
功能概述
在许多语言模型中，输入的词嵌入矩阵和输出的语言模型头（lm_head）可以共享相同的权重矩阵，这样可以减少模型参数量。这个特性通过配置文件中的 tie_word_embeddings 参数控制。
实现方式
1. ParallelLMHead 的 tie_weights 方法
ParallelLMHead 类提供了 tie_weights() 方法来实现权重绑定： 1
这个方法的逻辑是：
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:
Wiki pages you might want to explore:
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_tokens
else:
self.weight = embed_tokens.weight
return 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_tokens
else:
self.unpadded_vocab_size = config.vocab_size
self.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 False
if self.config.tie_word_embeddings:
self.lm_head = self.model.embed_tokens
else:
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 weights
if "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 = False
self.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_embeddings
else:
# TODO something wrong with ParallelLMHead with DP attention enabled
self.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 ：
与您的问题的关系
这段代码是在模型保存/转换时使用的，而不是在权重更新时使用的 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_weights
and 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 head
elif 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_weights
and 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 head
elif 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 stage
if 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 stage
if 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
return 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_weights
and 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 head
elif 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."},
)
更多推荐
26
0- 0
已为社区贡献49条内容
所有评论(0)