@@ -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