Implement auxiliary-loss-free load balancing

Author: lshpku

llm/run_pretrain.py

@@ -28,6 +28,7 @@
 )
 from paddlenlp.trainer import (
     FP8QuantWeightCallback,
+    MoECorrectionBiasAdjustCallback,
     PdArgumentParser,
     StepFlexToken,
     Trainer,
@@ -571,6 +572,9 @@ def main():
 
     callbacks = [StepFlexToken(), FP8QuantWeightCallback()]
 
+    if getattr(config, "topk_method", None) == "noaux_tc":
+        callbacks += [MoECorrectionBiasAdjustCallback()]
+
     trainer = PretrainingTrainer(
         model=model,
         args=training_args,

paddlenlp/trainer/trainer_callback.py

@@ -27,6 +27,11 @@
 import numpy as np
 from tqdm.auto import tqdm
 
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.fleet import fleet
+
+from paddlenlp.transformers.moe_gate import PretrainedMoEGate
 from paddlenlp.transformers.moe_utils import offload, reload
 from paddlenlp.utils.log import logger
 
@@ -44,6 +49,7 @@
     "EarlyStoppingCallback",
     "StepFlexToken",
     "FP8QuantWeightCallback",
+    "MoECorrectionBiasAdjustCallback",
 ]
 
 
@@ -671,3 +677,59 @@ def on_optimizer_begin(self, args, state, control, **kwargs):
         if (not g_shard_bypass_dygraph_optimizer) and hasattr(model, "fp8_quant_weight"):
             for name in self.moe_weights_name:
                 reload(optimizer._master_weights[name])
+
+
+class MoECorrectionBiasAdjustCallback(TrainerCallback):
+    """used for moe aux loss free balance"""
+
+    def __init__(self, lr=0.001, use_mp=False):
+        super().__init__()
+        self.update_lr = lr
+        self.use_mp = use_mp
+
+    def on_optimizer_end(self, args, state, control, **kwargs):
+        model = kwargs["model"]
+
+        biases = []
+        usages = []
+
+        def get_stat(layer):
+            if isinstance(layer, PretrainedMoEGate) and layer.topk_method == "noaux_tc":
+                biases.append(layer.e_score_correction_bias)
+                usages.append(layer.expert_usage)
+
+        model.apply(get_stat)
+
+        usages_tensor = paddle.stack(usages, 0)  # [num_layers, num_local_experts]
+        if not hasattr(fleet, "_hcg"):
+            dist.all_reduce(usages_tensor)
+            return
+
+        hcg = fleet.get_hybrid_communicate_group()
+        mp_group = hcg.get_model_parallel_group()
+        dp_group = hcg.get_data_parallel_group()
+        sd_group = hcg.get_sharding_parallel_group()
+
+        if self.use_mp and mp_group.nranks > 1:
+            dist.all_reduce(usages_tensor, group=mp_group)
+        if dp_group.nranks > 1:
+            dist.all_reduce(usages_tensor, group=dp_group)
+        if sd_group.nranks > 1:
+            dist.all_reduce(usages_tensor, group=sd_group)
+
+        usages_mean = usages_tensor.mean(-1, keepdim=True)
+        update = paddle.sign(usages_mean - usages_tensor) * self.update_lr
+        update_list = list(update)
+
+        print('on_optimizer_end bias:', paddle.stack(biases, 0).numpy())
+        print('on_optimizer_end usage:', usages_tensor.numpy())
+        print('on_optimizer_end update:', update.numpy())
+
+        def update_bias(layer):
+            if isinstance(layer, PretrainedMoEGate) and layer.topk_method == "noaux_tc":
+                with paddle.no_grad():
+                    if not layer.weight.stop_gradient:
+                        biases.pop(0).add_(update_list.pop(0))
+                    usages.pop(0).zero_()
+
+        model.apply(update_bias)

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -924,6 +924,11 @@ def __init__(
                 default_initializer=nn.initializer.Constant(0.0),
             )
             self.e_score_correction_bias.is_distributed = True
+            self.expert_usage = paddle.zeros(
+                shape=[num_experts],
+                dtype=paddle.int64,
+            )
+            self.expert_usage.stop_gradient = True
 
         if self.using_post_norm_recompute:
             assert norm_weight is not None and norm_eps is not None
@@ -969,6 +974,8 @@ def forward(self, hidden_states):
             scores, routing_map, exp_counts, l_aux, l_zloss = self.topkgating_nodrop(
                 scores
             )  # (scores, routing_map, exp_counts, l_aux, l_zloss)
+            with paddle.no_grad():
+                self.expert_usage += exp_counts
             ret = (scores, routing_map, l_aux, l_zloss)
         else:
             ret = self.topkgating(scores)  # (capacity, combine_weights, dispatch_mask, exp_counts, l_aux, l_zloss)

paddlenlp/transformers/moe_gate.py

@@ -301,7 +301,7 @@ def _topk_noaux_tc(
         assert n_experts % n_group == 0, "n_experts must be divisible by n_groups"
 
         assert self.e_score_correction_bias is not None, "e_score_correction_bias is None"
-        scores_for_choice = scores.reshape([bsz_seq_len, -1]) + self.e_score_correction_bias.unsqueeze(0)
+        scores_for_choice = scores.reshape([bsz_seq_len, -1]) + self.e_score_correction_bias.detach().unsqueeze(0)
         reshape_tmp_rst = scores_for_choice.reshape([bsz_seq_len, self.n_group, -1])
         top_k = min(reshape_tmp_rst.shape[2], 2)
         group_scores = reshape_tmp_rst.topk(top_k, axis=-1)[0].sum(axis=-1)  # fmt:skip [n, n_group]