[megatron] fix: return 3-tuple under calculate_per_token_loss to fix MoE aux/z-loss grad blowup at CP>1

tests/models/test_moe_zloss_per_token_loss.py

@@ -0,0 +1,156 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+"""MoE router z-loss must not depend on ``calculate_per_token_loss``.
+
+Under context parallelism the Megatron bridge enables ``calculate_per_token_loss``,
+which makes ``apply_z_loss`` scale the z-loss by the local token count, expecting
+``finalize_model_grads`` to divide it back out by ``num_tokens``. verl's loss path
+returns a 2-tuple ``(loss, output)``, so the schedule leaves ``num_tokens=0`` and
+finalize skips that division -- leaving the per-token factor (~thousands)
+uncancelled and blowing up the gradient. The flag must not change the gradient.
+"""
+
+import os
+
+os.environ["NCCL_DEBUG"] = "WARN"
+
+from functools import partial
+
+import numpy as np
+import ray
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen3MoeConfig
+
+from verl import DataProto
+from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
+from verl.utils import tensordict_utils as tu
+from verl.utils.model import compute_position_id_with_mask, create_random_mask
+from verl.workers.config import ActorConfig, HFModelConfig, McoreEngineConfig, McoreOptimizerConfig
+from verl.workers.engine_workers import TrainingWorker, TrainingWorkerConfig
+from verl.workers.utils.losses import ppo_loss
+from verl.workers.utils.padding import left_right_2_no_padding
+
+Z_LOSS_COEFF = 1e-3
+# A small local model dir to borrow a tokenizer from (same convention as test_engine.py).
+TOKENIZER_REF = os.path.expanduser("~/models/Qwen/Qwen2.5-0.5B")
+
+
+def _create_tiny_moe_model(path):
+    """Save a tiny Qwen3-MoE checkpoint so the router z-loss path is exercised. The
+    engine's forward needs a tokenizer (pad_token_id), so we borrow one and size the
+    vocab to match. Returns the vocab size."""
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_REF)
+    config = Qwen3MoeConfig(
+        vocab_size=len(tokenizer),
+        hidden_size=64,
+        intermediate_size=128,
+        moe_intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_experts=8,
+        num_experts_per_tok=2,
+        decoder_sparse_step=1,
+        norm_topk_prob=True,
+        max_position_embeddings=512,
+        tie_word_embeddings=False,
+    )
+    AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16).save_pretrained(path)
+    tokenizer.save_pretrained(path)
+    return len(tokenizer)
+
+
+def _build_data(vocab_size):
+    batch_size, seqlen = 4, 32
+    response_length = seqlen // 2
+    torch.manual_seed(1)
+    np.random.seed(1)
+    input_ids = torch.randint(0, vocab_size, (batch_size, seqlen))
+    attention_mask = create_random_mask(
+        input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6
+    )
+    position_ids = compute_position_id_with_mask(attention_mask)
+    responses = input_ids[:, response_length:]
+    response_mask = attention_mask[:, response_length:]
+    data = DataProto.from_single_dict(
+        {
+            "input_ids": input_ids,
+            "prompts": input_ids[:, :response_length],
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "responses": responses,
+            "response_mask": response_mask,
+            "old_log_probs": torch.randn_like(responses, dtype=torch.float32),
+            "advantages": torch.randn_like(responses, dtype=torch.float32),
+            "ref_log_prob": torch.randn_like(responses, dtype=torch.float32),
+        },
+        meta_info={"temperature": 1.0, "global_token_num": torch.sum(attention_mask, dim=-1).tolist()},
+    )
+    data_td = left_right_2_no_padding(data.to_tensordict())
+    tu.assign_non_tensor(data_td, global_batch_size=data_td.shape[0])
+    return data_td
+
+
+def _grad_norm(model_path, calculate_per_token_loss, data_td):
+    """Run one PPO update through a Megatron MoE engine and return the grad norm.
+    A fresh Ray session is used per call so the two single-GPU worker groups don't
+    contend for the device."""
+    config = TrainingWorkerConfig(
+        model_type="language_model",
+        model_config=HFModelConfig(path=model_path, use_remove_padding=True),
+        engine_config=McoreEngineConfig(
+            forward_only=False,
+            use_mbridge=True,
+            vanilla_mbridge=False,  # NVIDIA Megatron-Bridge (production path; ISEEKYAN mbridge is deprecated)
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=1,
+            use_dynamic_bsz=True,
+            use_remove_padding=True,
+            max_token_len_per_gpu=2048,
+            override_transformer_config={
+                "moe_z_loss_coeff": Z_LOSS_COEFF,
+                "calculate_per_token_loss": calculate_per_token_loss,
+            },
+        ),
+        optimizer_config=McoreOptimizerConfig(lr_decay_steps=10),
+        checkpoint_config=None,
+    )
+    ray.init()
+    try:
+        ray_cls = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config)
+        wg = RayWorkerGroup(resource_pool=RayResourcePool(process_on_nodes=[1]), ray_cls_with_init=ray_cls)
+        wg.reset()
+        actor_config = ActorConfig(strategy="megatron", rollout_n=1, ppo_micro_batch_size_per_gpu=-1)
+        wg.set_loss_fn(partial(ppo_loss, config=actor_config))
+        metrics = tu.get(wg.train_batch(data_td).get(), "metrics")
+        return float(metrics["grad_norm"])
+    finally:
+        ray.shutdown()
+
+
+def test_moe_zloss_invariant_to_per_token_loss(tmp_path):
+    model_path = str(tmp_path / "tiny_moe")
+    vocab_size = _create_tiny_moe_model(model_path)
+    data_td = _build_data(vocab_size)
+
+    grad_norm_default = _grad_norm(model_path, calculate_per_token_loss=False, data_td=data_td)
+    grad_norm_per_token = _grad_norm(model_path, calculate_per_token_loss=True, data_td=data_td)
+
+    # With the bug, calculate_per_token_loss=True leaves the ~num_tokens per-token
+    # factor uncancelled (finalize num_tokens=0), blowing grad_norm up by orders of
+    # magnitude. The z-loss gradient must not depend on the flag.
+    torch.testing.assert_close(grad_norm_per_token, grad_norm_default, rtol=5e-2, atol=1e-3)

