add dsa index aoa and log (#4490)

Author: xingmingyyj

paddleformers/trainer/trainer.py

@@ -2656,6 +2656,24 @@ def _maybe_log_save_evaluate(self, tr_loss, model, epoch, ignore_keys_for_eval,
             except (ImportError, AttributeError):
                 pass
 
+            # Add DSA indexer loss metrics if available
+            try:
+                from paddlefleet.transformer.dsa_attention import (
+                    DSAIndexerLossLoggingHelper,
+                )
+
+                if DSAIndexerLossLoggingHelper.tracker.get("values") is not None:
+                    loss_scale = 1.0 / self.args.gradient_accumulation_steps
+                    DSAIndexerLossLoggingHelper.reduce_loss_in_tracker()
+                    tracker = DSAIndexerLossLoggingHelper.tracker
+                    indexer_loss_values = tracker["values"] * loss_scale
+                    num_layers = indexer_loss_values.shape[0]
+                    avg_indexer_loss = indexer_loss_values.sum() / num_layers
+                    logs["indexer_loss"] = avg_indexer_loss.item()
+                    DSAIndexerLossLoggingHelper.clean_loss_in_tracker()
+            except (ImportError, AttributeError):
+                pass
+
             self._total_loss_scalar += tr_loss_scalar
             self._globalstep_last_logged = self.state.global_step
             self._globalstep_last_start_time = time.time()

paddleformers/trainer/training_args.py

@@ -690,6 +690,12 @@ class TrainingArguments:
             )
         },
     )
+
+    dsa_indexer_loss_coeff: bool = field(
+        default=0.01,
+        metadata={"help": "Loss coefficient for the DSA indexer; controls the weight of the indexer loss term."},
+    )
+
     sharding_comm_group_call_opt: bool = field(
         default=False,
         metadata={

paddleformers/transformers/aoa_config_base.py

@@ -63,6 +63,8 @@ class MoEAOAConfigParams:
     # Runtime config
     model_prefix: str = "model."
 
+    index_n_heads: int = 0
+
     # Extra statements to add
     extra_statements: List[str] = field(default_factory=list)
 
@@ -129,6 +131,7 @@ def _extract_params(cls, config: Any) -> MoEAOAConfigParams:
             use_qk_norm=getattr(config, "use_qk_norm", False),
             has_shared_experts=cls._has_shared_experts(config),
             model_prefix=cls._get_model_prefix(config),
+            index_n_heads=getattr(config, "index_n_heads", 0),
         )
 
     @classmethod
@@ -369,6 +372,23 @@ def _get_mla_attention_statements(cls, params: MoEAOAConfigParams, prefix: str,
                 ]
             )
 
+        if params.index_n_heads and params.index_n_heads > 0:
+            indexer_weights = [
+                "wq_b",
+                "wk",
+                "weights_proj",
+            ]
+            statements.extend(
+                [
+                    f"{prefix}.self_attn.indexer.{weight_name}.weight^T -> {prefix_offset}.self_attn.core_attention.indexer.{weight_name}.weight"
+                    for weight_name in indexer_weights
+                ]
+            )
+            statements += [
+                f"{prefix}.self_attn.indexer.k_norm.bias ->  {prefix_offset}.self_attn.core_attention.indexer.k_norm.bias",
+                f"{prefix}.self_attn.indexer.k_norm.weight ->  {prefix_offset}.self_attn.core_attention.indexer.k_norm.weight",
+            ]
+
         return statements
 
     # ==================== MoE Expert Weights ====================
@@ -725,6 +745,23 @@ def _get_inv_mla_attention_statements(
                 ]
             )
 
+        if params.index_n_heads and params.index_n_heads > 0:
+            indexer_weights = [
+                "wq_b",
+                "wk",
+                "weights_proj",
+            ]
+            statements.extend(
+                [
+                    f"{prefix_offset}.self_attn.core_attention.indexer.{weight_name}.weight^T -> {prefix}.self_attn.indexer.{weight_name}.weight"
+                    for weight_name in indexer_weights
+                ]
+            )
+            statements += [
+                f"{prefix_offset}.self_attn.core_attention.indexer.k_norm.bias -> {prefix}.self_attn.indexer.k_norm.bias",
+                f"{prefix_offset}.self_attn.core_attention.indexer.k_norm.weight -> {prefix}.self_attn.indexer.k_norm.weight",
+            ]
+
         return statements
 
     # ==================== Inverse MoE Expert Weights ====================