add moh

examples/experiments/deepseek_v3_pretrain/run_pretrain.py

@@ -32,6 +32,7 @@
 )
 from paddleformers.trainer import (
     FP8QuantWeightCallback,
+    IndexerBiasAdjustCallback,
     MoECorrectionBiasAdjustCallback,
     MoeExpertsGradScaleCallback,
     PdArgumentParser,
@@ -627,6 +628,15 @@ def main():
         moe_router_bias_update_rate = getattr(config, "moe_router_bias_update_rate", 0.001)
         callbacks += [MoECorrectionBiasAdjustCallback(moe_router_bias_update_rate)]
 
+    if getattr(config, "use_moh", False):
+        indexer_bias_update_rate = getattr(training_args, "indexer_bias_update_rate", 0.001)
+        callbacks += [
+            IndexerBiasAdjustCallback(
+                lr=indexer_bias_update_rate,
+                use_mp=getattr(training_args, "sequence_parallel", False),
+            )
+        ]
+
     def resume_from_custom_func(model):
         if training_args.resume_from_huggingface_ckpt:
             load_huggingface_ckpt(model, training_args.resume_from_huggingface_ckpt)

examples/experiments/paddlefleet/run_pretrain.py

@@ -30,6 +30,7 @@
 )
 from paddleformers.trainer import (
     FP8QuantWeightCallback,
+    IndexerBiasAdjustCallback,
     MoECorrectionBiasAdjustCallback,
     MoeExpertsGradScaleCallback,
     PdArgumentParser,
@@ -644,6 +645,15 @@ def main():
         moe_router_bias_update_rate = getattr(config, "moe_router_bias_update_rate", 0.001)
         callbacks += [MoECorrectionBiasAdjustCallback(moe_router_bias_update_rate)]
 
+    if getattr(config, "use_moh", False):
+        indexer_bias_update_rate = getattr(training_args, "indexer_bias_update_rate", 0.001)
+        callbacks += [
+            IndexerBiasAdjustCallback(
+                lr=indexer_bias_update_rate,
+                use_mp=getattr(training_args, "sequence_parallel", False),
+            )
+        ]
+
     def resume_from_custom_func(model):
         if training_args.resume_from_huggingface_ckpt:
             raise NotImplementedError("Resume from HuggingFace ckpt is not supported yet")

paddleformers/cli/train/deepseek_v3_pretrain/workflow.py

@@ -27,6 +27,7 @@
 )
 from paddleformers.trainer import (
     FP8QuantWeightCallback,
+    IndexerBiasAdjustCallback,
     MoECorrectionBiasAdjustCallback,
     MoeExpertsGradScaleCallback,
     StepFlexToken,
@@ -557,6 +558,18 @@ def run_dsv3_pretrain(model_args, data_args, generating_args, training_args):
         moe_router_bias_update_rate = getattr(config, "moe_router_bias_update_rate", 0.001)
         callbacks += [MoECorrectionBiasAdjustCallback(moe_router_bias_update_rate)]
 
+    # MoH: indexer_moh_bias load-balance update, active only when the model
+    # config enables MoH routing (dsv4-hybrid stack). The callback is a no-op
+    # if no CSAIndexer in the model carries the MoH buffers.
+    if getattr(config, "use_moh", False):
+        indexer_bias_update_rate = getattr(training_args, "indexer_bias_update_rate", 0.001)
+        callbacks += [
+            IndexerBiasAdjustCallback(
+                lr=indexer_bias_update_rate,
+                use_mp=getattr(training_args, "sequence_parallel", False),
+            )
+        ]
+
     def resume_from_custom_func(model):
         if training_args.resume_from_huggingface_ckpt:
             load_huggingface_ckpt(model, training_args.resume_from_huggingface_ckpt)

paddleformers/cli/train/sft/workflow.py

@@ -41,6 +41,7 @@
 from paddleformers.peft import LoRAConfig, LoRAModel
 from paddleformers.trainer import (
     FP8QuantWeightCallback,
+    IndexerBiasAdjustCallback,
     IntervalStrategy,
     MoECorrectionBiasAdjustCallback,
     MoeExpertsGradScaleCallback,
@@ -717,6 +718,17 @@ def fetch_and_serialize(generator, dtype):
     elif getattr(model_config.get_text_config(), "topk_method", None) == "quantile_balancing":
         callbacks += [MoEQuantileBalancingCallback()]
 
+    # MoH: indexer_moh_bias load-balance update. The callback is a no-op when
+    # no CSAIndexer in the model carries the MoH buffers, so it is safe to add
+    # whenever ``use_moh`` is on in the config.
+    if getattr(model_config.get_text_config(), "use_moh", False):
+        callbacks += [
+            IndexerBiasAdjustCallback(
+                lr=getattr(training_args, "indexer_bias_update_rate", 0.001),
+                use_mp=training_args.sequence_parallel,
+            )
+        ]
+
     if training_args.use_expert_parallel:
         callbacks += [MoeExpertsGradScaleCallback(training_args)]
 

paddleformers/trainer/__init__.py

@@ -78,6 +78,7 @@
         "StepFlexToken",
         "FP8QuantWeightCallback",
         "MoECorrectionBiasAdjustCallback",
+        "IndexerBiasAdjustCallback",
         "MoEQuantileBalancingCallback",
         "MoeExpertsGradScaleCallback",
         "MoEGateSpGradSyncCallBack",

paddleformers/trainer/trainer_callback.py

@@ -85,6 +85,7 @@ class StandardMoERouter:
     "StepFlexToken",
     "FP8QuantWeightCallback",
     "MoECorrectionBiasAdjustCallback",
+    "IndexerBiasAdjustCallback",
     "MoEQuantileBalancingCallback",
     "MoeExpertsGradScaleCallback",
     "MoEGateSpGradSyncCallBack",
@@ -864,6 +865,70 @@ def update_bias(layer):
         model.apply(update_bias)
 
 
+class IndexerBiasAdjustCallback(TrainerCallback):
+    """Per-step sign-based bias update for CSAIndexer MoH load balancing.
+
+    Mirrors ``MoECorrectionBiasAdjustCallback`` but targets
+    ``CSAIndexer.indexer_moh_bias`` / ``local_tokens_per_indexer_moh``.
+    """
+
+    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):
+        if getattr(args, "freeze_training", False):
+            logger.warning("freeze_training is enabled! indexer_moh_bias will NOT be updated.")
+            return
+
+        model = kwargs["model"]
+        lr_ratio_fn = get_lr_ratio_fn(kwargs.get("optimizer"))
+
+        modules_with_bias = [
+            m
+            for m in model.sublayers()
+            if hasattr(m, "indexer_moh_bias") and hasattr(m, "local_tokens_per_indexer_moh")
+        ]
+        if not modules_with_bias:
+            return
+
+        usages_tensor = paddle.stack(
+            [m.local_tokens_per_indexer_moh for m in modules_with_bias], axis=0
+        )  # [num_indexers, n_heads]
+
+        if hasattr(fleet, "_hcg"):
+            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)
+        else:
+            dist.all_reduce(usages_tensor)
+
+        usages_mean = usages_tensor.mean(-1, keepdim=True)  # [num_indexers, 1]
+        update = paddle.sign(usages_mean - usages_tensor) * self.update_lr  # [num_indexers, n_heads]
+        update = update.astype(paddle.float32)
+
+        with paddle.no_grad():
+            for i, m in enumerate(modules_with_bias):
+                # Skip if the indexer weights are frozen or lr ratio is 0.
+                # Use linear_wq_b.weight as the representative trainable param.
+                ref_param = getattr(getattr(m, "linear_weights_proj", None), "weight", None)
+                if ref_param is not None:
+                    frozen = ref_param.stop_gradient or (lr_ratio_fn is not None and not float(lr_ratio_fn(ref_param)))
+                    if frozen:
+                        m.local_tokens_per_indexer_moh.zero_()
+                        continue
+                m.indexer_moh_bias.add_(update[i])
+                m.local_tokens_per_indexer_moh.zero_()
+
+
 class MoEQuantileBalancingCallback(TrainerCallback):
     """PaddleFormers adapter for PaddleFleet's optimizer-step QB update."""
 

paddleformers/trainer/training_args.py

@@ -521,6 +521,15 @@ class TrainingArguments:
         and decreased for the experts with more assigned tokens."""
         },
     )
+    indexer_bias_update_rate: float = field(
+        default=0.001,
+        metadata={
+            "help": """The MoH indexer bias (indexer_moh_bias) is updated based on the number of tokens
+        routed to each indexer head in a global batch, where the bias is increased for the heads with
+        fewer assigned tokens and decreased for the heads with more assigned tokens. Only effective
+        when the model config sets use_moh=True."""
+        },
+    )
     freeze_training: bool = field(
         default=False,
         metadata={

paddleformers/transformers/configuration_utils.py

@@ -417,6 +417,27 @@ class LlmMetaConfig:
             0.01,
             "Loss coefficient for the DSA indexer; controls the weight of the indexer loss term.",
         ),
+        (
+            "use_moh",
+            bool,
+            False,
+            "Whether to enable Mixture-of-Heads (MoH) routing over the CSA indexer heads. "
+            "Requires `num_activated_heads` to be set.",
+        ),
+        (
+            "num_activated_heads",
+            Optional[int],
+            None,
+            "Number of indexer heads kept per token by MoH routing. Only read when `use_moh=True`; "
+            "must satisfy 1 <= num_activated_heads <= dsa_index_n_heads.",
+        ),
+        (
+            "indexer_bias_update_rate",
+            float,
+            0.001,
+            "Update rate for the MoH indexer_moh_bias (used by IndexerBiasAdjustCallback). Controls the "
+            "magnitude of the per-step sign-based load-balancing adjustment. Defaults to 0.001.",
+        ),
     ]
 
     mtp_attributes = [

paddleformers/transformers/deepseek_v4/configuration.py

@@ -152,6 +152,9 @@ def __init__(
         dsa_index_topk=512,
         dsa_indexer_loss_coeff=0.01,
         dsa_indexer_use_sparse_loss=True,
+        # === MoH (Mixture-of-Heads) indexer routing ===
+        use_moh=False,
+        num_activated_heads=None,
         # === mHC (Hyper-Connection) ===
         enable_hyper_connections=True,
         num_residual_streams=4,
@@ -253,6 +256,13 @@ def __init__(
         self.dsa_indexer_loss_coeff = dsa_indexer_loss_coeff
         self.dsa_indexer_use_sparse_loss = dsa_indexer_use_sparse_loss
 
+        # MoH (Mixture-of-Heads) indexer routing. ``num_activated_heads`` is the
+        # number of indexer heads kept per token; it is only read when
+        # ``use_moh`` is set, and PaddleFleet's TransformerConfig rejects the
+        # pair if it is missing or exceeds ``dsa_index_n_heads``.
+        self.use_moh = use_moh
+        self.num_activated_heads = num_activated_heads
+
         # mHC (Hyper-Connection)
         self.enable_hyper_connections = enable_hyper_connections
         self.num_residual_streams = num_residual_streams
@@ -318,6 +328,8 @@ def __init__(
             tensor_model_parallel_size=tensor_model_parallel_size,
             pipeline_model_parallel_size=pipeline_model_parallel_size,
             context_parallel_size=context_parallel_size,
+            use_moh=use_moh,
+            num_activated_heads=num_activated_heads,
             **kwargs,
         )
 

paddleformers/transformers/deepseek_v4/modeling.py

@@ -398,6 +398,7 @@ def _gen_aoa_config(cls, config: DeepseekV4Config):
         n_shared_experts = getattr(config, "n_shared_experts", 1)
         moe_n_hash_layers = getattr(config, "moe_n_hash_layers", 3)
         dense_mode = getattr(config, "csa_dense_mode", False)
+        use_moh = getattr(config, "use_moh", False)
         csa_compress_ratios = config.csa_compress_ratios
         num_head_empty_layers = (
             config.num_empty_layers_add_in_head
@@ -518,6 +519,12 @@ def _gen_aoa_config(cls, config: DeepseekV4Config):
                     f"{idx_src}.weights_proj.weight^T -> {idx_tgt}.linear_weights_proj.weight",
                     f"{idx_src}.wq_b.weight^T -> {idx_tgt}.linear_wq_b.weight",
                 ]
+                # V4_INDEXER_MOH: indexer_moh_bias (persistable buffer) is not in the pretrained checkpoint;
+                # randomly initialize them via add primitive.
+                if use_moh:
+                    stmts += [
+                        f"_ -> {idx_tgt}.indexer_moh_bias",
+                    ]
 
             # --- MoE Gate ---
             stmts += [f"{src}.ffn.gate.weight -> {tgt}.mlp.gate.weight, dtype='float32'"]
@@ -651,6 +658,12 @@ def _gen_aoa_config(cls, config: DeepseekV4Config):
                         f"{idx_src}.weights_proj.weight^T -> {idx_tgt}.linear_weights_proj.weight",
                         f"{idx_src}.wq_b.weight^T -> {idx_tgt}.linear_wq_b.weight",
                     ]
+                    # V4_INDEXER_MOH: indexer_moh_bias (persistable buffer) is not in the pretrained checkpoint;
+                    # randomly initialize them via add primitive.
+                    if use_moh:
+                        stmts += [
+                            f"_ -> {idx_tgt}.indexer_moh_bias",
+                        ]
 
             # --- MoE Gate (MTP layers are always non-hash, so always have bias) ---
             stmts += [