verl/workers/engine/megatron/transformer_impl.py

@@ -644,6 +644,18 @@ def load_checkpoint(
         if self._is_offload_optimizer:
             offload_megatron_optimizer(self.optimizer)
 
+    def _routed_num_tokens(self, data: TensorDict) -> torch.Tensor:
+        """Real (unpadded) tokens fed to the MoE router: attention_mask in the padded RL
+        path, else the packed input_ids count in the no-padding SFT path. Not loss_mask,
+        which counts response tokens only and would under-normalize the router loss."""
+        attention_mask = data.get("attention_mask", None)
+        if attention_mask is not None:
+            return attention_mask.sum()
+        input_ids = data["input_ids"]
+        if input_ids.is_nested:
+            return input_ids.offsets()[-1]
+        return torch.tensor(input_ids.numel(), device=input_ids.device)
+
     def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:
         self._distillation_use_topk_active = tu.get_non_tensor_data(data, key="distillation_use_topk", default=False)
         tu.assign_non_tensor(data, sp_size=self.engine_config.context_parallel_size)
@@ -656,6 +668,16 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw
         tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())
         tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())
 
+        # Global routed-token count for the per-token-loss regime (consumed in
+        # postprocess_micro_batch_func). Real tokens are CP-replicated, so a single
+        # all-reduce over the DP group gives the global value.
+        if self.tf_config is not None and self.tf_config.calculate_per_token_loss:
+            routed_num_tokens = self._routed_num_tokens(data).to(get_device_id())
+            torch.distributed.all_reduce(
+                routed_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group()
+            )
+            tu.assign_non_tensor(data, routed_num_tokens=routed_num_tokens.item())
+
         # BSHD path only: pad every micro-batch to the mini-batch's global max seq_len so the
         # padded `s_q` is shared -> cuDNN plan built once per shape. Raw (unaligned)
         # max; TP/CP/FP8 alignment is applied inside preprocess_bshd_engine.
