bugfix: Bind embed_token and lm_head weights during DCP weight loading when tie_word

已合并

bugfix: Bind embed_token and lm_head weights during DCP weight loading when tie_word_embeddings is True #3039

bugfix: Bind embed_token and lm_head weights during DCP weight loading when tie_word_embeddings is True #3039

已合并

LKONE创建于 28 天前

共 4 个文件变更+8-25

	@@ -121,29 +121,6 @@ pip list \| grep fla_npu
121		121
122	如果使用fsdp2的meta init初始化模型或MoE模型需要支持mtp，都需要先根据模型配置完成以下权重转换：	122	如果使用fsdp2的meta init初始化模型或MoE模型需要支持mtp，都需要先根据模型配置完成以下权重转换：
123		123
124	-(1) 模型配置文件config.json中的`tie_word_embeddings`字段为`true`时（例如0.8B，2B，4B模型），使用以下转换脚本：
125	-
126	-```bash
127	-mm-convert Qwen35Converter hf_to_dcp \
128	---hf_dir ckpt/hf_path/xxxxxxx \
129	---dcp_dir ckpt/dcp_path/xxxxxxx \
130	---tie_weight_mapping '{"lm_head.weight":"model.language_model.embed_tokens.weight"}' \
131	---num_workers 0
132	-
133	-# 其中：
134	-# hf_dir: huggingface权重目录
135	-# dcp_dir: 转换后DCP格式的权重保存目录
136	-# tie_weight_mapping: 权重绑定映射关系
137	-# num_workers: 并行工作线程数，0表示串行执行，若存储IO性能允许，可适当调大并发数以提升转换效率，推荐设置为4
138	-
139	-# 转换后的目录结构为：
140	-# ———— xxxxxxx
141	-# \|—— release
142	-# \|—— latest_checkpointed_iteration.txt
143	-```
144	-
145	-(2) 其它场景:
146	-
147	```bash	124	```bash
148	mm-convert Qwen35Converter hf_to_dcp \	125	mm-convert Qwen35Converter hf_to_dcp \
149	--hf_dir ckpt/hf_path/xxxxxxx \	126	--hf_dir ckpt/hf_path/xxxxxxx \

	@@ -478,6 +478,10 @@ class Qwen3TTSConfig(PretrainedConfig):
478	tts_eos_token_id=151673,	478	tts_eos_token_id=151673,
479	**kwargs,	479	**kwargs,
480	):	480	):
		481	+ # transformers 4.57.3 defaults tie_word_embeddings to True when the field
		482	+ # is not explicitly provided, which later causes _retie_embeddings to fail.
		483	+ # Force-set it to False before forwarding to the parent config.
		484	+ kwargs["tie_word_embeddings"] = False
481	super().__init__(**kwargs)	485	super().__init__(**kwargs)
482		486
483	if talker_config is None:	487	if talker_config is None:

	@@ -42,6 +42,7 @@ from mindspeed_mm.fsdp.utils.lora_utils import (
42	from mindspeed_mm.fsdp.utils.lora_weight_manager import LoraWeightManager	42	from mindspeed_mm.fsdp.utils.lora_weight_manager import LoraWeightManager
43	from mindspeed_mm.config.config_manager import ConfigManager	43	from mindspeed_mm.config.config_manager import ConfigManager
44	from mindspeed_mm.fsdp.utils.dtype import get_dtype	44	from mindspeed_mm.fsdp.utils.dtype import get_dtype
		45	+from mindspeed_mm.fsdp.checkpoint.hf_load_utils import _retie_embeddings
45		46
46		47
47	logger = logging.getLogger(__name__)	48	logger = logging.getLogger(__name__)
	@@ -212,6 +213,9 @@ class Trainer:
212	self.lora_weight_manager = LoraWeightManager(model)	213	self.lora_weight_manager = LoraWeightManager(model)
213	self.lora_weight_manager.verify_lora_weights()	214	self.lora_weight_manager.verify_lora_weights()
214		215
		216	+ # Re-tie embed_tokens and lm_head weights when tie_word_embeddings is true
		217	+ _retie_embeddings(model)
		218	+
215	return model	219	return model
216		220
217	def enable_lora(self, model: torch.nn.Module) -> torch.nn.Module:	221	def enable_lora(self, model: torch.nn.Module) -> torch.nn.Module: