diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 061e2f6e29..53ab344079 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/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,13 +1082,13 @@ 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 @@ -1095,12 +1096,12 @@ def fully_shard( 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( diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 702984aa03..8287c2274f 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/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( diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py index dc42366068..edfb238062 100644 --- a/xtuner/v1/module/mtp/mtp_block.py +++ b/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(