[Fix] Improve MTP compatibility and mitigate MoE load imbalance by skipping pad tokens

xtuner/v1/data_proto/sequence_context.py

@@ -31,7 +31,9 @@ class SequenceContext:
     cu_seq_lens_k: torch.IntTensor
     max_length_q: torch.Tensor
     max_length_k: torch.Tensor
-    num_padding: int
+    _num_padding: int
+    _nonpad_indices: torch.LongTensor | None
+    _pad_indices: torch.LongTensor | None
     sequence_parallel_mesh: DeviceMesh | None
     block_table: torch.Tensor | None
     device: str | torch.device  # TODO: 这个地方有点乱，到处是 device
@@ -98,6 +100,8 @@ def __init__(
             self.max_length_k = torch.tensor(max_length_k, device="cpu")
         else:
             self.max_length_k = max_length_k
+        self._nonpad_indices = None
+        self._pad_indices = None
         self.num_padding = num_padding
         self.sequence_parallel_mesh = sequence_parallel_mesh
         self.block_table = block_table
@@ -334,6 +338,53 @@ def mask(self) -> torch.BoolTensor:
             mask[..., -self.num_padding :] = False
         return mask
 
+    def _invalidate_padding_indices(self) -> None:
+        self._nonpad_indices = None
+        self._pad_indices = None
+
+    def _cache_padding_indices(self) -> None:
+        if self.input_ids is not None:
+            seq_len = self.input_ids.shape[1]
+            device = self.input_ids.device
+        else:
+            assert self.inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
+            seq_len = self.inputs_embeds.shape[1]
+            device = self.inputs_embeds.device
+
+        # SequenceContext.mask only represents a contiguous padding suffix, so slicing a
+        # single arange is equivalent to running nonzero on mask and ~mask, without the
+        # mask allocation and two nonzero kernels. If arbitrary-position padding is
+        # supported in the future, derive both index tensors from mask instead.
+        num_non_padding = seq_len if self.num_padding <= 0 else max(seq_len - self.num_padding, 0)
+        token_indices = torch.arange(seq_len, dtype=torch.long, device=device)
+        self._nonpad_indices = cast(torch.LongTensor, token_indices[:num_non_padding])
+        self._pad_indices = cast(torch.LongTensor, token_indices[num_non_padding:])
+
+    @property
+    def num_padding(self) -> int:
+        return self._num_padding
+
+    @num_padding.setter
+    def num_padding(self, value: int) -> None:
+        self._num_padding = value
+        self._invalidate_padding_indices()
+
+    @property
+    def nonpad_indices(self) -> torch.LongTensor:
+        """Indices of non-padding tokens, cached for reuse across decoder
+        layers."""
+        if self._nonpad_indices is None:
+            self._cache_padding_indices()
+        return cast(torch.LongTensor, self._nonpad_indices)
+
+    @property
+    def pad_indices(self) -> torch.LongTensor:
+        """Indices of padding tokens, cached for reuse across decoder
+        layers."""
+        if self._pad_indices is None:
+            self._cache_padding_indices()
+        return cast(torch.LongTensor, self._pad_indices)
+
     @property
     def seq_lens_q(self) -> torch.LongTensor:
         return self.cu_seq_lens_q[1:] - self.cu_seq_lens_q[:-1]  # type: ignore
@@ -531,6 +582,7 @@ def to(self, device: torch.device | str):
         if self.rollout_routed_experts is not None and hasattr(self.rollout_routed_experts, "to"):
             self.rollout_routed_experts = self.rollout_routed_experts.to(device)  # type: ignore
 
+        self._invalidate_padding_indices()
         self.device = device
 
         return self

xtuner/v1/model/moe/moe.py

@@ -155,6 +155,7 @@ class MoEConfig(TransformerConfig):
     freeze_routers: bool = False
     router_async_offload: bool = False
     aux_loss_cfg: AuxLossConfig = AuxLossConfig()
+    skip_dispatch_pad_tokens: Annotated[bool, Parameter(group="moe")] = False
     # TODO: `FSDPConfig` should be model-specific; temporarily keep
     # `embed_reshard_after_forward` here until per-submodule FSDP config is supported.
     # Compose models call `self.embed_tokens` multiple times per step, so default to
@@ -687,9 +688,9 @@ def _forward(
         self._mark_dynamic(seq_ctx)
         balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx)
         # Hoisted out of the per-layer accumulate path: mask is constant across layers.
-        nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
+        nonpad_indices = seq_ctx.nonpad_indices
         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, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, nonpad_indices.device)
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
@@ -699,7 +700,9 @@ def _forward(
                     seq_ctx=seq_ctx,
                 )
             else:
