Skip to content

Commit eaa39fe

Browse files
huangjiyicodex
andcommitted
[DSV4] Fix Muon slice config empty layer offset
Offset generated DSV4 layer keys by num_empty_layers_add_in_head so Muon slice configs target the physical checkpoint layer ids. Co-authored-by: Codex <codex@openai.com>
1 parent 1e3643f commit eaa39fe

1 file changed

Lines changed: 13 additions & 3 deletions

File tree

paddleformers/transformers/deepseek_v4/modeling.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,9 @@ def _moe_experts(matrix_3d_global, ortho_fn):
142142
moe_grouped_gemm = getattr(config, "moe_grouped_gemm", False)
143143
use_gated_attn = getattr(config, "use_gated_attn", False)
144144
csa_compress_ratios = getattr(config, "csa_compress_ratios", None)
145+
num_empty_layers_add_in_head = getattr(
146+
config, "num_empty_layers_add_in_head", 0
147+
) or 0
145148

146149
# Get Muon configuration from muon_configs
147150
muon_qkv_update_mode = muon_configs.get("muon_qkv_update_mode", "split_head")
@@ -293,18 +296,25 @@ def _add_layer_slice_config(prefix, layer_idx):
293296

294297
# Main layers
295298
for layer_idx in range(num_hidden_layers):
296-
_add_layer_slice_config(f"model.layers.{layer_idx}", layer_idx)
299+
_add_layer_slice_config(
300+
f"model.layers.{layer_idx + num_empty_layers_add_in_head}",
301+
layer_idx,
302+
)
297303

298304
# MTP layers
299305
if config.mtp_num_layers > 0:
300306
num_nextn_predict_layers = config.mtp_num_layers
301307
else:
302308
num_nextn_predict_layers = config.num_nextn_predict_layers if config.num_nextn_predict_layers else 0
303309
for layer_idx in range(num_nextn_predict_layers):
304-
_add_layer_slice_config(f"model.layers.{num_hidden_layers + layer_idx}", num_hidden_layers + layer_idx)
310+
_add_layer_slice_config(
311+
f"model.layers.{num_hidden_layers + num_empty_layers_add_in_head + layer_idx}",
312+
num_hidden_layers + layer_idx,
313+
)
305314
for layer_idx in range(num_nextn_predict_layers):
306315
_add_layer_slice_config(
307-
f"model.layers.{num_hidden_layers + layer_idx}.transformer_layer", num_hidden_layers + layer_idx
316+
f"model.layers.{num_hidden_layers + num_empty_layers_add_in_head + layer_idx}.transformer_layer",
317+
num_hidden_layers + layer_idx,
308318
)
309319

310320
return slice_config

0 commit comments

Comments
 (0)