diff --git a/tests/models/test_moe_zloss_per_token_loss.py b/tests/models/test_moe_zloss_per_token_loss.py new file mode 100644 index 00000000000..495dedc936d --- /dev/null +++ b/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) diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index cbfd6d9ea6d..b5c9be0eeed 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/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