diff --git a/examples/experiments/deepseek_v3_pretrain/run_pretrain.py b/examples/experiments/deepseek_v3_pretrain/run_pretrain.py index 3948b5a0033..e1de5aa3ca4 100644 --- a/examples/experiments/deepseek_v3_pretrain/run_pretrain.py +++ b/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) diff --git a/examples/experiments/paddlefleet/run_pretrain.py b/examples/experiments/paddlefleet/run_pretrain.py index d941721a717..573e40586f9 100644 --- a/examples/experiments/paddlefleet/run_pretrain.py +++ b/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") diff --git a/paddleformers/cli/train/deepseek_v3_pretrain/workflow.py b/paddleformers/cli/train/deepseek_v3_pretrain/workflow.py index a9d69077e70..e058c70109e 100644 --- a/paddleformers/cli/train/deepseek_v3_pretrain/workflow.py +++ b/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) diff --git a/paddleformers/cli/train/sft/workflow.py b/paddleformers/cli/train/sft/workflow.py index f07c0305686..a00cc4d4fb7 100644 --- a/paddleformers/cli/train/sft/workflow.py +++ b/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)] diff --git a/paddleformers/trainer/__init__.py b/paddleformers/trainer/__init__.py index 540c1638c0b..ce813075cf7 100644 --- a/paddleformers/trainer/__init__.py +++ b/paddleformers/trainer/__init__.py @@ -78,6 +78,7 @@ "StepFlexToken", "FP8QuantWeightCallback", "MoECorrectionBiasAdjustCallback", + "IndexerBiasAdjustCallback", "MoEQuantileBalancingCallback", "MoeExpertsGradScaleCallback", "MoEGateSpGradSyncCallBack", diff --git a/paddleformers/trainer/trainer.py b/paddleformers/trainer/trainer.py index b20047a1a39..b780f4543eb 100644 --- a/paddleformers/trainer/trainer.py +++ b/paddleformers/trainer/trainer.py @@ -272,6 +272,35 @@ DIST_MODEL_PATH = "dist_model" +def _read_flex_checkpoint_keys(ckpt_path): + """Best-effort read of the HF safetensors shard index to enumerate keys. + + Used by the flex-checkpoint load path (via ``_gen_aoa_config``) to + distinguish HF checkpoints that already carry persisted trainer state + (e.g. ``indexer_moh_bias``) from fresh HF releases. Returns ``None`` on + any failure -- callers must treat ``None`` as "unknown" and fall back to + the historical zero-init behavior. + """ + if not ckpt_path or not os.path.isdir(ckpt_path): + return None + # ``model.safetensors.index.json`` is the standard HF sharded-checkpoint + # index file. Non-sharded (single .safetensors) checkpoints won't have + # it, and we don't want to eagerly open every shard just to build the + # key set -- ``None`` is a safe fallback there too. + index_path = os.path.join(ckpt_path, "model.safetensors.index.json") + if not os.path.isfile(index_path): + return None + try: + with open(index_path, "r") as f: + index = json.load(f) + except (OSError, ValueError): + return None + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + return None + return set(weight_map.keys()) + + class Trainer: """ Trainer is a simple but feature-complete training and eval loop for PaddlePaddle, optimized for PaddleFormers. @@ -1256,7 +1285,17 @@ def get_metadata_file_name(path): worker_groups = None if self.args.load_from_hf: - hf_aoa_config = self.model._gen_aoa_config(self.model.config) + # V4_INDEXER_MOH round-trip: read the HF checkpoint's key set so + # _gen_aoa_config can emit ``named -> named`` for persisted trainer + # state (e.g. ``indexer_moh_bias``) and fall back to ``_ -> ...`` + # only for genuinely missing keys. Read the shard index first; + # older classmethods that don't accept ``checkpoint_keys`` keep + # their old signature via the TypeError fallback below. + _aoa_checkpoint_keys = _read_flex_checkpoint_keys(resume_from_checkpoint) + try: + hf_aoa_config = self.model._gen_aoa_config(self.model.config, checkpoint_keys=_aoa_checkpoint_keys) + except TypeError: + hf_aoa_config = self.model._gen_aoa_config(self.model.config) assert ( self.args.ignore_load_lr_and_optim ), "Loading from HuggingFace format is only allowed when learning rate and optimizer state are ignored." @@ -1848,7 +1887,15 @@ def train( if resume_from_checkpoint is not None: if self.args.convert_from_hf: model_sharded_state_dict = model.sharded_state_dict() - aoa_config = model._gen_aoa_config(model.config) + # See _load_flex_checkpoint for the V4_INDEXER_MOH round-trip + # rationale -- feed the HF checkpoint's key set through so + # persisted trainer state (e.g. ``indexer_moh_bias``) is + # loaded named->named instead of being zero-init'd. + _aoa_checkpoint_keys = _read_flex_checkpoint_keys(resume_from_checkpoint) + try: + aoa_config = model._gen_aoa_config(model.config, checkpoint_keys=_aoa_checkpoint_keys) + except TypeError: + aoa_config = model._gen_aoa_config(model.config) dist.load_state_dict( model_sharded_state_dict, resume_from_checkpoint, diff --git a/paddleformers/trainer/trainer_callback.py b/paddleformers/trainer/trainer_callback.py index f47f5f6f938..5c06b8d1ffa 100644 --- a/paddleformers/trainer/trainer_callback.py +++ b/paddleformers/trainer/trainer_callback.py @@ -85,6 +85,7 @@ class StandardMoERouter: "StepFlexToken", "FP8QuantWeightCallback", "MoECorrectionBiasAdjustCallback", + "IndexerBiasAdjustCallback", "MoEQuantileBalancingCallback", "MoeExpertsGradScaleCallback", "MoEGateSpGradSyncCallBack", @@ -864,6 +865,73 @@ 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_weights_proj.weight`` as the representative + # trainable param: it is the head-scoring projection that + # actually drives head selection, so freezing it is what + # semantically means "this indexer is frozen". + 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.""" diff --git a/paddleformers/trainer/training_args.py b/paddleformers/trainer/training_args.py index 7eee5632617..434b79104b2 100644 --- a/paddleformers/trainer/training_args.py +++ b/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={ diff --git a/paddleformers/transformers/configuration_utils.py b/paddleformers/transformers/configuration_utils.py index 5720caad04f..1a320c3ec63 100644 --- a/paddleformers/transformers/configuration_utils.py +++ b/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 = [ diff --git a/paddleformers/transformers/deepseek_v4/configuration.py b/paddleformers/transformers/deepseek_v4/configuration.py index d08cd732530..1d1ca429f2e 100644 --- a/paddleformers/transformers/deepseek_v4/configuration.py +++ b/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, ) diff --git a/paddleformers/transformers/deepseek_v4/modeling.py b/paddleformers/transformers/deepseek_v4/modeling.py index a28c6aff070..7c325ede011 100644 --- a/paddleformers/transformers/deepseek_v4/modeling.py +++ b/paddleformers/transformers/deepseek_v4/modeling.py @@ -382,7 +382,7 @@ def build_muon_param_info_map(cls, model, config): return info_map @classmethod - def _gen_aoa_config(cls, config: DeepseekV4Config): + def _gen_aoa_config(cls, config: DeepseekV4Config, checkpoint_keys=None): """Weight conversion: HuggingFace DSv4 checkpoint -> PaddleFleet internal format. Maps open-source DeepSeek-V4 HuggingFace parameter names to PaddleFleet names. @@ -392,12 +392,40 @@ def _gen_aoa_config(cls, config: DeepseekV4Config): HF naming convention: layers.{L}.attn.*, layers.{L}.ffn.*, embed.weight, etc. PF naming convention: model.layers.{L}.self_attn.*, model.layers.{L}.mlp.*, etc. + + Args: + config: Model config. + checkpoint_keys: Optional iterable of HF-side state-dict keys that the + loader has determined actually exist in the on-disk checkpoint (e.g. + ``sharded_metadata["all_checkpoint_keys"]``). Used **only** to decide + the load rule for the trained-but-persistable ``indexer_moh_bias`` + buffer (see V4_INDEXER_MOH below): + + * When ``indexer_moh_bias`` **is** present in the checkpoint (i.e. + a prior ``save_pretrained`` wrote it back via + ``_gen_inv_aoa_config``), the forward AOA emits a named + ``hf_key -> fleet_key`` mapping so the trained aux-loss-free + bias is loaded, not overwritten. + * When the key is **absent** (fresh HuggingFace release, or a + legacy pre-round-trip save), we fall back to the ``_ -> + fleet_key`` add primitive so the bias is zero-initialized. + + ``checkpoint_keys=None`` preserves the historical behavior of + always zero-initializing -- safe but drops the round-trip. Pass + a concrete set/list when the caller knows the checkpoint contents + (from_pretrained / _load_flex_checkpoint do; merge tools may not). """ + # Normalize to a set for O(1) membership checks; ``None`` sentinel is + # preserved as "caller doesn't know", which keeps the pre-existing + # zero-init fallback. + if checkpoint_keys is not None and not isinstance(checkpoint_keys, (set, frozenset)): + checkpoint_keys = set(checkpoint_keys) num_hidden_layers = config.num_hidden_layers num_experts = config.n_routed_experts 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 +546,26 @@ 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 is a *persistable* buffer + # mutated every step by the aux-loss-free callback. Two cases: + # * fresh HF release (or legacy pre-round-trip save): + # ``indexer_moh_bias`` is absent -> use the add primitive + # so the buffer is zero-initialized on load. + # * checkpoint saved by our own ``_gen_inv_aoa_config``: + # ``indexer_moh_bias`` is present -> map named->named so + # the trained load-balancing state actually loads. + # ``checkpoint_keys=None`` means the caller doesn't know, so + # we conservatively zero-init (backward-compatible). + if use_moh: + hf_bias_key = f"{src}.attn.indexer.indexer_moh_bias" + if checkpoint_keys is not None and hf_bias_key in checkpoint_keys: + stmts += [ + f"{hf_bias_key} -> {idx_tgt}.indexer_moh_bias", + ] + else: + stmts += [ + f"_ -> {idx_tgt}.indexer_moh_bias", + ] # --- MoE Gate --- stmts += [f"{src}.ffn.gate.weight -> {tgt}.mlp.gate.weight, dtype='float32'"] @@ -651,6 +699,19 @@ 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: same rationale as the decoder branch -- + # named->named when the checkpoint carries the trained bias + # (round-trip case), add primitive otherwise. + if use_moh: + hf_bias_key = f"{mtp_src}.attn.indexer.indexer_moh_bias" + if checkpoint_keys is not None and hf_bias_key in checkpoint_keys: + stmts += [ + f"{hf_bias_key} -> {idx_tgt}.indexer_moh_bias", + ] + else: + stmts += [ + f"_ -> {idx_tgt}.indexer_moh_bias", + ] # --- MoE Gate (MTP layers are always non-hash, so always have bias) --- stmts += [ @@ -702,6 +763,7 @@ def _gen_inv_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 @@ -824,6 +886,13 @@ def _gen_inv_aoa_config(cls, config: DeepseekV4Config): f"{idx_src}.linear_weights_proj.weight^T -> {idx_tgt}.weights_proj.weight", f"{idx_src}.linear_wq_b.weight^T -> {idx_tgt}.wq_b.weight", ] + # V4_INDEXER_MOH: same rationale as the decoder path -- the + # aux-loss-free bias is persistable state and must survive a + # save/load round-trip. + if use_moh: + stmts += [ + f"{idx_src}.indexer_moh_bias -> {idx_tgt}.indexer_moh_bias", + ] # --- MoE Gate --- stmts += [ @@ -949,6 +1018,15 @@ def _gen_inv_aoa_config(cls, config: DeepseekV4Config): f"{idx_src}.linear_weights_proj.weight^T -> {idx_tgt}.weights_proj.weight", f"{idx_src}.linear_wq_b.weight^T -> {idx_tgt}.wq_b.weight", ] + # V4_INDEXER_MOH: persist the trained aux-loss-free bias back to + # HF. Forward path uses ``_ -> ...indexer_moh_bias`` to zero-init + # on load, so if we skip this side of the round-trip the trained + # load-balancing state is dropped on every save_pretrained and + # replaced with zeros on the next from_pretrained. + if use_moh: + stmts += [ + f"{idx_src}.indexer_moh_bias -> {idx_tgt}.indexer_moh_bias", + ] # --- MoE Gate --- stmts += [f"{src}.mlp.gate.weight -> {tgt}.ffn.gate.weight,dtype='float32'"] diff --git a/paddleformers/transformers/model_utils.py b/paddleformers/transformers/model_utils.py index 05b54f215c3..4c23a9cbbed 100644 --- a/paddleformers/transformers/model_utils.py +++ b/paddleformers/transformers/model_utils.py @@ -2906,7 +2906,21 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): "When using flex_checkpoint to load Hugging Face open-source weights, " "the model must implement the _gen_aoa_config function to provide checkpoint conversion rules." ) - aoa_config = cls._gen_aoa_config(config) + # V4_INDEXER_MOH round-trip: some models (e.g. DeepSeek-V4 with MoH) + # need to know which keys are physically present in the HF + # checkpoint so they can emit ``named -> named`` for persisted + # trainer state (e.g. ``indexer_moh_bias``) and fall back to the + # ``_ -> ...`` add primitive only for keys that are actually + # missing. Pass ``checkpoint_keys`` when the classmethod accepts + # it; older classmethods without the kwarg keep their old + # signature and behavior. + _aoa_checkpoint_keys = None + if sharded_metadata is not None and "all_checkpoint_keys" in sharded_metadata: + _aoa_checkpoint_keys = set(sharded_metadata["all_checkpoint_keys"]) + try: + aoa_config = cls._gen_aoa_config(config, checkpoint_keys=_aoa_checkpoint_keys) + except TypeError: + aoa_config = cls._gen_aoa_config(config) sharded_state_dict = model.sharded_state_dict() metadata_path = os.path.join(ckpt_path, FLEX_CKPT_AUTO_GENERATED_METADATA) diff --git a/tests/trainer/test_indexer_bias_adjust_callback.py b/tests/trainer/test_indexer_bias_adjust_callback.py new file mode 100644 index 00000000000..6c3d66a90c5 --- /dev/null +++ b/tests/trainer/test_indexer_bias_adjust_callback.py @@ -0,0 +1,237 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for IndexerBiasAdjustCallback (MoH load balancing).""" + +import types +import unittest +from unittest.mock import MagicMock, patch + +import paddle +from paddle import nn + +from paddleformers.trainer.trainer_callback import IndexerBiasAdjustCallback + +# --------------------------------------------------------------------------- +# Helpers: minimal fake model that carries the MoH buffers. +# --------------------------------------------------------------------------- + + +class _FakeLinear(nn.Layer): + """Minimal stand-in for a ``nn.Linear``-shaped submodule of CSAIndexer.""" + + def __init__(self, in_size=16, out_size=16): + super().__init__() + self.weight = self.create_parameter( + shape=[in_size, out_size], default_initializer=nn.initializer.Constant(0.0) + ) + + +class _FakeCSAIndexer(nn.Layer): + """Fakes the two MoH buffers and the representative trainable projection. + + The callback keys "is this indexer frozen?" off + ``linear_weights_proj.weight`` (the head-scoring projection), so that's + the exact attribute the fake must expose -- otherwise the freeze branch + is never actually taken in tests. + """ + + def __init__(self, n_heads=8): + super().__init__() + self.register_buffer( + "indexer_moh_bias", + paddle.zeros([n_heads], dtype="float32"), + persistable=True, + ) + self.register_buffer( + "local_tokens_per_indexer_moh", + paddle.zeros([n_heads], dtype="float32"), + persistable=False, + ) + # Must be named ``linear_weights_proj`` to match the production + # ``getattr(m, "linear_weights_proj", None)`` gate in + # ``IndexerBiasAdjustCallback``. + self.linear_weights_proj = _FakeLinear() + + +class _NoMoHLayer(nn.Layer): + """Layer with no MoH buffers -- the callback should skip it entirely.""" + + def __init__(self): + super().__init__() + self.linear = _FakeLinear() + + +class _FakeModel(nn.Layer): + """Model containing a mix of MoH-carrying indexers and unrelated layers.""" + + def __init__(self, n_indexers=2, n_heads=8, include_no_moh_layer=True): + super().__init__() + self.indexers = nn.LayerList([_FakeCSAIndexer(n_heads=n_heads) for _ in range(n_indexers)]) + if include_no_moh_layer: + self.other = _NoMoHLayer() + + +def _run_callback(callback, model, freeze_training=False, optimizer=None): + """Invoke ``on_optimizer_end`` with the arg shape the callback expects.""" + args = types.SimpleNamespace(freeze_training=freeze_training) + state = MagicMock() + control = MagicMock() + + # Bypass the distributed all_reduce path: ``fleet._hcg`` is absent in the + # test env, so the callback would fall through to ``dist.all_reduce``. + # Patch it into an identity op so the callback can run single-process. + with patch( + "paddleformers.trainer.trainer_callback.dist.all_reduce", + side_effect=lambda t, group=None: t, + ): + callback.on_optimizer_end(args, state, control, model=model, optimizer=optimizer) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestIndexerBiasAdjustCallbackNoMoHModel(unittest.TestCase): + """A model without MoH buffers is a silent no-op (no error, no writes).""" + + def test_no_moh_module_is_noop(self): + model = nn.LayerList([_NoMoHLayer(), _NoMoHLayer()]) + callback = IndexerBiasAdjustCallback(lr=0.001) + # Should return cleanly without raising -- no all_reduce call needed. + _run_callback(callback, model) + + +class TestIndexerBiasAdjustCallbackUpdatesBias(unittest.TestCase): + """Standard update path: bias shifts toward mean, counter is zeroed.""" + + def test_sign_based_update_and_counter_reset(self): + paddle.seed(0) + model = _FakeModel(n_indexers=1, n_heads=4, include_no_moh_layer=True) + indexer = model.indexers[0] + # Head 0 got the most tokens, head 3 the fewest. + indexer.local_tokens_per_indexer_moh.set_value(paddle.to_tensor([10.0, 5.0, 5.0, 0.0], dtype="float32")) + # Snapshot pre-update bias. + before = indexer.indexer_moh_bias.numpy().copy() + + lr = 0.01 + callback = IndexerBiasAdjustCallback(lr=lr) + _run_callback(callback, model) + + after = indexer.indexer_moh_bias.numpy() + # mean = 5.0; sign(mean - usage) = [-1, 0, 0, +1] -> update = [-lr, 0, 0, +lr] + expected_delta = [-lr, 0.0, 0.0, lr] + for i in range(4): + self.assertAlmostEqual(float(after[i] - before[i]), expected_delta[i], places=6) + # Counter must be zeroed for the next accumulation window. + self.assertEqual(float(indexer.local_tokens_per_indexer_moh.sum().item()), 0.0) + + def test_multiple_indexers_updated_independently(self): + model = _FakeModel(n_indexers=3, n_heads=4) + for i, indexer in enumerate(model.indexers): + # Give each indexer a distinct imbalance. + indexer.local_tokens_per_indexer_moh.set_value(paddle.to_tensor([i + 1.0, 0.0, 0.0, 0.0], dtype="float32")) + callback = IndexerBiasAdjustCallback(lr=0.001) + _run_callback(callback, model) + # Every indexer's head-0 (over-used) bias must have decreased and the + # others increased, independently per indexer. + for indexer in model.indexers: + bias = indexer.indexer_moh_bias.numpy() + self.assertLess(float(bias[0]), 0.0) + for k in (1, 2, 3): + self.assertGreater(float(bias[k]), 0.0) + + +class TestIndexerBiasAdjustCallbackFreezeTraining(unittest.TestCase): + """``freeze_training`` skips both the bias update and the counter reset.""" + + def test_freeze_training_skips_everything(self): + model = _FakeModel(n_indexers=1, n_heads=4) + indexer = model.indexers[0] + indexer.local_tokens_per_indexer_moh.set_value(paddle.to_tensor([10.0, 0.0, 0.0, 0.0], dtype="float32")) + before_bias = indexer.indexer_moh_bias.numpy().copy() + before_counter = indexer.local_tokens_per_indexer_moh.numpy().copy() + + callback = IndexerBiasAdjustCallback(lr=0.01) + _run_callback(callback, model, freeze_training=True) + + # Neither the bias nor the counter changes when freeze_training is on. + self.assertTrue((indexer.indexer_moh_bias.numpy() == before_bias).all()) + self.assertTrue((indexer.local_tokens_per_indexer_moh.numpy() == before_counter).all()) + + +class TestIndexerBiasAdjustCallbackFrozenIndexer(unittest.TestCase): + """A frozen indexer (linear_weights_proj.weight.stop_gradient) is skipped.""" + + def test_stop_gradient_skips_update_but_still_zeroes_counter(self): + model = _FakeModel(n_indexers=1, n_heads=4) + indexer = model.indexers[0] + # Freeze the *same* param the callback keys off of. If this ever + # drifts (e.g. someone renames the ref_param in production), the + # assertion below on ``bias == before_bias`` will fail loudly instead + # of silently exercising the not-frozen branch. + indexer.linear_weights_proj.weight.stop_gradient = True + indexer.local_tokens_per_indexer_moh.set_value(paddle.to_tensor([10.0, 0.0, 0.0, 0.0], dtype="float32")) + before_bias = indexer.indexer_moh_bias.numpy().copy() + + callback = IndexerBiasAdjustCallback(lr=0.01) + _run_callback(callback, model) + + # Bias untouched, but counter still reset (mirrors MoECorrectionBias + # callback: frozen layers should not accumulate stale usage counts). + self.assertTrue((indexer.indexer_moh_bias.numpy() == before_bias).all()) + self.assertEqual(float(indexer.local_tokens_per_indexer_moh.sum().item()), 0.0) + + def test_stop_gradient_on_non_ref_param_does_not_freeze(self): + """Sanity: freezing a *different* param must NOT trip the freeze branch. + + Guards against the reviewer's exact complaint -- if the test freezes + something other than ``linear_weights_proj.weight``, ``ref_param`` is + untouched and the callback should proceed with the normal update. + This locks the ref_param contract from the other side. + """ + model = _FakeModel(n_indexers=1, n_heads=4) + indexer = model.indexers[0] + # Add an unrelated sibling projection and freeze *it*; the callback + # must NOT interpret that as "indexer is frozen". + indexer.some_other_proj = _FakeLinear() + indexer.some_other_proj.weight.stop_gradient = True + indexer.local_tokens_per_indexer_moh.set_value(paddle.to_tensor([10.0, 0.0, 0.0, 0.0], dtype="float32")) + before_bias = indexer.indexer_moh_bias.numpy().copy() + + callback = IndexerBiasAdjustCallback(lr=0.01) + _run_callback(callback, model) + + # The bias MUST have moved -- freezing an unrelated param is not a + # freeze signal for the callback. + self.assertFalse((indexer.indexer_moh_bias.numpy() == before_bias).all()) + + +class TestIndexerBiasAdjustCallbackExports(unittest.TestCase): + """The callback must be importable via the top-level trainer package.""" + + def test_top_level_import(self): + from paddleformers.trainer import IndexerBiasAdjustCallback as cb # noqa: F401 + + self.assertIs(cb, IndexerBiasAdjustCallback) + + def test_all_export(self): + from paddleformers.trainer.trainer_callback import __all__ + + self.assertIn("IndexerBiasAdjustCallback", __all__) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/transformers/deepseek_v4/__init__.py b/tests/transformers/deepseek_v4/__init__.py new file mode 100644 index 00000000000..290f972cf31 --- /dev/null +++ b/tests/transformers/deepseek_v4/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/transformers/deepseek_v4/test_deepseek_v4_aoa_moh_bias.py b/tests/transformers/deepseek_v4/test_deepseek_v4_aoa_moh_bias.py new file mode 100644 index 00000000000..df496641adc --- /dev/null +++ b/tests/transformers/deepseek_v4/test_deepseek_v4_aoa_moh_bias.py @@ -0,0 +1,403 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression tests for ``indexer_moh_bias`` save/load round-trip. + +``indexer_moh_bias`` is a *persistable* buffer that the aux-loss-free callback +mutates every optimizer step. The HF↔Fleet AOA config must therefore export it +in *both* directions: + + * HF -> Fleet (``_gen_aoa_config``): zero-init on load, since a fresh HF + checkpoint doesn't carry it. Uses the ``_ -> ...indexer_moh_bias`` add + primitive. + + * Fleet -> HF (``_gen_inv_aoa_config``): persist the trained bias back into + the HF checkpoint on ``save_pretrained``. Without this side of the pair, + every save/load round-trip resets the load-balancing state to zero and + silently loses training progress. + +These tests parse the AOA statement lists produced by both classmethods and +assert the symmetry directly, so any future refactor that drops one side +fails here. +""" + +import re +import unittest + +from paddleformers.transformers.deepseek_v4.configuration import DeepseekV4Config +from paddleformers.transformers.deepseek_v4.modeling import DeepseekV4PreTrainedModel + +# --------------------------------------------------------------------------- +# Config factory tuned to build both a decoder-side and an MTP-side CSAIndexer, +# so both indexer branches of the AOA config are exercised. +# --------------------------------------------------------------------------- + + +def _moh_config(**overrides): + """A minimal DSv4 config with MoH ON and both branches populated. + + * ``csa_compress_ratios[0] = 4`` -> layer 0 has a CSAIndexer (decoder branch). + * ``mtp_num_layers = 1`` and the MTP slot's compress ratio is ``4`` + -> the MTP branch also has a CSAIndexer. + """ + kwargs = dict( + num_hidden_layers=1, + n_routed_experts=2, + # csa_compress_ratios has length = num_hidden_layers + mtp_num_layers. + # index 0 is the decoder layer; the tail is consumed by MTP indexer + # branch (only compress_ratio > 0 and <= 4 builds an indexer). + csa_compress_ratios=[4, 4], + csa_dense_mode=False, + mtp_num_layers=1, + use_moh=True, + num_activated_heads=8, + dsa_index_n_heads=64, + # Keep expert count trivial so per-expert AOA lines don't drown out + # the indexer lines we're asserting on. + moe_n_hash_layers=0, + ) + kwargs.update(overrides) + return DeepseekV4Config(**kwargs) + + +def _find_indexer_lines(stmts, direction): + """Return the subset of AOA statements that touch ``indexer_moh_bias``. + + ``direction`` is 'fwd' (HF -> Fleet, expect ``_`` on the LHS) or 'inv' + (Fleet -> HF, expect the Fleet name on the LHS and HF on the RHS). + """ + out = [] + for s in stmts: + if "indexer_moh_bias" not in s: + continue + out.append(s.strip()) + return out + + +def _aoa_stmts(config, direction, checkpoint_keys=None): + """Extract the statement list from the ``_gen_[inv_]aoa_config`` return. + + Both methods return ``{"aoa_statements": [...]}`` (the dict form is what + ``PaddleFormers`` hands to ``AoAExecutor``); we only care about the flat + statement strings for these regressions. + + ``checkpoint_keys`` is forwarded to ``_gen_aoa_config`` (fwd only) so the + tests can exercise the "HF checkpoint carries the trained bias" branch; + ``None`` preserves the historical zero-init fallback. Ignored for the + inverse direction, which has no dependency on the loaded key set. + """ + fn = ( + DeepseekV4PreTrainedModel._gen_aoa_config + if direction == "fwd" + else DeepseekV4PreTrainedModel._gen_inv_aoa_config + ) + if direction == "fwd": + try: + out = fn(config, checkpoint_keys=checkpoint_keys) + except TypeError: + # Historical classmethod without the kwarg -- covered by the + # ``None`` case anyway; re-raise if a real key set was passed + # since that means the code hasn't picked up the kwarg yet. + if checkpoint_keys is not None: + raise + out = fn(config) + else: + out = fn(config) + if isinstance(out, dict): + return out["aoa_statements"] + # Historical form: a flat list. + return list(out) + + +class TestIndexerMoHBiasRoundTrip(unittest.TestCase): + """Both AOA directions must carry ``indexer_moh_bias``.""" + + def test_forward_aoa_zero_inits_bias(self): + """HF -> Fleet: ``_ -> ....indexer_moh_bias`` on both decoder & MTP. + + This is the *fresh HF release* / legacy caller path (no + ``checkpoint_keys`` info), so the forward AOA must still fall back + to the add primitive so the buffer is zero-initialized. + """ + cfg = _moh_config() + stmts = _aoa_stmts(cfg, "fwd") # checkpoint_keys=None (legacy path) + bias_lines = _find_indexer_lines(stmts, "fwd") + # One line for the decoder indexer, one for the MTP indexer. + self.assertEqual( + len(bias_lines), + 2, + f"expected 2 indexer_moh_bias entries in HF->Fleet AOA, got {bias_lines}", + ) + for line in bias_lines: + # LHS must be the add-primitive '_'. + self.assertRegex(line, r"^_\s*->\s*.*indexer_moh_bias\b") + + def test_forward_aoa_loads_bias_when_checkpoint_carries_it(self): + """HF -> Fleet: ``hf_key -> fleet_key`` when the HF checkpoint has the bias. + + This is the *round-trip load* path: the previous ``save_pretrained`` + wrote ``layers.*.attn.indexer.indexer_moh_bias`` via + ``_gen_inv_aoa_config``, so the next ``from_pretrained`` must load + that trained state instead of overwriting it with zeros. Regression + guard for the P1 the reviewer flagged. + """ + cfg = _moh_config() + + # Enumerate the HF-side prefixes the inverse config exports (this is + # the exact key set the round-tripped checkpoint will contain). + inv_stmts = _aoa_stmts(cfg, "inv") + inv_re = re.compile(r"^[\w\.]+\.indexer_moh_bias\s*->\s*([\w\.]+\.indexer_moh_bias)\b") + hf_bias_keys = set() + for s in inv_stmts: + m = inv_re.match(s.strip()) + if m: + hf_bias_keys.add(m.group(1)) + self.assertGreater( + len(hf_bias_keys), + 0, + "test setup: inverse AOA must export at least one indexer_moh_bias key", + ) + + stmts = _aoa_stmts(cfg, "fwd", checkpoint_keys=hf_bias_keys) + bias_lines = _find_indexer_lines(stmts, "fwd") + self.assertEqual( + len(bias_lines), + 2, + f"expected 2 indexer_moh_bias entries in HF->Fleet AOA, got {bias_lines}", + ) + pattern = re.compile(r"^([\w\.]+)\.indexer_moh_bias\s*->\s*([\w\.]+)\.indexer_moh_bias\b") + for line in bias_lines: + m = pattern.match(line) + self.assertIsNotNone( + m, + f"forward AOA with checkpoint_keys must be a named->named " f"mapping (not '_ -> ...'), got: {line!r}", + ) + hf_prefix, fleet_prefix = m.group(1), m.group(2) + # HF side (LHS): ``layers.*.attn.indexer`` (decoder or MTP). + self.assertTrue( + hf_prefix.endswith("attn.indexer"), + f"unexpected HF LHS for indexer_moh_bias: {hf_prefix!r}", + ) + # Fleet side (RHS): ``...self_attn.core_attention.indexer``. + self.assertIn("self_attn.core_attention.indexer", fleet_prefix) + # The '_' add primitive must NOT appear when the key is present. + self.assertFalse( + line.strip().startswith("_"), + f"AOA still zero-inits a bias that IS in the checkpoint: {line!r}", + ) + + def test_forward_aoa_mixed_checkpoint(self): + """Partial round-trip: only some indexer sites carry the trained bias. + + E.g. a checkpoint saved by an older Fleet where only the decoder + branch had ``indexer_moh_bias`` -- the MTP branch is still fresh. + The forward AOA must emit named->named for the present key AND + ``_ -> ...`` for the missing one, not one rule for both. + """ + cfg = _moh_config() + inv_stmts = _aoa_stmts(cfg, "inv") + inv_re = re.compile(r"^[\w\.]+\.indexer_moh_bias\s*->\s*([\w\.]+\.indexer_moh_bias)\b") + all_hf_keys = [inv_re.match(s.strip()).group(1) for s in inv_stmts if inv_re.match(s.strip())] + self.assertEqual(len(all_hf_keys), 2, "test setup: need exactly 2 HF bias keys") + # Keep only the decoder-side key (the one that does NOT contain + # ``transformer_layer`` on the Fleet side -- but we're keying by HF + # names here, so filter by MTP prefix instead). + mtp_key = next((k for k in all_hf_keys if "mtp" in k or k.startswith("mtp")), None) + # Fallback: HF-side MTP layers live at ``layers.{num_decoder+i}.attn.indexer`` + # in this codebase, so treat the second key as MTP if the first isn't. + if mtp_key is None: + mtp_key = all_hf_keys[1] + decoder_key = next(k for k in all_hf_keys if k != mtp_key) + + # Only the decoder-side bias is present in this "partial" checkpoint. + stmts = _aoa_stmts(cfg, "fwd", checkpoint_keys={decoder_key}) + bias_lines = _find_indexer_lines(stmts, "fwd") + self.assertEqual(len(bias_lines), 2, f"got {bias_lines}") + + named_lines = [ln for ln in bias_lines if not ln.strip().startswith("_")] + add_lines = [ln for ln in bias_lines if ln.strip().startswith("_")] + self.assertEqual( + len(named_lines), + 1, + f"exactly one named->named line expected for the present key, got {named_lines}", + ) + self.assertEqual( + len(add_lines), + 1, + f"exactly one '_ -> ...' line expected for the missing key, got {add_lines}", + ) + # The named line must use the decoder key on the LHS. + self.assertTrue( + named_lines[0].strip().startswith(decoder_key), + f"named line should route from {decoder_key!r}, got: {named_lines[0]!r}", + ) + + def test_inverse_aoa_persists_bias(self): + """Fleet -> HF: named -> named mapping on both decoder & MTP. + + Regression guard against the bug where ``indexer_moh_bias`` was + missing from ``_gen_inv_aoa_config``, so every ``save_pretrained`` + silently dropped the trained aux-loss-free bias. + """ + cfg = _moh_config() + stmts = _aoa_stmts(cfg, "inv") + bias_lines = _find_indexer_lines(stmts, "inv") + self.assertEqual( + len(bias_lines), + 2, + f"expected 2 indexer_moh_bias entries in Fleet->HF AOA, got {bias_lines}", + ) + pattern = re.compile(r"^([\w\.]+)\.indexer_moh_bias\s*->\s*([\w\.]+)\.indexer_moh_bias\b") + for line in bias_lines: + m = pattern.match(line) + self.assertIsNotNone( + m, + f"inverse AOA line must be a named->named mapping, got: {line!r}", + ) + fleet_prefix, hf_prefix = m.group(1), m.group(2) + # Fleet side: ``...self_attn.core_attention.indexer``. + self.assertIn("self_attn.core_attention.indexer", fleet_prefix) + # HF side: ``...attn.indexer`` (decoder or MTP), NOT the Fleet form. + self.assertTrue( + hf_prefix.endswith("attn.indexer"), + f"unexpected HF prefix for indexer_moh_bias: {hf_prefix!r}", + ) + self.assertNotIn("core_attention", hf_prefix) + + def test_round_trip_pairs_line_up(self): + """Every HF->Fleet target must have a matching Fleet->HF source. + + This is the direct assertion the reviewer asked for: for each + ``_ -> X.indexer_moh_bias`` in the forward config, there must be a + corresponding ``X.indexer_moh_bias -> _`` (any HF target) in the + inverse config with the *same* Fleet path. + """ + cfg = _moh_config() + fwd = _aoa_stmts(cfg, "fwd") + inv = _aoa_stmts(cfg, "inv") + + # Extract Fleet-side prefixes touched by the forward add-primitive. + fwd_prefixes = set() + fwd_re = re.compile(r"^_\s*->\s*([\w\.]+)\.indexer_moh_bias") + for s in fwd: + m = fwd_re.match(s.strip()) + if m: + fwd_prefixes.add(m.group(1)) + + # Extract Fleet-side prefixes on the LHS of the inverse mapping. + inv_prefixes = set() + inv_re = re.compile(r"^([\w\.]+)\.indexer_moh_bias\s*->") + for s in inv: + m = inv_re.match(s.strip()) + if m: + inv_prefixes.add(m.group(1)) + + self.assertEqual( + fwd_prefixes, + inv_prefixes, + "HF->Fleet and Fleet->HF must cover the same set of " + "indexer_moh_bias sites; drift here means save_pretrained will " + "silently drop the trained bias.", + ) + self.assertGreater( + len(fwd_prefixes), + 0, + "sanity: the test config should exercise at least one indexer", + ) + + def test_save_load_round_trip_end_to_end(self): + """Full save-then-load: no site is zero-init'd after a round-trip. + + Directly models the reviewer's P1: the previous fix only closed the + Fleet -> HF export side; a subsequent ``from_pretrained`` of *that* + checkpoint had no way to know it should load the exported bias + instead of zeroing it back out via the ``_ -> ...`` add primitive. + + This test simulates the save/load pair: + 1. Ask the inverse config which HF keys ``save_pretrained`` would + write for the bias (build ``checkpoint_keys``). + 2. Ask the forward config what it would emit given exactly that + key set (i.e. the ``from_pretrained`` right after save). + 3. Assert every emitted bias line for a site the checkpoint DID + persist is a named->named mapping -- NEVER ``_ -> ...``. + + Failure mode this catches: bias is written by save_pretrained but + the load path still uses the add primitive, so ``from_pretrained`` + of a fresh save silently resets the trained load-balancing state. + """ + cfg = _moh_config() + inv_stmts = _aoa_stmts(cfg, "inv") + # Set of HF-side keys that save_pretrained will actually persist. + inv_re = re.compile(r"^[\w\.]+\.indexer_moh_bias\s*->\s*([\w\.]+\.indexer_moh_bias)\b") + persisted_hf_keys = set() + for s in inv_stmts: + m = inv_re.match(s.strip()) + if m: + persisted_hf_keys.add(m.group(1)) + self.assertGreater(len(persisted_hf_keys), 0) + + # Simulate ``from_pretrained`` right after save -- pass those keys in. + fwd_stmts = _aoa_stmts(cfg, "fwd", checkpoint_keys=persisted_hf_keys) + fwd_bias_lines = _find_indexer_lines(fwd_stmts, "fwd") + self.assertEqual( + len(fwd_bias_lines), + len(persisted_hf_keys), + f"expected one forward bias line per persisted site, " + f"got {fwd_bias_lines} for keys {persisted_hf_keys}", + ) + + # None of them may be the zero-init add primitive. + offenders = [ln for ln in fwd_bias_lines if ln.strip().startswith("_")] + self.assertEqual( + offenders, + [], + "round-trip broken: save_pretrained persisted these HF keys " + f"{persisted_hf_keys}, but the forward AOA still zero-inits at " + f"least one of them: {offenders}. This is the exact regression " + "the reviewer flagged -- trained aux-loss-free bias is lost on " + "the next load.", + ) + + # Also verify every LHS is actually one of the persisted HF keys + # (not some fabricated name that doesn't line up with the save side). + pattern = re.compile(r"^([\w\.]+\.indexer_moh_bias)\s*->") + emitted_lhs = set() + for line in fwd_bias_lines: + m = pattern.match(line.strip()) + self.assertIsNotNone(m, f"malformed forward bias line: {line!r}") + emitted_lhs.add(m.group(1)) + self.assertEqual( + emitted_lhs, + persisted_hf_keys, + "forward AOA reads a different HF key set than the inverse " + "AOA writes; save/load will silently mismatch.", + ) + + +class TestIndexerMoHBiasOnlyWhenEnabled(unittest.TestCase): + """No ``indexer_moh_bias`` lines when ``use_moh=False`` (both directions).""" + + def test_forward_no_bias_when_moh_disabled(self): + cfg = _moh_config(use_moh=False, num_activated_heads=None) + stmts = _aoa_stmts(cfg, "fwd") + self.assertEqual(_find_indexer_lines(stmts, "fwd"), []) + + def test_inverse_no_bias_when_moh_disabled(self): + cfg = _moh_config(use_moh=False, num_activated_heads=None) + stmts = _aoa_stmts(cfg, "inv") + self.assertEqual(_find_indexer_lines(stmts, "inv"), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/transformers/deepseek_v4/test_deepseek_v4_config_moh.py b/tests/transformers/deepseek_v4/test_deepseek_v4_config_moh.py new file mode 100644 index 00000000000..43302e749c8 --- /dev/null +++ b/tests/transformers/deepseek_v4/test_deepseek_v4_config_moh.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for DeepseekV4Config MoH (use_moh / num_activated_heads) plumbing.""" + +import unittest + +from paddleformers.transformers.deepseek_v4.configuration import DeepseekV4Config + + +class TestDeepseekV4ConfigMoHDefaults(unittest.TestCase): + """MoH fields default to off / None.""" + + def test_use_moh_defaults_false(self): + config = DeepseekV4Config() + self.assertFalse(config.use_moh) + + def test_num_activated_heads_defaults_none(self): + config = DeepseekV4Config() + self.assertIsNone(config.num_activated_heads) + + +class TestDeepseekV4ConfigMoHAccepted(unittest.TestCase): + """When both fields are given they are stored and accessible.""" + + def test_fields_stored(self): + config = DeepseekV4Config(use_moh=True, num_activated_heads=8) + self.assertTrue(config.use_moh) + self.assertEqual(config.num_activated_heads, 8) + + def test_num_activated_heads_equal_dsa_index_n_heads(self): + config = DeepseekV4Config(use_moh=True, num_activated_heads=64, dsa_index_n_heads=64) + self.assertEqual(config.num_activated_heads, 64) + self.assertEqual(config.dsa_index_n_heads, 64) + + +class TestDeepseekV4ConfigMoHSerialization(unittest.TestCase): + """Fields survive round-trip through to_dict.""" + + def test_to_dict_includes_moh_fields(self): + config = DeepseekV4Config(use_moh=True, num_activated_heads=16) + d = config.to_dict() + self.assertTrue(d["use_moh"]) + self.assertEqual(d["num_activated_heads"], 16) + + def test_default_to_dict_includes_moh_fields(self): + config = DeepseekV4Config() + d = config.to_dict() + self.assertFalse(d["use_moh"]) + self.assertIsNone(d["num_activated_heads"]) + + +class TestDeepseekV4ConfigMoHHFMapping(unittest.TestCase): + """``_HF_TO_FLEET_FIELD_MAP`` does NOT remap these fields (same name).""" + + def test_no_hf_remapping_for_moh(self): + # use_moh / num_activated_heads should NOT be in the remap dict since + # the HF and Fleet names are identical. If they were remapped, their + # default would be overwritten by the pop() in __init__. + mapping = DeepseekV4Config._HF_TO_FLEET_FIELD_MAP + self.assertNotIn("use_moh", mapping) + self.assertNotIn("num_activated_heads", mapping) + + +if __name__ == "__main__": + unittest.main()