[train] moe support aux_loss (#5187)

Author: Jintao-Huang

docs/source/Instruction/命令行参数.md

@@ -162,6 +162,7 @@
 - 🔥report_to: 默认值为`tensorboard`。你也可以指定`--report_to tensorboard wandb swanlab`、`--report_to all`。
 - logging_first_step: 是否记录第一个step的日志，默认为True。
 - logging_steps: 日志打印间隔，默认为5。
+- router_aux_loss_coef: 用于moe模型训练时，设置 aux_loss 的权重。默认为None，使用config中值。若设置为0，则不计算 aux_loss。
 - logging_dir: tensorboard日志路径。默认为None，即设置为`f'{self.output_dir}/runs'`。
 - predict_with_generate: 验证时使用生成式的方式，默认为False。
 - metric_for_best_model: 默认为None，即当`predict_with_generate`设置为False时，设置为'loss'，否则设置为'rouge-l'（在PPO训练时，不进行默认值设置；GRPO训练设置为'reward'）。

docs/source_en/Instruction/Command-line-parameters.md

@@ -165,6 +165,7 @@ This parameter list inherits from transformers `Seq2SeqTrainingArguments`, with
 - 🔥report_to: Default value is `tensorboard`. You can also specify `--report_to tensorboard wandb swanlab` or `--report_to all`.
 - logging_first_step: Whether to log the first step, defaults to True.
 - logging_steps: Interval for logging, defaults to 5.
+- router_aux_loss_coef: Weight for aux_loss when training MoE models. Defaults to None, meaning the value from the config is used. If set to 0, aux_loss is not computed.
 - logging_dir: The path for TensorBoard logs. Defaults to None, which means it is set to `f'{self.output_dir}/runs'`.
 - predict_with_generate: Whether to use generative method during validation, default is False.
 - metric_for_best_model: Default is None, which means that when predict_with_generate is set to False, it is set to 'loss'; otherwise, it is set to 'rouge-l' (during PPO training, the default value is not set; in GRPO training, it is set to 'reward').

swift/trainers/arguments.py

@@ -30,6 +30,7 @@ class TrainArgumentsMixin:
     gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None
     logging_first_step: bool = True
     logging_steps: int = 5
+    router_aux_loss_coef: Optional[float] = None
 
     weight_decay: float = 0.1
     adam_beta2: float = 0.95

swift/trainers/trainers.py

@@ -1,5 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 # Part of the implementation is borrowed from huggingface/transformers.
+import inspect
 import os
 from contextlib import contextmanager, nullcontext
 from functools import wraps
@@ -15,6 +16,7 @@
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.utils import is_peft_available
 
+from swift.plugin import MeanMetric
 from swift.utils import JsonlWriter, Serializer, gc_collect, get_logger, unwrap_model_for_generation
 from .arguments import Seq2SeqTrainingArguments, TrainingArguments
 from .mixin import DataLoaderMixin, SwiftMixin
@@ -335,6 +337,16 @@ def _prepare_inputs(self, inputs):
                 if self.args.tuner_backend == 'unsloth':
                     inputs['logits_to_keep'] = int(logits_to_keep.sum())
 
+        if self.model.model_info.is_moe_model:
+            base_model = self.template.get_base_model(self.model)
+            router_aux_loss_coef = self.args.router_aux_loss_coef
+            if router_aux_loss_coef is None:
+                router_aux_loss_coef = getattr(base_model.config, 'router_aux_loss_coef', None)
+            if router_aux_loss_coef is not None:
+                base_model.config.router_aux_loss_coef = router_aux_loss_coef
+                if router_aux_loss_coef > 0 and 'output_router_logits' in inspect.signature(
+                        base_model.forward).parameters:
+                    inputs['output_router_logits'] = True
         inputs['compute_loss_func'] = compute_loss_func
         inputs['loss_kwargs'] = loss_kwargs
         return inputs
@@ -346,8 +358,11 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
 
         if (self.label_smoother is not None or compute_loss_func is not None) and 'labels' in inputs:
             labels = inputs.pop('labels')
-
         outputs = model(**inputs)
+        if getattr(outputs, 'aux_loss', None) is not None:
+            if 'aux_loss' not in self._custom_metrics:
+                self._custom_metrics['aux_loss'] = MeanMetric(nan_value=None)
+            self._custom_metrics['aux_loss'].update(outputs.aux_loss)
         # Save past state if it exists
         # TODO: this needs to be fixed and made cleaner later.
         if self.args.past_index >= 0: