diff --git a/xtuner/v1/loss/aux_loss.py b/xtuner/v1/loss/aux_loss.py
index a6ac7179f..d19f59633 100644
--- a/xtuner/v1/loss/aux_loss.py
+++ b/xtuner/v1/loss/aux_loss.py
@@ -91,7 +91,6 @@ def accumulate(
         z_ctx: list[ZLossContext] | ZLossContext | None = None,
         num_tokens_local: int = 0,
         num_tokens_global: torch.Tensor | None = None,
-        world_size: int = 1,
     ) -> torch.Tensor:
         """Accumulate routing statistics for one layer and inject z-loss into
         the main graph.
@@ -113,7 +112,6 @@ def accumulate(
             num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across
                 ranks (int64 scalar). Pass ``None`` when ``z_loss_global_average`` is off or no
                 process group is initialized.
-            world_size (int): World size that produced ``num_tokens_global``.
 
         Returns:
             torch.Tensor: ``hidden_states`` augmented with the per-layer z-loss autograd hook.
@@ -139,7 +137,6 @@ def accumulate(
                 router_logits=selected_router_logits,
                 num_tokens_local=num_tokens_local,
                 num_tokens_global=num_tokens_global,
-                world_size=world_size,
             )
             hidden_states = AuxLossScaler.apply(hidden_states, z_loss_l)
 
diff --git a/xtuner/v1/loss/moe_loss.py b/xtuner/v1/loss/moe_loss.py
index 3ee6f38f8..c70d9e17b 100644
--- a/xtuner/v1/loss/moe_loss.py
+++ b/xtuner/v1/loss/moe_loss.py
@@ -5,7 +5,6 @@
 from cyclopts import Parameter
 from pydantic import BaseModel, ConfigDict
 from torch import distributed as dist
-from torch.distributed._functional_collectives import all_reduce
 
 from xtuner.v1.utils.device import get_device
 
@@ -13,24 +12,6 @@
 DEVICE = get_device()
 
 
-class _AllReduce(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, op, group, tensor):
-        ctx.group = group
-        ctx.op = op
-        tensor = tensor.clone(memory_format=torch.contiguous_format)
-        tensor = all_reduce(tensor, op, group=group)
-        return tensor
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return (None, None) + (_AllReduce.apply(ctx.op, ctx.group, grad_output),)
-
-
-def all_reduce_autograd(tensor, op, group):
-    return _AllReduce.apply(op, group, tensor)
-
-
 class BalancingLossConfig(BaseModel):
     """Balancing loss configuration for MoE models.
 
@@ -139,11 +120,11 @@ def finalize(
             non_pad_token (int): Number of non-padding tokens on this rank.
 
         Returns:
-            tuple[torch.Tensor, torch.Tensor]: ``(balancing_loss, local_balancing_loss)``. The first
-            carries the autograd graph used for backward; in the global-average branch it is reduced
-            across ranks via ``all_reduce_autograd``. The second is a detached, per-rank local
-            component for logging whose cross-rank SUM reproduces the first (used by the display
-            pipeline so that logged curves stay global once backward switches to the local component).
+            tuple[torch.Tensor, torch.Tensor]: ``(balancing_loss, local_balancing_loss)``. Both are
+            this rank's per-rank local component (computed from its own ``local_gating_sum`` with the
+            global detached statistics); the first carries the autograd graph for backward, aggregated
+            across ranks by the SUM gradient reduction, while the second is detached for the display
+            pipeline whose cross-rank SUM restores the global balancing loss.
         """
         routing_weights_sum_list = self.routing_weights_sum_list
         self.routing_weights_sum_list = []
@@ -260,7 +241,6 @@ def accumulate(
         router_logits: torch.Tensor,
         num_tokens_local: int,
         num_tokens_global: torch.Tensor | None,
-        world_size: int,
     ) -> torch.Tensor:
         """Compute z-loss for one layer and return it as a scalar with autograd
         attached.
