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

xtuner/v1/model/moe/moe.py

@@ -699,7 +699,8 @@ 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,
@@ -1081,26 +1082,26 @@ def fully_shard(
             offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None,
             module=self.lm_head,
         )
-
+        layer_next.set_modules_to_forward_prefetch([self.lm_head])
         # Shard MTP block if it exists
         if self.mtp_block is not None:
             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
+                ) or True:  # share mtp head must recompute
                     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,
+                    reshard_after_forward=True,
                     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
+                    self.lm_head.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:
                 for prev_mtp_layer, next_mtp_layer in zip(

xtuner/v1/module/decoder_layer/moe_decoder_layer.py

@@ -1,5 +1,6 @@
 from functools import partial
 from typing import Literal, Protocol, TypeAlias, cast
+import os
 
 import torch
 import torch.nn as nn
@@ -381,7 +382,13 @@ def _forward(
             position_embeddings=position_embeddings,
             state=ForwardState.TRAINING,
         )
-
+        skip_pad_tokens = (os.environ.get("SKIP_PAD_TOKENS", "False") == "True")
+        if skip_pad_tokens:
+            nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
+            pad_indices = torch.nonzero(~seq_ctx.mask, as_tuple=True)[1]
+            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,11 @@ 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 skip_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 skip_pad_tokens else router_results["topk_weights"][nonpad_indices, :],
             decoding=False,
         )  # type: ignore[call-overload]
         post_dispatched = self.dispatcher.dispatch_postprocess(
@@ -456,9 +463,14 @@ def _forward(
 
         hidden_states = self._post_moe_forward(
             combined_hidden_states=combined_hidden_states,
-            residual=residual,
+            residual=residual if not skip_pad_tokens else residual[:,nonpad_indices,:],
             shared_experts_out=shared_experts_out,
         )
+        if skip_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

@@ -171,6 +171,7 @@ def _forward(
         mtp_outputs: list[MTPDepthOutput] = []
         current_hidden_states = hidden_states.detach() if self.mtp_config.detach_mtp_inputs else hidden_states
         current_seq_ctx = seq_ctx
+        shared_layer = self.layers[0] if self.mtp_config.share_weights else None
 
         num_steps = self.mtp_config.num_layers
         for step in range(num_steps):
@@ -191,6 +192,12 @@ def _forward(
             )
             mtp_outputs.append((current_hidden_states, router_logits, router_weights))
 
+
+        # Shared MTP reuses one physical FSDP layer across multiple steps.
+        # Keep it unsharded during inner steps, then reshard once at block end.
+        if shared_layer is not None:
+            shared_layer.reshard()
+
         return mtp_outputs
 
     def _micro_batch_forward(