diff --git a/tests/unit_tests/cpu/test_aux_loss.py b/tests/unit_tests/cpu/test_aux_loss.py index b59684997b7..ab40e50b1e4 100644 --- a/tests/unit_tests/cpu/test_aux_loss.py +++ b/tests/unit_tests/cpu/test_aux_loss.py @@ -37,7 +37,10 @@ make_routed_experts_config, make_router_config, ) -from torchtitan.models.common.moe import MicrobatchWiseLoadBalanceLoss +from torchtitan.models.common.moe import ( + GlobalBatchLoadBalanceLoss, + MicrobatchWiseLoadBalanceLoss, +) _COEFF = 0.1 _METRIC_KEY = ("dp", "microbatch_wise_load_balance_loss") @@ -75,6 +78,23 @@ def _reference_loss( return (f_E * p_E).sum() * T * coeff +def _reference_global_loss( + scores_TE: torch.Tensor, + routing_map_TE: torch.Tensor, + global_routing_map_TE: torch.Tensor, + top_k: int, + *, + coeff: float = 1.0, +) -> torch.Tensor: + """Global-batch reference: ``f_i`` from every DP rank's tokens + (``global_routing_map_TE``), ``p_i`` summed over this rank's tokens.""" + E = scores_TE.size(-1) + counts_E = global_routing_map_TE.sum(dim=0).to(scores_TE.dtype) + f_E = counts_E * (E / (top_k * global_routing_map_TE.size(0))) + probs_TE = scores_TE / scores_TE.sum(dim=-1, keepdim=True) + return (f_E * probs_TE.sum(dim=0)).sum() * coeff + + def _routing_map(ids_TK: torch.Tensor, num_experts: int) -> torch.Tensor: """One-hot routing map for ``(T, K)`` expert ids.""" return torch.zeros(ids_TK.shape[0], num_experts, dtype=torch.bool).scatter_( @@ -273,6 +293,62 @@ def test_config_builds_the_loss(self): type(moe_cfg.router.aux_loss.build()), MicrobatchWiseLoadBalanceLoss ) + def test_config_selects_the_global_batch_loss(self): + kwargs = dict( + num_experts=4, + router=make_router_config( + dim=8, + num_experts=4, + gate_param_init={}, + score_func=Sigmoid.Config(), + ), + routed_experts=make_routed_experts_config( + dim=8, + hidden_dim=16, + num_experts=4, + top_k=1, + param_init={}, + ), + load_balance_coeff=None, + global_batch_aux_loss=True, + ) + moe_cfg = make_moe_config(**kwargs, aux_loss_coeff=_COEFF) + self.assertIs(type(moe_cfg.router.aux_loss.build()), GlobalBatchLoadBalanceLoss) + self.assertIsNone(moe_cfg.load_balance_coeff) + # Selecting the global-batch loss without a coefficient is an error, not a no-op. + with self.assertRaises(ValueError): + make_moe_config(**kwargs) + + def test_qwen3_30b_a3b_global_lbl_recipe(self): + """The opt-in recipe swaps every MoE layer to the global-batch loss.""" + from torchtitan_recipes.tests.models.qwen3 import ( + qwen3_30b_a3b, + qwen3_30b_a3b_global_lbl, + ) + + moes = [l.moe for l in qwen3_30b_a3b_global_lbl().model.layers] + self.assertEqual(len(moes), 48) + for moe in moes: + self.assertIsNone(moe.load_balance_coeff) + self.assertIs(type(moe.router.aux_loss), GlobalBatchLoadBalanceLoss.Config) + self.assertEqual(moe.router.aux_loss.coeff, 1e-3) + # The default recipe keeps the expert bias and no aux loss. + for layer in qwen3_30b_a3b().model.layers: + self.assertEqual(layer.moe.load_balance_coeff, 1e-3) + self.assertIsNone(layer.moe.router.aux_loss) + + def test_qwen3_registers_the_aux_loss_hook(self): + """Without the zero hook the metric reads 0 and instance_acc grows.""" + from torchtitan.models.qwen3.model import Qwen3Model + + with patch( + "torchtitan.models.common.aux_loss.register_aux_loss_zero_hook" + ) as aux_hook, patch( + "torchtitan.models.common.moe.register_moe_load_balancing_hook" + ): + Qwen3Model._register_optimizer_hooks("opt", ["part"], "ctx") + aux_hook.assert_called_once_with("opt", ["part"], "ctx") + class TestMicrobatchWiseLossSpmdTypes(DTensorTestBase): """8-rank spmd_types cases: dp2/cp2/tp2 on CPU float64 + gloo. @@ -343,7 +419,7 @@ def test_pp_reduction_sums_stages(self): self.assertAlmostEqual(metrics[f"{_METRIC_KEY[1]}/mean"], 2.0, places=6) _clear_aux_loss_registry() - def _run_reduction_case(self, *, use_typecheck: bool): + def _run_reduction_case(self, *, use_typecheck: bool, global_batch: bool = False): """Compare one distributed layout with the per-DP-rank reference.""" parallelism_context, dense_mesh = self._setup_mesh() from torchtitan.distributed.spmd_types import set_current_spmd_mesh @@ -383,9 +459,12 @@ def _run_reduction_case(self, *, use_typecheck: bool): spmd.assert_type(local_ids_TK, placement) spmd.assert_type(local_map, placement) - loss = MicrobatchWiseLoadBalanceLoss( - MicrobatchWiseLoadBalanceLoss.Config(coeff=_COEFF) + loss_cls = ( + GlobalBatchLoadBalanceLoss + if global_batch + else MicrobatchWiseLoadBalanceLoss ) + loss = loss_cls(loss_cls.Config(coeff=_COEFF)) local_scores.requires_grad_(True) carrier_TK = local_scores.gather(dim=-1, index=local_ids_TK) out_TK = loss(local_scores, local_map, carrier=carrier_TK) @@ -398,12 +477,21 @@ def _run_reduction_case(self, *, use_typecheck: bool): dp_scores = global_scores_TE[dp_start : dp_start + t_dp] dp_ids_TK = global_ids_TK[dp_start : dp_start + t_dp] dp_map = _routing_map(dp_ids_TK, E) + global_map = _routing_map(global_ids_TK, E) + + def ref_fn(scores, coeff=1.0): + if global_batch: + return _reference_global_loss( + scores, dp_map, global_map, K, coeff=coeff + ) + return _reference_loss(scores, dp_map, K, coeff=coeff) + # The metric accumulator is float32 (the reference float64). - ref_raw = _reference_loss(dp_scores, dp_map, K).item() + ref_raw = ref_fn(dp_scores).item() self.assertAlmostEqual(loss.instance_acc.item(), ref_raw, places=4) ref_scores = dp_scores.detach().clone().requires_grad_(True) - ref_aux = _reference_loss(ref_scores, dp_map, K, coeff=_COEFF) + ref_aux = ref_fn(ref_scores, coeff=_COEFF) (ref_aux + (ref_scores * dp_map).sum()).backward() ref_local_grad = ref_scores.grad[shard * t_blk : (shard + 1) * t_blk] self.assertLess( @@ -418,13 +506,27 @@ def _run_reduction_case(self, *, use_typecheck: bool): for stream in range(dp): stream_scores = global_scores_TE[stream * t_dp : (stream + 1) * t_dp] stream_ids_TK = global_ids_TK[stream * t_dp : (stream + 1) * t_dp] - ref_total += _reference_loss( - stream_scores, _routing_map(stream_ids_TK, E), K + stream_map = _routing_map(stream_ids_TK, E) + ref_total += ( + _reference_global_loss( + stream_scores, stream_map, _routing_map(global_ids_TK, E), K + ) + if global_batch + else _reference_loss(stream_scores, stream_map, K) ).item() metrics = collect_aux_loss_metrics(parallelism_context) - self.assertAlmostEqual(metrics[f"{_METRIC_KEY[1]}/mean"], ref_total, places=4) + key = loss.metric_name + self.assertAlmostEqual(metrics[f"{key}/mean"], ref_total, places=4) _clear_aux_loss_registry() + @with_comms + def test_global_batch_reduction_matches_reference(self): + """``f_i`` comes from both DP ranks' tokens; each rank's gradient is + the global-count reference on its own tokens.""" + for use_typecheck in (True, False): + with self.subTest(use_typecheck=use_typecheck): + self._run_reduction_case(use_typecheck=use_typecheck, global_batch=True) + @with_comms def test_reduction_matches_reference(self): """Loss, gradient and collected metric match the per-DP-rank reference: diff --git a/torchtitan/models/common/__init__.py b/torchtitan/models/common/__init__.py index 88a1f3229f7..7c5d1692399 100644 --- a/torchtitan/models/common/__init__.py +++ b/torchtitan/models/common/__init__.py @@ -42,7 +42,7 @@ RowParallelLinear, SharedExpertRowParallelLinear, ) -from .moe import MicrobatchWiseLoadBalanceLoss, MoE +from .moe import GlobalBatchLoadBalanceLoss, MicrobatchWiseLoadBalanceLoss, MoE from .multimodal import MultimodalModel from .nn_modules import Conv1d, Conv2d, GELU, GroupNorm, Identity, LayerNorm, RMSNorm from .norm import GatedRMSNorm @@ -77,6 +77,7 @@ "LayerNorm", "Linear", "MoE", + "GlobalBatchLoadBalanceLoss", "MicrobatchWiseLoadBalanceLoss", "MultimodalModel", "RMSNorm", diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py index 60d35c1adf4..de1c3ed47fe 100644 --- a/torchtitan/models/common/config_utils.py +++ b/torchtitan/models/common/config_utils.py @@ -34,6 +34,7 @@ SharedExpertRowParallelLinear, ) from torchtitan.models.common.moe import ( + GlobalBatchLoadBalanceLoss, MicrobatchWiseLoadBalanceLoss, MoE, RoutedExperts, @@ -314,12 +315,24 @@ def make_moe_config( shared_experts: FeedForward.Config | None = None, load_balance_coeff: float | None = 1e-3, aux_loss_coeff: float | None = None, + global_batch_aux_loss: bool = False, ) -> MoE.Config: - """Build a fully-specified MoE.Config.""" + """Build a fully-specified MoE.Config. + + ``global_batch_aux_loss`` selects ``GlobalBatchLoadBalanceLoss`` (Qwen3) + instead of the per-microbatch loss; Qwen3's recipe also sets + ``load_balance_coeff=None`` (no expert bias). + """ + if global_batch_aux_loss and aux_loss_coeff is None: + raise ValueError("global_batch_aux_loss=True requires aux_loss_coeff") if aux_loss_coeff is not None: + loss_cls = ( + GlobalBatchLoadBalanceLoss + if global_batch_aux_loss + else MicrobatchWiseLoadBalanceLoss + ) router = dataclasses.replace( - router, - aux_loss=MicrobatchWiseLoadBalanceLoss.Config(coeff=aux_loss_coeff), + router, aux_loss=loss_cls.Config(coeff=aux_loss_coeff) ) return MoE.Config( num_experts=num_experts, diff --git a/torchtitan/models/common/moe.py b/torchtitan/models/common/moe.py index a39381a980d..bbdaf86a82d 100644 --- a/torchtitan/models/common/moe.py +++ b/torchtitan/models/common/moe.py @@ -580,8 +580,10 @@ def forward( # degree and follows any masking the router applies to the map. # The map is cast to float before the reduction: casting a Partial # tensor is non-linear and rejected by spmd_types. - counts_E = self._reduce_token_partials( - routing_map_TE.to(scores_TE.dtype).sum(dim=0), axes + counts_E = self._balancing_counts( + self._reduce_token_partials( + routing_map_TE.to(scores_TE.dtype).sum(dim=0), axes + ) ) f_E = F.normalize(counts_E, p=1, dim=0) * E @@ -597,6 +599,48 @@ def forward( loss = (f_E * p_E).sum() return self.inject(loss, carrier=carrier) + def _balancing_counts(self, counts_E: torch.Tensor) -> torch.Tensor: + """Expert counts that define ``f_i``: this forward's token stream.""" + return counts_E + + +class GlobalBatchLoadBalanceLoss(MicrobatchWiseLoadBalanceLoss): + """Global-batch MoE load-balance loss (Qiu et al. 2025, used by Qwen3). + + Same ``L_bal = sum_i f_i * p_i`` as ``MicrobatchWiseLoadBalanceLoss``, but + ``f_i`` comes from the counts of every data-parallel rank's tokens, not + one rank's. ``p_i`` stays this rank's score sum: the injected gradients + of all DP ranks add up (the step denominator is the dp-summed token + count), so the step optimizes ``sum_i f_i^global * p_i^global``. + + Balancing over the global batch instead of one microbatch lets experts + specialize on domains that are unevenly spread over sequences, which a + per-microbatch constraint penalizes (https://arxiv.org/abs/2501.11873). + + Limitation: under gradient accumulation each microbatch uses its own + (DP-global) counts; accumulating counts across microbatches is not + implemented. + """ + + @dataclass(kw_only=True, slots=True) + class Config(MicrobatchWiseLoadBalanceLoss.Config): + """Same fields as ``AuxLoss.Config``.""" + + def _balancing_counts(self, counts_E: torch.Tensor) -> torch.Tensor: + # The counts are one-hot sums with no gradient, so a plain all-reduce + # over DP (outside the typechecker: DP is local here) is enough. + group = spmd_mesh_group("dp") + if group is None: + return counts_E + local_type = spmd.get_local_type(counts_E) + with spmd.no_typecheck(): + global_counts_E = counts_E.detach().clone() + torch.distributed.all_reduce(global_counts_E, group=group) + # Same (Invariant) type as the per-rank counts; only the values differ. + if local_type is not None: + spmd.assert_type(global_counts_E, local_type) + return global_counts_E + class MoE(Module): """Mixture of Experts layer. diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py index aeac6d33f6c..e1830d8bff1 100644 --- a/torchtitan/models/qwen3/model.py +++ b/torchtitan/models/qwen3/model.py @@ -76,9 +76,11 @@ class Qwen3Model(Decoder): def _register_optimizer_hooks( cls, optimizers, model_parts, parallelism_context ) -> None: + from torchtitan.models.common.aux_loss import register_aux_loss_zero_hook from torchtitan.models.common.moe import register_moe_load_balancing_hook register_moe_load_balancing_hook(optimizers, model_parts, parallelism_context) + register_aux_loss_zero_hook(optimizers, model_parts, parallelism_context) """ Qwen3Model Module diff --git a/torchtitan_recipes/tests/models/qwen3.py b/torchtitan_recipes/tests/models/qwen3.py index 238bd0efcdb..dea0c8d17af 100644 --- a/torchtitan_recipes/tests/models/qwen3.py +++ b/torchtitan_recipes/tests/models/qwen3.py @@ -39,6 +39,7 @@ decoder_vocab_size, DEFAULT_DEBUG_MODEL_SEQ_LEN, ) +from torchtitan.models.common.moe import GlobalBatchLoadBalanceLoss from torchtitan.models.common.token_dispatcher import DeepEPTokenDispatcher from torchtitan.models.qwen3 import build_model_config @@ -410,6 +411,20 @@ def qwen3_30b_a3b(seq_len: int | None = None) -> Trainer.Config: ) +def qwen3_30b_a3b_global_lbl(seq_len: int | None = None) -> Trainer.Config: + """``qwen3_30b_a3b`` with Qwen3's global-batch load-balancing loss. + + Every MoE layer uses ``GlobalBatchLoadBalanceLoss`` (coeff 1e-3) and no + expert bias, as in Qwen3's recipe (https://arxiv.org/abs/2505.09388). + """ + config = qwen3_30b_a3b(seq_len=seq_len) + for layer in config.model.layers: + if layer.moe is not None: + layer.moe.load_balance_coeff = None + layer.moe.router.aux_loss = GlobalBatchLoadBalanceLoss.Config(coeff=1e-3) + return config + + def qwen3_32b(seq_len: int | None = None) -> Trainer.Config: model_config = build_model_config("32B", seq_len=seq_len) return Trainer.Config(