Implement auxiliary-loss-free load balancing

Author: lshpku

llm/model_config/DeepSeek-V3/config.json

@@ -10,6 +10,7 @@
       "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM"
     },
     "aux_loss_alpha": 0.001,
+    "aux_loss_free_gamma": 0.001,
     "bos_token_id": 0,
     "eos_token_id": 1,
     "ep_size": 1,

llm/run_pretrain.py

@@ -573,7 +573,8 @@ def main():
     callbacks = [StepFlexToken(), FP8QuantWeightCallback()]
 
     if getattr(config, "topk_method", None) == "noaux_tc":
-        callbacks += [MoECorrectionBiasAdjustCallback()]
+        aux_loss_free_gamma = getattr(config, "aux_loss_free_gamma", 0.0)
+        callbacks += [MoECorrectionBiasAdjustCallback(aux_loss_free_gamma)]
 
     trainer = PretrainingTrainer(
         model=model,

paddlenlp/trainer/trainer_callback.py

@@ -700,7 +700,7 @@ def get_stat(layer):
 
         model.apply(get_stat)
 
-        usages_tensor = paddle.stack(usages, 0)  # [num_layers, num_local_experts]
+        usages_tensor = paddle.stack(usages, 0)  # [num_layers, num_experts]
         if not hasattr(fleet, "_hcg"):
             dist.all_reduce(usages_tensor)
             return
@@ -721,10 +721,6 @@ def get_stat(layer):
         update = paddle.sign(usages_mean - usages_tensor) * self.update_lr
         update_list = list(update)
 
-        print('on_optimizer_end bias:', [bias.tolist() for bias in biases])
-        print('on_optimizer_end usage:', usages_tensor.tolist())
-        print('on_optimizer_end update:', update.tolist())
-
         def update_bias(layer):
             if isinstance(layer, PretrainedMoEGate) and layer.topk_method == "noaux_tc":
                 with paddle.no_grad():