From ffc3657f05ce9287803ad74080224d43919ad396 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Tue, 30 Jun 2026 15:28:04 +0000 Subject: [PATCH 1/5] =?UTF-8?q?MTP=E5=85=BC=E5=AE=B9=E6=80=A7=E9=80=82?= =?UTF-8?q?=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- xtuner/v1/model/moe/moe.py | 11 ++++++----- xtuner/v1/module/mtp/mtp_block.py | 7 +++++++ 2 files changed, 13 insertions(+), 5 deletions(-) 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/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( From 1a729fc276061372630f01a32f62358f915badc0 Mon Sep 17 00:00:00 2001 From: "houyufeng4@huawei.com" Date: Mon, 1 Jun 2026 06:51:15 +0000 Subject: [PATCH 2/5] =?UTF-8?q?=E7=BC=93=E8=A7=A3MoE=E8=B4=9F=E8=BD=BD?= =?UTF-8?q?=E4=B8=8D=E5=9D=87=EF=BC=9Askip=20pad=5Ftoken?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../module/decoder_layer/moe_decoder_layer.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) 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( From 7ee76765f53b4ad1c15b4147d17446a5e2494e07 Mon Sep 17 00:00:00 2001 From: chenchiyu Date: Mon, 13 Jul 2026 12:38:52 +0000 Subject: [PATCH 3/5] Optimize MTP forward memory scheduling --- xtuner/v1/model/moe/moe.py | 42 +++++++++++++++++++++++++++----------- 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 53ab344079..a2d5677ddc 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -739,14 +739,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 ( @@ -806,6 +798,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, @@ -1082,7 +1085,6 @@ 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): @@ -1100,9 +1102,6 @@ def fully_shard( offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=mtp_layer, ) - if mtp_idx == 0: - 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( list(self.mtp_block.layers)[:-1], @@ -1110,6 +1109,25 @@ def fully_shard( ): 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 + + assert self.config.mtp_config is not None + # 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 self.config.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, From 05dd841e9bc2666d4e4fcb02946b3452f8ce41f6 Mon Sep 17 00:00:00 2001 From: chenchiyu Date: Wed, 15 Jul 2026 08:29:00 +0000 Subject: [PATCH 4/5] Optimize padding token dispatch in MoE layers --- xtuner/v1/data_proto/sequence_context.py | 54 ++++++++++++++++++- xtuner/v1/model/moe/moe.py | 24 +++++---- xtuner/v1/model/moe/qwen3vl_text.py | 4 +- .../module/decoder_layer/moe_decoder_layer.py | 26 +++++---- 4 files changed, 85 insertions(+), 23 deletions(-) diff --git a/xtuner/v1/data_proto/sequence_context.py b/xtuner/v1/data_proto/sequence_context.py index 84ac054ba9..20456377b7 100644 --- a/xtuner/v1/data_proto/sequence_context.py +++ b/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 diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index a2d5677ddc..64e827d343 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/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,8 +700,9 @@ def _forward( seq_ctx=seq_ctx, ) else: - 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)): + 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, @@ -752,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 @@ -891,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) @@ -956,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 @@ -1088,9 +1092,11 @@ def fully_shard( # 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 - ) or True: # share mtp head must recompute + if ( + self._should_recompute(None, mtp_idx=mtp_idx) + or (self.config.mtp_config is not None and self.config.mtp_config.share_weights) + 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 diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py index fafb70a80e..075ea10a59 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/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 diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 8287c2274f..80287b8fb3 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -1,6 +1,5 @@ from functools import partial from typing import Literal, Protocol, TypeAlias, cast -import os import torch import torch.nn as nn @@ -222,6 +221,7 @@ 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 @@ -229,6 +229,7 @@ def __init__( 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, @@ -382,13 +383,12 @@ 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] + 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,:] + 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) @@ -397,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"] if not skip_pad_tokens else router_results["topk_ids"][nonpad_indices, :], + 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"] if not skip_pad_tokens else router_results["topk_weights"][nonpad_indices, :], + 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( @@ -463,10 +467,10 @@ def _forward( hidden_states = self._post_moe_forward( combined_hidden_states=combined_hidden_states, - residual=residual if not skip_pad_tokens else residual[:,nonpad_indices,:], + residual=residual if not self.skip_dispatch_pad_tokens else residual[:, nonpad_indices, :], shared_experts_out=shared_experts_out, ) - if skip_pad_tokens: + if self.skip_dispatch_pad_tokens: result = torch.zeros_like(origin_hidden_states) result[:, nonpad_indices, :] = hidden_states result[:, pad_indices, :] = pad_hidden_states From f7efe1e825a2d26a11d5110d569613d83cf07813 Mon Sep 17 00:00:00 2001 From: chenchiyu Date: Tue, 28 Jul 2026 08:17:56 +0000 Subject: [PATCH 5/5] Refine MTP recomputation and FSDP scheduling --- xtuner/v1/model/moe/moe.py | 22 +++++++------- xtuner/v1/module/mtp/mtp_block.py | 49 ++++++++++++++++++------------- 2 files changed, 40 insertions(+), 31 deletions(-) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 64e827d343..862984cc8c 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1091,24 +1091,27 @@ def fully_shard( ) # 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) - or True - ): # 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=True, + # 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 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:], @@ -1119,14 +1122,13 @@ def fully_shard( last_decoder_layer = list(self.layers.values())[-1] last_decoder_layer.set_modules_to_forward_prefetch([first_mtp_layer]) # type: ignore - assert self.config.mtp_config is not None # 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 self.config.mtp_config.share_weights: + 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 diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py index edfb238062..3d4ebd2ee9 100644 --- a/xtuner/v1/module/mtp/mtp_block.py +++ b/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,6 +113,7 @@ 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" @@ -119,26 +121,38 @@ def forward( 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, @@ -171,7 +185,6 @@ 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): @@ -192,12 +205,6 @@ 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(