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 061e2f6e29..862984cc8c 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,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, 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 702984aa03..80287b8fb3 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -221,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 @@ -228,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, @@ -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( diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py index dc42366068..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,