-                if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
+                if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1 and (
+                    self.mtp_block is None or (self.mtp_block is not None and int(idx) > 0)
+                ):
                     with async_save_on_cpu(
                         h2d_stream=self.offload_stream,
                         d2h_stream=self.offload_stream,
@@ -738,14 +741,6 @@ def _forward(
                 output["hidden_states"].append(hidden_states)
 
         layer_hidden_states = hidden_states
-        hidden_states = self.norm(hidden_states)
-
-        # Get LM loss context from dict
-        lm_loss_ctx = loss_ctx["lm"] if loss_ctx is not None else None
-        loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx)  # type: ignore
-        output["loss"] = loss
-        output["logits"] = logits
-        output["extra_info"] = extra_info
 
         # MTP forward pass and loss computation
         if (
@@ -759,10 +754,10 @@ def _forward(
                 inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None,
             )
             # 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_nonpad_indices = mtp_seq_ctx.nonpad_indices
             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
+                z_ctx, mtp_non_pad_token, mtp_nonpad_indices.device
             )
 
             # Forward through MTP block
@@ -805,6 +800,17 @@ def _forward(
             # Add to total loss
             output["mtp_loss"] = scaled_mtp_loss
 
+        # Keep the main LM branch after MTP so the final normalized states, logits,
+        # and unsharded LM-head parameters are not resident across the whole MTP forward.
+        hidden_states = self.norm(layer_hidden_states)
+
+        # Get LM loss context from dict
+        lm_loss_ctx = loss_ctx["lm"] if loss_ctx is not None else None
+        loss, (logits, extra_info) = self.lm_head(hidden_states, lm_loss_ctx)  # type: ignore
+        output["loss"] = loss
+        output["logits"] = logits
+        output["extra_info"] = extra_info
+
         split_aux_output = self.aux_loss.finalize(
             balancing_ctx=balancing_ctx,
             z_ctx=z_ctx,
@@ -887,6 +893,7 @@ def build_layers(self, config: MoEConfig) -> nn.ModuleDict:
                     layer_idx=layer_idx,
                     dispatcher=config.dispatcher,
                     ep_mesh=self.ep_mesh,
+                    skip_dispatch_pad_tokens=config.skip_dispatch_pad_tokens,
                 )
                 if self.config.freeze_routers:
                     layers[str(layer_idx)].gate.requires_grad_(False)
@@ -952,6 +959,7 @@ def build_mtp_block(self, config: MoEConfig) -> MTPBlock:
                 layer_idx=config.num_hidden_layers + i,
                 dispatcher=config.dispatcher,
                 ep_mesh=self.ep_mesh,
+                skip_dispatch_pad_tokens=config.skip_dispatch_pad_tokens,
             )
 
             # Wrap decoder layer in MTPLayer
@@ -1081,34 +1089,53 @@ def fully_shard(
             offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None,
             module=self.lm_head,
         )
-
         # Shard MTP block if it exists
         if self.mtp_block is not None:
+            assert self.config.mtp_config is not None
+            mtp_config = self.config.mtp_config
             for mtp_idx, mtp_layer in enumerate(self.mtp_block.layers):
-                if self._should_recompute(None, mtp_idx=mtp_idx) or (
-                    self.config.mtp_config is not None and self.config.mtp_config.share_weights
-                ):  # share mtp head must recompute
+                # One shared physical layer serves every logical MTP depth. Always
+                # checkpoint it to avoid retaining activations from all depths.
+                if self._should_recompute(None, mtp_idx=mtp_idx) or mtp_config.share_weights:
                     mtp_layer = checkpoint_wrapper(mtp_layer, checkpoint_impl=CheckpointImpl.REENTRANT)
                 self.mtp_block.layers[mtp_idx] = mtp_layer
 
-                reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1
                 self._fully_shard(
                     mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh,
                     mp_policy=mp_policy,
-                    reshard_after_forward=reshard_after_forward,
+                    # Reuse the unsharded shared layer across logical depths. Explicit
+                    # resharding at the end of MTPBlock.forward is currently disabled and
+                    # can be enabled if keeping the shared layer unsharded causes excessive
+                    # peak memory usage.
+                    reshard_after_forward=not mtp_config.share_weights,
                     offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None,
                     module=mtp_layer,
                 )
-                if mtp_idx == 0:
-                    layer_next.set_modules_to_forward_prefetch([mtp_layer])  # type: ignore
-
-            if self.config.mtp_config is not None and self.config.mtp_config.num_layers > 0:
+            if mtp_config.num_layers > 0:
                 for prev_mtp_layer, next_mtp_layer in zip(
                     list(self.mtp_block.layers)[:-1],
                     list(self.mtp_block.layers)[1:],
                 ):
                     prev_mtp_layer.set_modules_to_forward_prefetch([next_mtp_layer])  # type: ignore
 
+            first_mtp_layer = self.mtp_block.layers[0]
+            last_decoder_layer = list(self.layers.values())[-1]
+            last_decoder_layer.set_modules_to_forward_prefetch([first_mtp_layer])  # type: ignore
+
+            # Shared-weight MTP reuses one physical FSDP layer at every logical depth. A static
+            # forward-prefetch hook on that layer would therefore materialize the LM head after
+            # the first depth and keep it resident for all remaining depths. We intentionally
+            # leave the first post-MTP LM-head call to unshard on demand for now. A future
+            # optimization can explicitly prefetch the LM head in MTPBlock only after its final
+            # logical depth finishes.
+            if not mtp_config.share_weights:
+                # MTP outputs are projected by the shared LM head before the main LM branch,
+                # so non-shared MTP can safely prefetch it from the last physical MTP layer.
+                self.mtp_block.layers[-1].set_modules_to_forward_prefetch([self.lm_head])  # type: ignore
+        else:
+            last_decoder_layer = list(self.layers.values())[-1]
+            last_decoder_layer.set_modules_to_forward_prefetch([self.lm_head])  # type: ignore
+
         self._fully_shard(
             mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh,
             mp_policy=mp_policy,

xtuner/v1/model/moe/qwen3vl_text.py

@@ -146,9 +146,9 @@ def _forward(
         self._mark_dynamic(seq_ctx)
         balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx)
         # Hoisted out of the per-layer accumulate path: mask is constant across layers.
-        nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
+        nonpad_indices = seq_ctx.nonpad_indices
         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, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, nonpad_indices.device)
 
         # =====================================================
         deepstack_visual_embeds = seq_ctx.deepstack_visual_embeds

xtuner/v1/module/decoder_layer/moe_decoder_layer.py

@@ -221,13 +221,15 @@ def __init__(
         layer_idx: int = 0,
         dispatcher: Literal["deepep", "all2all", "agrs"] | None,
         ep_mesh: DeviceMesh | None = None,
+        skip_dispatch_pad_tokens: bool = False,
     ):
         super().__init__()
         self.ep_mesh = ep_mesh
         self.hidden_size = hidden_size
         self.n_routed_experts = n_routed_experts
         self.n_shared_experts = n_shared_experts
         self.hidden_factor = hidden_factor
+        self.skip_dispatch_pad_tokens = skip_dispatch_pad_tokens
 
         self.self_attn: MultiHeadAttention | MultiLatentAttention | GatedDeltaNet = attention_config.build(
             hidden_size=hidden_size,
@@ -381,7 +383,12 @@ def _forward(
             position_embeddings=position_embeddings,
             state=ForwardState.TRAINING,
         )
-
+        if self.skip_dispatch_pad_tokens:
+            nonpad_indices = seq_ctx.nonpad_indices
+            pad_indices = seq_ctx.pad_indices
+            origin_hidden_states = hidden_states
+            hidden_states = origin_hidden_states[:, nonpad_indices, :]
+            pad_hidden_states = origin_hidden_states[:, pad_indices, :]
         origin_shape = hidden_states.shape
 
         # reshape hidden_states to (batch_size * seq_len, hidden_size)
@@ -390,11 +397,15 @@ def _forward(
         # )
         pre_dispatched = self.dispatcher.dispatch_preprocess(
             hidden_states=hidden_states.view(-1, hidden_states.shape[-1]),
-            topk_ids=router_results["topk_ids"],
+            topk_ids=router_results["topk_ids"]
+            if not self.skip_dispatch_pad_tokens
+            else router_results["topk_ids"][nonpad_indices, :],
         )
         dispatched = self.dispatcher.dispatch(
             pre_dispatched=pre_dispatched,
-            topk_weights=router_results["topk_weights"],
+            topk_weights=router_results["topk_weights"]
+            if not self.skip_dispatch_pad_tokens
+            else router_results["topk_weights"][nonpad_indices, :],
             decoding=False,
         )  # type: ignore[call-overload]
         post_dispatched = self.dispatcher.dispatch_postprocess(
@@ -456,9 +467,14 @@ def _forward(
 
         hidden_states = self._post_moe_forward(
             combined_hidden_states=combined_hidden_states,
-            residual=residual,
+            residual=residual if not self.skip_dispatch_pad_tokens else residual[:, nonpad_indices, :],
             shared_experts_out=shared_experts_out,
         )
+        if self.skip_dispatch_pad_tokens:
+            result = torch.zeros_like(origin_hidden_states)
+            result[:, nonpad_indices, :] = hidden_states
+            result[:, pad_indices, :] = pad_hidden_states
+            hidden_states = result
         return hidden_states, router_results["logits"], router_results["router_weights"]
 
     def _micro_batch_forward(

xtuner/v1/module/mtp/mtp_block.py

@@ -4,6 +4,7 @@
 
 import torch
 import torch.nn as nn
+from torch.distributed.fsdp import FSDPModule
 
 from xtuner.v1.data_proto import SequenceContext
 
@@ -112,33 +113,46 @@ def forward(
                 For ``N`` micro-batches, ``list[list[(hidden, router_logits, router_weights)]]``
                 with outer length ``N`` and inner length ``D``: ``outputs[mb_idx][depth_idx]``.
         """
+        outputs: list[MTPDepthOutput] | list[list[MTPDepthOutput]]
         if len(hidden_states) == 1:
             assert isinstance(seq_ctx, SequenceContext), (
                 "seq_ctx should be a SequenceContext instance in single-microbatch mode"
             )
             assert isinstance(position_embeddings, tuple) and len(position_embeddings) == 2, (
                 "position_embeddings should be a (cos, sin) tuple in single-microbatch mode"
             )
-            return self._forward(
+            outputs = self._forward(
                 hidden_states=hidden_states[0],
                 embed_tokens_fn=embed_tokens_fn,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
             )
+        else:
+            n = len(hidden_states)
+            assert isinstance(seq_ctx, list) and len(seq_ctx) == n, (
+                "seq_ctx should be a list aligned with hidden_states in multi-microbatch mode"
+            )
+            assert isinstance(position_embeddings, list) and len(position_embeddings) == n, (
+                "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
+            )
+            outputs = self._micro_batch_forward(
+                hidden_states_list=list(hidden_states),
+                embed_tokens_fn=embed_tokens_fn,
+                position_embeddings_list=position_embeddings,
+                seq_ctx_list=seq_ctx,
+            )
 
-        n = len(hidden_states)
-        assert isinstance(seq_ctx, list) and len(seq_ctx) == n, (
-            "seq_ctx should be a list aligned with hidden_states in multi-microbatch mode"
-        )
-        assert isinstance(position_embeddings, list) and len(position_embeddings) == n, (
-            "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
-        )
-        return self._micro_batch_forward(
-            hidden_states_list=list(hidden_states),
-            embed_tokens_fn=embed_tokens_fn,
-            position_embeddings_list=position_embeddings,
-            seq_ctx_list=seq_ctx,
-        )
+        # Explicit resharding is currently disabled to avoid an extra reshard/unshard
+        # cycle. Consider enabling it if keeping the shared MTP layer unsharded causes
+        # excessive peak memory usage.
+        # self._reshard_shared_layer()
+        return outputs
+
+    def _reshard_shared_layer(self) -> None:
+        if self.mtp_config.share_weights:
+            shared_layer = self.layers[0]
+            if isinstance(shared_layer, FSDPModule):
+                shared_layer.reshard()
 
     def _forward(
         self,