@@ -1081,29 +1103,6 @@ def postprocess_micro_batch_func(
             # scale loss by num_micro_batch because megatron will scale loss
             # by n_micro_batch inside pp schedule
             scaled_loss = loss * data["num_micro_batch"]
-            if self.tf_config.calculate_per_token_loss and not forward_only:
-                # verl losses are already normalized over the global DP batch. MCore's
-                # legacy two-item loss callback still multiplies by CP, while per-token
-                # mode disables DDP's usual 1/(DP*CP) gradient pre-scaling. Compensate
-                # for that combined reduction until the callback returns MCore's
-                # three-item (loss sum, token count, metrics) contract.
-                num_moe_experts = getattr(self.tf_config, "num_moe_experts", None)
-                has_moe_aux_loss = bool(num_moe_experts) and (
-                    bool(getattr(self.tf_config, "moe_aux_loss_coeff", 0.0))
-                    or bool(getattr(self.tf_config, "moe_z_loss_coeff", None))
-                )
-                has_auxiliary_loss = (
-                    has_moe_aux_loss
-                    or bool(getattr(self.tf_config, "mtp_num_layers", None))
-                    or getattr(self.tf_config, "experimental_attention_variant_loss_scale_func", None) is not None
-                    or getattr(self.tf_config, "experimental_attention_variant", None) == "dsa"
-                )
-                # Auxiliary-loss autograd scalers bypass scaled_loss, and dynamic CP
-                # uses per-microbatch groups. Preserve their existing behavior until
-                # verl adopts the three-item callback that normalizes every gradient.
-                if not self.engine_config.dynamic_context_parallel and not has_auxiliary_loss:
-                    dp_cp_world_size = mpu.get_data_parallel_world_size(with_context_parallel=True)
-                    scaled_loss /= dp_cp_world_size
         else:
             assert forward_only, "forward_only must be True when loss_function is None"
             loss = torch.tensor(1.0, device=device)
@@ -1126,6 +1125,54 @@ def postprocess_micro_batch_func(
             "metrics": metrics,
         }
 
+        # calculate_per_token_loss=True (auto-enabled by Megatron-Bridge at CP>1) puts
+        # Megatron in its per-token regime: loss_func must return (loss_sum, num_tokens,
+        # output), and finalize_model_grads divides every gradient by the accumulated
+        # total_num_tokens. That division is what cancels the MoE router's pre-multiplication
+        # of the aux/z loss by num_tokens; a 2-tuple leaves total_num_tokens=0, so the factor
+        # is never cancelled (the ~1e4 grad_norm blow-up at CP>1).
+        if self.tf_config is not None and self.tf_config.calculate_per_token_loss and loss_function is not None:
+            # seq-mean-token-mean is the one incompatible agg mode: its per-sequence 1/n_s
+            # uses CP-local shard counts that diverge from the global normalization. The
+            # other modes compose correctly across CP shards.
+            if hasattr(loss_function, "keywords") and "config" in loss_function.keywords:
+                _agg_mode = getattr(loss_function.keywords["config"], "loss_agg_mode", None)
+                if _agg_mode == "seq-mean-token-mean":
+                    raise ValueError(
+                        "loss_agg_mode='seq-mean-token-mean' is incompatible with "
+                        "calculate_per_token_loss=True (auto-enabled by Megatron-Bridge "
+                        "under CP>1). The per-sequence inner division by n_s requires "
+                        "local-shard counts that diverge from global under CP. Use one "
+                        "of: 'token-mean', 'seq-mean-token-sum', 'seq-mean-token-sum-norm'."
+                    )
+            # verl never passes a router padding_mask, so the MoE router normalizes the
+            # aux/z loss by logits.shape[0]. THD packs padding out -> that equals the real
+            # token count; BSHD leaves it at B*S (padding-inclusive), while gradients are
+            # divided by the real token count -> a padding-ratio mis-normalization.
+            if not self.engine_config.use_remove_padding:
+                raise ValueError(
+                    "calculate_per_token_loss=True requires use_remove_padding=True. "
+                    "verl does not pass a padding_mask to the MoE router, so in BSHD it "
+                    "normalizes the aux/z loss by the padding-inclusive token count (B*S) "
+                    "while gradients are divided by the real token count. Use THD "
+                    "(use_remove_padding=True) or disable CP."
+                )
+            # finalize_model_grads all-reduces the returned token count over the DP*CP group
+            # and divides every gradient by it. Real tokens are CP-replicated across the CP
+            # ranks, so report the per-CP-rank share (/cp_size); otherwise that DP*CP sum
+            # over-counts by cp_size and every gradient comes out 1/cp_size too small.
+            cp_size = self.engine_config.context_parallel_size
+            local_num_tokens = (self._routed_num_tokens(data) // cp_size).to(torch.int)
+            # n_i is the global routed-token count (all-reduced in forward_backward_batch);
+            # scaling loss by the same value makes Sum(L_i)/Sum(n_i) recover the loss. Falls
+            # back to local counts when not plumbed (single-rank / tests).
+            routed_num_tokens = data["routed_num_tokens"] if "routed_num_tokens" in data.keys() else None
+            if routed_num_tokens is None:
+                routed_num_tokens = self._routed_num_tokens(data)
+            dp_size = data["dp_size"] if "dp_size" in data.keys() else 1
+            local_sum = loss * routed_num_tokens / dp_size
+            return local_sum, local_num_tokens, output
+
         # return loss and stats
         return scaled_loss, output