@@ -278,8 +258,6 @@ def accumulate(
             num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across
                 ranks, as an int64 scalar tensor. ``None`` when ``z_loss_global_average`` is off
                 or the process group is not initialized.
-            world_size (int): Number of ranks contributing to ``num_tokens_global``. Ignored when
-                ``num_tokens_global`` is ``None``.
 
         Returns:
             torch.Tensor: Per-layer z-loss as a 0-d tensor with autograd graph back to
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index c9f337084..d58c63639 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -248,24 +248,23 @@ def _z_loss_dist_token_count(
         z_ctx: list[ZLossContext] | ZLossContext | None,
         num_tokens_local: int,
         device: torch.device | str | int,
-    ) -> tuple[torch.Tensor | None, int]:
+    ) -> torch.Tensor | None:
         """Compute the cross-rank non-padding token count needed by the z-loss
         inline path.
 
-        Returns ``(num_tokens_global, world_size)``. ``num_tokens_global`` is ``None`` (i.e. skip
-        global averaging) when there is no z-loss context, when the configured z-loss is not
-        global-average, or when no process group is initialized.
+        Returns the global non-padding token count, or ``None`` (i.e. skip global averaging) when
+        there is no z-loss context, when the configured z-loss is not global-average, or when no
+        process group is initialized.
         """
         if z_ctx is None:
-            return None, 1
+            return None
         first = z_ctx[0] if isinstance(z_ctx, list) else z_ctx
         if not first.loss_cfg.z_loss_global_average or not dist.is_initialized():
-            return None, 1
+            return None
         n = torch.tensor(num_tokens_local, device=device, dtype=torch.int64)
         group = dist.group.WORLD
         assert group is not None
-        n_global = all_reduce(n, "sum", group)
-        return n_global, dist.get_world_size()
+        return all_reduce(n, "sum", group)
 
     def _extract_aux_loss_ctx(
         self,
@@ -487,7 +486,7 @@ def _micro_batch_forward(
             [{} for _ in range(len(seq_ctx_list))] if keep_router else []
         )
         balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx_list)
-        num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device)
+        num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device)
 
         # Process through layers
         cat_seq_ctx: SequenceContext | None = None
@@ -569,7 +568,6 @@ def _micro_batch_forward(
                     z_ctx=z_ctx,
                     num_tokens_local=non_pad_token,
                     num_tokens_global=num_tokens_global,
-                    world_size=z_world_size,
                 )
 
         assert hidden_states_list, "XTuner Internal Error, found empty hidden states for domino EP"
@@ -754,7 +752,7 @@ def _forward(
         # Hoisted out of the per-layer accumulate path: mask is constant across layers.
         nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
         non_pad_token = nonpad_indices.numel()
-        num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
+        num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
@@ -796,7 +794,6 @@ def _forward(
                     z_ctx=z_ctx,
                     num_tokens_local=non_pad_token,
                     num_tokens_global=num_tokens_global,
-                    world_size=z_world_size,
                 )
 
             if self.config.return_hidden_states:
@@ -827,9 +824,7 @@ def _forward(
             # MTP uses its own mask; main mask's non-pad indices do not apply.
             mtp_nonpad_indices = torch.nonzero(mtp_seq_ctx.mask, as_tuple=True)[1]
             mtp_non_pad_token = mtp_nonpad_indices.numel()
-            mtp_num_tokens_global, mtp_z_world_size = self._z_loss_dist_token_count(
-                z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device
-            )
+            mtp_num_tokens_global = self._z_loss_dist_token_count(z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device)
 
             # Forward through MTP block
             mtp_outputs = self.mtp_block(
@@ -861,7 +856,6 @@ def _forward(
                     z_ctx=z_ctx,
                     num_tokens_local=mtp_non_pad_token,
                     num_tokens_global=mtp_num_tokens_global,
-                    world_size=mtp_z_world_size,
                 )
                 mtp_loss, (_, mtp_extra) = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
                 mtp_losses += mtp_loss
diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py
index 2c5aee6a8..a7dbb7393 100644
--- a/xtuner/v1/model/moe/qwen3vl_text.py
+++ b/xtuner/v1/model/moe/qwen3vl_text.py
@@ -148,7 +148,7 @@ def _forward(
         # Hoisted out of the per-layer accumulate path: mask is constant across layers.
         nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
         non_pad_token = nonpad_indices.numel()
-        num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
+        num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
 
         # =====================================================
         deepstack_visual_embeds = seq_ctx.deepstack_visual_embeds
@@ -196,7 +196,6 @@ def _forward(
                     z_ctx=z_ctx,
                     num_tokens_local=non_pad_token,
                     num_tokens_global=num_tokens_global,
-                    world_size=z_world_size,
                 )
 
             if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))):