diff --git a/docs/user-guide/features/multi_token_prediction.md b/docs/user-guide/features/multi_token_prediction.md index 4059fa5326e..64c89a8000a 100644 --- a/docs/user-guide/features/multi_token_prediction.md +++ b/docs/user-guide/features/multi_token_prediction.md @@ -18,6 +18,31 @@ We can train GPTModel like models with Multi-Token Prediction (MTP) by setting m | mtp_num_layers | Number of Multi-Token Prediction (MTP) Layers. MTP extends the prediction scope to multiple future tokens at each position. This MTP implementation sequentially predict additional tokens by using D sequential modules to predict D additional tokens. Default is None. | | mtp_loss_scaling_factor | Scaling factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP losses across all depths, and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. Default is 0.1. | +## Pipeline Parallel Layout for MTP + +MTP supports flexible placement of MTP layers across pipeline stages using a custom `pipeline_model_parallel_layout`. By default, all MTP layers are placed on the last pipeline stage, but you can customize their placement. + +### MTP Standalone Mode + +When MTP layers are placed in a separate virtual pipeline (vpp) stage that is not on the last pipeline rank, the `mtp_standalone` flag is automatically set to `True`. This mode enables MTP to run independently in its own pipeline stage. + +### Layout Format + +Use `m` to represent MTP layers in the pipeline layout string. For example: +- `"E|t*3|(t|)*5mL"` - MTP in the last stage +- `"E|t*3|(t|)*4tm|L"` - MTP in the second-to-last stage with a decoder layer +- `"E|t*3|(t|)*3tt|m|L"` - MTP in a standalone stage (second-to-last) with no other layers + +### Constraints + +- All MTP layers must be placed in the same one virtual pipeline stage. +- MTP layers cannot be placed on the first pipeline rank. + +## Implementation Notes + +- For models with MTP layers, the final layernorm is placed in the stage that contains the last decoder layer, rather than in the post-process stage. This may cause small numerical differences in gradient norm reduction when final layernorm is placed in different pipeline stages in deterministic mode. Bitwise alignment can be achieved by disabling gradient norm clipping. +- MTP loss is computed in the post-processing stage. + ## Precautions Please do not use Context Parallel (CP), or arbitrary AttnMaskType, or learned absolute position embedding type with MTP. These use cases are not yet supported. diff --git a/gpt_builders.py b/gpt_builders.py index 4b86e30e597..13c914acc56 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -7,6 +7,7 @@ get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_with_inference_spec, get_gpt_mtp_block_spec, + get_gpt_decoder_layer_specs, ) from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import ( get_gpt_heterogeneous_layer_spec, @@ -69,7 +70,12 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ # Only happens with block spec (TransformerBlockSubmodules) when using MoE. transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config) else: - transformer_layer_spec_for_mtp = transformer_layer_spec + # Define the decoder block spec + decoder_layer_specs = get_gpt_decoder_layer_specs( + config, use_transformer_engine=use_te, normalization=args.normalization, qk_l2_norm=args.qk_l2_norm, vp_stage=vp_stage + ) + transformer_layer_spec_for_mtp = decoder_layer_specs[-1] + # Use spec of the last layer in decoder block as spec of the transformer layer in MTP mtp_block_spec = get_gpt_mtp_block_spec( config, transformer_layer_spec_for_mtp, @@ -101,12 +107,12 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ def _get_transformer_layer_spec(use_te, config): """Get transformer layer specification based on configuration. - + Args: use_te (bool): Whether to use Transformer Engine args: Training arguments config: Model configuration - + Returns: transformer_layer_spec: The transformer layer specification """ diff --git a/megatron/core/datasets/blended_megatron_dataset_builder.py b/megatron/core/datasets/blended_megatron_dataset_builder.py index 2f10bbade1c..f728fe10d03 100644 --- a/megatron/core/datasets/blended_megatron_dataset_builder.py +++ b/megatron/core/datasets/blended_megatron_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index 55663acdc10..ddaeb7e8d84 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/megatron/core/distributed/finalize_model_grads.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from typing import Callable, List, Optional, Union @@ -193,7 +193,11 @@ def _allreduce_word_embedding_grads( pp_group = parallel_state.get_pipeline_model_parallel_group() _allreduce_embedding_grad( - model, embd_group, pp_group, partial(_get_shared_word_embedding_weight, config=config) + model, + embd_group, + pp_group, + partial(_get_shared_word_embedding_weight, config=config), + config=config, ) @@ -203,6 +207,7 @@ def _allreduce_embedding_grad( pp_group: torch.distributed.ProcessGroup, weight_getter: Callable[[torch.nn.Module], Optional[torch.nn.Parameter]], skip_if_none: bool = True, + config: TransformerConfig = None, ): """Unified helper to all-reduce embedding parameters across pipeline stages. @@ -229,6 +234,9 @@ def _allreduce_embedding_grad( model_module = model[0] elif is_pp_last_stage(pp_group): model_module = model[-1] + elif getattr(config, 'mtp_num_layers', None) is not None and config.mtp_num_layers > 0: + # Embedding for MTP layers is in the last virtual pipeline model parallel stage. + model_module = model[-1] else: # We do not support an interleaved schedule for models with encoders yet. model_module = model[0] diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 184e972476c..44e4dd52efe 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import warnings from dataclasses import dataclass @@ -329,6 +329,10 @@ class ModelParallelConfig: rank 1 | 0 1 2 0 1 2 3 4 3 4 """ + mtp_standalone: bool = False + """This will be set automatically according to the pipeline layout, + and will be set to True if MTP is in a separate vpp stage.""" + ################### # CPU Offloading ################### diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 64b86f869e1..b0fa6126b63 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -68,6 +68,8 @@ def _is_in_embd_group(self): if torch.distributed.get_rank() in torch.distributed.get_process_group_ranks( self.embd_group ): + if getattr(self, 'mtp_process', False): + return True if ( torch.distributed.get_rank() == torch.distributed.get_process_group_ranks(self.embd_group)[0] @@ -207,7 +209,10 @@ def setup_embeddings_and_output_layer(self) -> None: ): self.shared_embedding_or_output_weight().shared_embedding = True - if (self.post_process or getattr(self, 'mtp_process', False)) and not self.pre_process: + if ( + (self.post_process and self.share_embeddings_and_output_weights) + or getattr(self, 'mtp_process', False) + ) and not self.pre_process: assert not ( is_vp_first_stage(self.vp_stage, self.vp_size) and is_pp_first_stage(self.pp_group) ) diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index 9b2eb9cdde2..02312c1c59b 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/megatron/core/models/common/model_chunk_schedule_plan.py @@ -14,7 +14,6 @@ get_comm_stream, get_comp_stream, ) -from megatron.core.transformer.multi_token_prediction import get_mtp_num_layers_to_build class ModelChunkState: @@ -352,37 +351,40 @@ def __init__( self._model_chunk_state.context_mask = None self._model_chunk_state.attention_bias = None - transformer_num_layers = model.decoder.num_layers_per_pipeline_rank - mtp_num_layers = get_mtp_num_layers_to_build(model.config, vp_stage=self.vp_stage) - # build preprocess self.pre_process = PreProcessNode(model, self._model_chunk_state, self._event, comp_stream) - # build layer schedule plan for each layer - for layer_idx in range(transformer_num_layers): - layer = model.decoder._get_layer(layer_idx) - layer_plan = TransformerLayerSchedulePlan( - layer, self._event, self._model_chunk_state, comp_stream, comm_stream + + # build layer schedule plan for each layer. + # The methods to obtain layers are different for MTP so we need the other build plan for + # MTP. Also, this can help annotate MTP layer so that it can know where MTP is. + self._build_layer_schedule_plan(model.decoder, comp_stream, comm_stream) + self._build_layer_schedule_plan(getattr(model, "mtp", None), comp_stream, comm_stream) + + # build post process + if model.post_process: + self.post_process = PostProcessNode( + model, self._model_chunk_state, self._event, comp_stream ) - self._transformer_layers.append(layer_plan) - # build mtp layers - for layer_idx in range(mtp_num_layers): + def _build_layer_schedule_plan(self, module, comp_stream, comm_stream): + if module is None: + return + num_layers = len(module.layers) + for layer_idx in range(num_layers): extra_args = { "is_first_layer": layer_idx == 0, - "is_last_layer": layer_idx == mtp_num_layers - 1, + "is_last_layer": layer_idx == num_layers - 1, } - layer = model.mtp.layers[layer_idx] layer_plan = TransformerLayerSchedulePlan( - layer, self.event, self.state, comp_stream, comm_stream, extra_args + module.layers[layer_idx], + self.event, + self.state, + comp_stream, + comm_stream, + extra_args, ) self._transformer_layers.append(layer_plan) - # build post process - if model.post_process: - self.post_process = PostProcessNode( - model, self._model_chunk_state, self._event, comp_stream - ) - @property def event(self): """Gets the CUDA event for synchronization.""" diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 1f64d3e1a0d..61fa3b95f27 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -157,9 +157,8 @@ def forward_impl(self, hidden_states): """Implements the forward pass for postprocessing. This method handles: - 1. Final layer normalization - 2. Output layer computation - 3. Loss computation if labels are provided + 1. Output layer computation + 2. Loss computation if labels are provided Args: hidden_states: The hidden states from the transformer layers. @@ -167,12 +166,11 @@ def forward_impl(self, hidden_states): Returns: The logits or loss depending on whether labels are provided. """ - # Final layer norm from Decoder - if self.gpt_model.decoder.final_layernorm and not self.gpt_model.mtp_process: - hidden_states = self.gpt_model.decoder.final_layernorm(hidden_states) - # TENorm produces a "viewed" tensor. This will result in schedule.py's - # deallocate_output_tensor() throwing an error, so a viewless tensor is - # created to prevent this. + + empty_decoder = len(self.gpt_model.decoder.layers) == 0 + layer_norm = self.gpt_model.decoder.final_layernorm + if not self.gpt_model.config.mtp_num_layers and empty_decoder and layer_norm: + hidden_states = layer_norm(hidden_states) hidden_states = make_viewless_tensor( inp=hidden_states, requires_grad=True, keep_graph=True ) @@ -251,6 +249,7 @@ def __init__( self.submodule = submodule self.detached = tuple() self.before_detached = tuple() + self.is_mtp = extra_args.get("is_mtp", False) # Create flags to indicate first and last layer self.is_first_layer = extra_args.get("is_first_layer", False) @@ -470,6 +469,12 @@ def submodule_combine_forward( # release tensor reference after use node.layer_state.residual = None + + # final layer norm from decoder + final_layernorm = node.chunk_state.model.decoder.final_layernorm + if not node.is_mtp and final_layernorm and node.is_last_layer: + output = final_layernorm(output) + output = make_viewless_tensor(inp=output, requires_grad=True, keep_graph=True) return output def mlp_wrapper(node: ScheduleNode, *args, **kwargs): @@ -509,15 +514,7 @@ def build_mtp_layer_callables(layer): def submodule_mtp_attn_forward(node, hidden_states): # MTP Block Preprocess if node.is_first_layer: - # Final layer norm from Decoder - final_layernorm = node.chunk_state.model.decoder.final_layernorm - if final_layernorm: - hidden_states = final_layernorm(hidden_states) - hidden_states = make_viewless_tensor( - inp=hidden_states, requires_grad=True, keep_graph=True - ) - hidden_states = node.detach(hidden_states) - offset = get_mtp_layer_offset(layer.config) + offset = get_mtp_layer_offset(layer.config, node.chunk_state.model.vp_stage) node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0)) hidden_states = node.chunk_state.mtp_hidden_states[offset] diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index b61a544c5df..8615298a1e2 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -515,7 +515,7 @@ def get_mlp_module_spec_for_backend( ) -def get_gpt_decoder_block_spec( +def get_gpt_decoder_layer_specs( config: TransformerConfig, use_transformer_engine: bool, normalization: Optional[str] = None, @@ -607,6 +607,21 @@ def get_gpt_decoder_block_spec( else: raise ValueError(f"Invalid layer pattern: {moe_layer_pattern}") + return layer_specs + + +def get_gpt_decoder_block_spec( + config: TransformerConfig, + use_transformer_engine: bool, + normalization: Optional[str] = None, + qk_l2_norm: Optional[bool] = False, + vp_stage: Optional[int] = None, + pp_rank: Optional[int] = None, +) -> TransformerBlockSubmodules: + """GPT block spec.""" + layer_specs = get_gpt_decoder_layer_specs( + config, use_transformer_engine, normalization, qk_l2_norm + ) # Slice the layer specs to only include the layers that are built in this pipeline stage. # Note: MCore layer_number starts at 1 num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage, pp_rank=pp_rank) @@ -624,6 +639,10 @@ def get_gpt_decoder_block_spec( offset = get_transformer_layer_offset(config, vp_stage=vp_stage, pp_rank=pp_rank) local_layer_specs = layer_specs[offset : offset + num_layers_to_build] + if use_transformer_engine: + layer_norm_impl = TENorm + else: + layer_norm_impl = LNImpl # Block spec. block_spec = TransformerBlockSubmodules( layer_specs=local_layer_specs, layer_norm=layer_norm_impl @@ -691,7 +710,7 @@ def get_gpt_mtp_block_spec_for_backend( mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0 mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers - offset = get_mtp_layer_offset(config) + offset = get_mtp_layer_offset(config, vp_stage=vp_stage) # split the mtp layer specs to only include the layers that are built in this pipeline stage. mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build] if len(mtp_layer_specs) > 0: diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 95854d07ec1..7b4bdf4a8c5 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -27,7 +27,6 @@ MTPLossLoggingHelper, MultiTokenPredictionBlock, roll_tensor, - tie_output_layer_state_dict, tie_word_embeddings_state_dict, ) from megatron.core.transformer.spec_utils import ModuleSpec @@ -250,7 +249,7 @@ def __init__( tp_group=self.pg_collection.tp, ) - if self.pre_process or self.post_process: + if self.pre_process or self.post_process or self.mtp_process: self.setup_embeddings_and_output_layer() if has_config_logger_enabled(self.config): @@ -530,7 +529,6 @@ def _postprocess( output_weight = None if self.share_embeddings_and_output_weights: output_weight = self.shared_embedding_or_output_weight() - if mtp_in_postprocess: hidden_states = self.mtp( input_ids=input_ids, @@ -550,7 +548,7 @@ def _postprocess( if not self.post_process: return hidden_states - if self.mtp_process: + if self.config.mtp_num_layers is not None: mtp_labels = labels.clone() hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0) hidden_states = hidden_states_list[0] @@ -602,6 +600,7 @@ def _postprocess( hidden_states, mtp_loss_scale * mtp_loss / num_tokens ) sequence_parallel_override = False + if in_inference_mode and inference_context.materialize_only_last_token_logits: if inference_context.is_static_batching(): hidden_states = hidden_states[-1:, :, :] @@ -761,14 +760,11 @@ def sharded_state_dict( output_extra_state and output_extra_state.data ), f'Expected output layer extra state to be empty, got: {output_extra_state}' - # Multi-Token Prediction (MTP) need both embedding layer and output layer in - # mtp process stage. + # Multi-Token Prediction (MTP) need embedding layer in mtp process stage. # If MTP is not placed in the pre processing stage, we need to maintain a copy of # embedding layer in the mtp process stage and tie it to the embedding in the pre # processing stage. - # Also, if MTP is not placed in the post processing stage, we need to maintain a copy - # of output layer in the mtp process stage and tie it to the output layer in the post - # processing stage. + # Now MTP loss is computed in post processing stage, so the output_layer is not needed. if self.mtp_process and not self.pre_process: emb_weight_key = f'{prefix}embedding.word_embeddings.weight' emb_weight = self.embedding.word_embeddings.weight @@ -779,19 +775,5 @@ def sharded_state_dict( tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], ) - if self.mtp_process and not self.post_process: - # We only need to tie the output layer weight if share_embeddings_and_output_weights - # is False. Because if share_embeddings_and_output_weights is True, the shared weight - # will be stored in embedding layer, and output layer will not have any weight. - if not self.share_embeddings_and_output_weights: - output_layer_weight_key = f'{prefix}output_layer.weight' - output_layer_weight = self.output_layer.weight - tie_output_layer_state_dict( - sharded_state_dict, - output_layer_weight, - output_layer_weight_key, - tp_group=self.tp_group, - dp_cp_group=metadata['dp_cp_group'], - ) return sharded_state_dict diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py index 63ee9d1f537..ac839c21f18 100644 --- a/megatron/core/pipeline_parallel/p2p_communication.py +++ b/megatron/core/pipeline_parallel/p2p_communication.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from typing import List, Optional, Tuple, Union @@ -214,22 +214,22 @@ def _communicate_shapes(self, tensor_send_next, tensor_send_prev, recv_prev, rec ops = [] if send_prev_shape_tensor is not None: send_prev_op = torch.distributed.P2POp( - torch.distributed.isend, send_prev_shape_tensor, self.prev_rank + torch.distributed.isend, send_prev_shape_tensor, self.prev_rank, self.pp_group ) ops.append(send_prev_op) if recv_prev_shape_tensor is not None: recv_prev_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank + torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank, self.pp_group ) ops.append(recv_prev_op) if send_next_shape_tensor is not None: send_next_op = torch.distributed.P2POp( - torch.distributed.isend, send_next_shape_tensor, self.next_rank + torch.distributed.isend, send_next_shape_tensor, self.next_rank, self.pp_group ) ops.append(send_next_op) if recv_next_shape_tensor is not None: recv_next_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_next_shape_tensor, self.next_rank + torch.distributed.irecv, recv_next_shape_tensor, self.next_rank, self.pp_group ) ops.append(recv_next_op) if len(ops) > 0: @@ -298,13 +298,13 @@ def _communicate( tensor_recv_prev_func = None tensor_recv_next_func = None - if not config.variable_seq_lengths: - recv_prev_shape = tensor_shape - recv_next_shape = tensor_shape - else: + if config.variable_seq_lengths or config.mtp_standalone: recv_prev_shape, recv_next_shape = self._communicate_shapes( tensor_send_next, tensor_send_prev, recv_prev, recv_next ) + else: + recv_prev_shape = tensor_shape + recv_next_shape = tensor_shape def create_tensor_recv_prev(): return torch.empty( diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 124430e107a..3b62be7dd64 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -14,17 +14,17 @@ from megatron.core.fp8_utils import get_fp8_context from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.pipeline_parallel.utils import is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel import ( gather_from_tensor_model_parallel_region, scatter_to_sequence_parallel_region, ) -from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_block import TransformerBlockSubmodules from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from megatron.core.utils import ( get_pg_rank, is_torch_min_version, @@ -444,25 +444,100 @@ def get_mtp_layer_spec_for_backend( return mtp_layer_spec -def get_mtp_layer_offset(config: TransformerConfig) -> int: +def mtp_on_this_rank( + config: TransformerConfig, ignore_virtual: Optional[bool] = True, vp_stage: Optional[int] = None +) -> bool: + """ + Check if there is MTP on the current rank. + + Behavior: + - If a custom pipeline model parallel layout is provided in the config: + - If virtual pipeline parallelism is enabled (and `ignore_virtual` is False), checks + whether any MTP layers are present on this (pp_rank, vp_stage) pair. + - Otherwise, checks all virtual pipeline ranks of the current pipeline rank. Returns + True if any virtual sub-rank includes at least one MTP layer. + - If no custom layout is provided, assumes all MTP layers (if any) are placed on the last + pipeline stage. The function returns True only on the last pipeline stage. + """ + mtp_on_this_rank = False + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + if config.pipeline_model_parallel_layout is not None: + # with custom PP layout, we support put MTP layers on any pipeline stage + layout = config.pipeline_model_parallel_layout.layout + if ( + not ignore_virtual + and parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None + ): + assert vp_stage is not None, "vp_stage must be passed if virtual pipeline is enabled" + num_layers_to_build = layout[pp_rank][vp_stage].count(LayerType.mtp) + mtp_on_this_rank = num_layers_to_build > 0 + else: + for vpp_rank in range(len(layout[pp_rank])): + num_layers_to_build = layout[pp_rank][vpp_rank].count(LayerType.mtp) + if num_layers_to_build > 0: + mtp_on_this_rank = True + break + else: + # without custom PP layout, we only support put all of MTP layers on the last pipeline stage + if config.mtp_num_layers is not None: + mtp_on_this_rank = parallel_state.is_pipeline_last_stage( + ignore_virtual=ignore_virtual, vp_stage=vp_stage + ) + else: + mtp_on_this_rank = False + return mtp_on_this_rank + + +def get_mtp_ranks(pp_ranks: List[int], config: TransformerConfig) -> List[int]: + """Get the ranks of the MTP layers.""" + mtp_ranks = set() + if config.mtp_num_layers is None: + return [] + if config.pipeline_model_parallel_layout is None: + return [pp_ranks[-1]] + layout = config.pipeline_model_parallel_layout.layout + for pp_rank in range(len(layout)): + for vpp_rank in range(len(layout[pp_rank])): + num_layers_to_build = layout[pp_rank][vpp_rank].count(LayerType.mtp) + if num_layers_to_build: + mtp_ranks.add(pp_ranks[pp_rank]) + return list(mtp_ranks) + + +def get_mtp_layer_offset(config: TransformerConfig, vp_stage: Optional[int] = None) -> int: """Get the offset of the MTP layer.""" - # Currently, we only support put all of MTP layers on the last pipeline stage. - return 0 + if config.pipeline_model_parallel_size > 1: + if config.pipeline_model_parallel_layout: + offset = config.pipeline_model_parallel_layout.get_layer_offset( + layer_type=LayerType.mtp, vp_stage=vp_stage + ) + else: + offset = 0 + else: + offset = 0 + return offset def get_mtp_num_layers_to_build( config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None ) -> int: """Get the number of MTP layers to build.""" - # Currently, we only support put all of MTP layers on the last pipeline stage. - vp_size = config.virtual_pipeline_model_parallel_size - if pp_rank is None: - pp_rank = parallel_state.get_pipeline_model_parallel_rank() - is_last_pp_stage = pp_rank == config.pipeline_model_parallel_size - 1 - if is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) and is_last_pp_stage: - return config.mtp_num_layers if config.mtp_num_layers else 0 + if config.pipeline_model_parallel_layout is not None: + # If we have a custom PP layout, get the number of mtp layers in the layout array. + num_layers_to_build = config.pipeline_model_parallel_layout.get_num_layers_to_build( + layer_type=LayerType.mtp, vp_stage=vp_stage + ) + assert num_layers_to_build == config.mtp_num_layers or num_layers_to_build == 0, ( + f"Currently, we only support put all of MTP layers on the last pipeline stage, " + f"so the number of MTP layers to build ({num_layers_to_build}) must match " + f"mtp_num_layers ({config.mtp_num_layers}) or be 0." + ) else: - return 0 + if parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage): + num_layers_to_build = config.mtp_num_layers if config.mtp_num_layers else 0 + else: + num_layers_to_build = 0 + return num_layers_to_build class MTPLossAutoScaler(torch.autograd.Function): @@ -542,7 +617,7 @@ def __init__( super().__init__(config=config) self.sequence_parallel = config.sequence_parallel self.submodules = submodules - self.layer_number = layer_number + self.layer_number = layer_number + get_mtp_layer_offset(self.config, vp_stage) self.vp_stage = vp_stage self.cp_group = pg_collection.cp @@ -584,8 +659,15 @@ def __init__( skip_bias_add=False, is_expert=False, ) + + diff_transformer_layer_offset = self.config.num_layers - get_transformer_layer_offset( + self.config, vp_stage + ) self.transformer_layer = build_module( - self.submodules.transformer_layer, config=self.config, vp_stage=vp_stage + self.submodules.transformer_layer, + config=self.config, + vp_stage=vp_stage, + layer_number=self.layer_number + diff_transformer_layer_offset, ) self.final_layernorm = build_module( @@ -1023,7 +1105,7 @@ def forward( (Tensor): The mtp loss tensor of shape [b, s]. """ # get hidden states from previous mtp stages - offset = get_mtp_layer_offset(self.config) + offset = get_mtp_layer_offset(self.config, self.vp_stage) hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0)) hidden_states = hidden_states_list[offset] for layer_number in range(len(self.layers)): @@ -1068,7 +1150,7 @@ def sharded_state_dict( sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) layer_prefix = f'{prefix}layers.' for layer in self.layers: - offset = get_mtp_layer_offset(self.config) + offset = get_mtp_layer_offset(self.config, self.vp_stage) sharded_prefix = f'{layer_prefix}{layer.layer_number - 1 }.' state_dict_prefix = f'{layer_prefix}{layer.layer_number - 1 - offset}.' diff --git a/megatron/core/transformer/pipeline_parallel_layer_layout.py b/megatron/core/transformer/pipeline_parallel_layer_layout.py index 56467bf0e9d..7a8195e1bee 100644 --- a/megatron/core/transformer/pipeline_parallel_layer_layout.py +++ b/megatron/core/transformer/pipeline_parallel_layer_layout.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy import logging @@ -127,15 +127,28 @@ def validate_layer_layout(self, num_layers: int, mtp_num_layers: int): if LayerType.mtp in self.layout[pp_rank][-1]: assert ( self.layout[pp_rank][-1].count(LayerType.mtp) == mtp_num_layers - ), "All of the MTP layers must be in the same stage" - assert ( - pp_rank == self.pipeline_model_parallel_size - 1 - and LayerType.loss in self.layout[pp_rank][-1] - ), "MTP layers must be in the last stage together with Loss stage." + ), "All of the MTP layers must be in the same one virtual pipeline stage" + for vpp_rank in range(self.virtual_pipeline_model_parallel_size - 1): + assert LayerType.mtp not in self.layout[0][vpp_rank], ( + f"Currently we restrict that the MTP should not be in the first pp rank." + f"But got {self.layout[0]} for the first pp rank." + ) + ## Detect MTP standalone usage. + mtp_standalone = False + for pp_rank in range(self.pipeline_model_parallel_size): + if ( + LayerType.mtp in self.layout[pp_rank][-1] + and pp_rank != self.pipeline_model_parallel_size - 1 + ): + mtp_standalone = True + break + # TODO: remove them in the future once they are supported if self.flatten_layout.count(LayerType.encoder) > 0: raise NotImplementedError("Encoder layer is not supported for flexible pipeline layout") + return mtp_standalone + def get_num_layers_to_build( self, layer_type: LayerType = LayerType.decoder, diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index b44fe75898b..b16d88a83cd 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -374,7 +374,7 @@ def build_layer(layer_spec, layer_number): # @TODO: add back account_for_embedding_in_pipeline_split (see issue #293) # In pipeline parallelism, we want to add this LN only to the last stage of the pipeline # self.post_process and self.post_layer_norm guide this behavior - if self.submodules.layer_norm and self.post_process and self.post_layer_norm: + if self.has_final_layernorm_in_this_stage(): self.final_layernorm = build_module( self.submodules.layer_norm, config=self.config, @@ -387,6 +387,35 @@ def build_layer(layer_spec, layer_number): if self.config.inference_fuse_tp_communication: self._setup_fused_tp_communication() + def has_final_layernorm_in_this_stage(self): + """ + Check if this vpp stage contains the final layernorm. + + Note: + Final layernorm now has been moved from the post-process stage to the last decoder + layer by using this function. + There will be a small numeric difference because of grad norm reduction when final + layernorm is placed in different pipeline stages in deterministic mode. It can still + be bitwise aligned by disabling grad norm clipping. + """ + if self.config.mtp_num_layers is None: + # for model without MTPLayer, the final layernorm is set in the stage which does + # post_process + return self.submodules.layer_norm and self.post_process and self.post_layer_norm + else: + # for model with MTPLayer, the final layernorm is set in the stage which has the + # last layer of the decoder + has_final_layernorm_in_this_stage = False + for layer in self.layers: + if layer.layer_number == self.config.num_layers: + has_final_layernorm_in_this_stage = True + break + return ( + self.submodules.layer_norm + and has_final_layernorm_in_this_stage + and self.post_layer_norm + ) + def _setup_fused_tp_communication(self): """Setup fused TP communication for all layers. We have a fused reduce-scatter + add + layer-norm + all-gather operation. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index f677dc9ac9b..56f6be32e8c 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1163,7 +1163,7 @@ def __post_init__(self): self.virtual_pipeline_model_parallel_size = detected_vpp_size # Check whether the layout is valid. - self.pipeline_model_parallel_layout.validate_layer_layout( + self.mtp_standalone = self.pipeline_model_parallel_layout.validate_layer_layout( num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers ) diff --git a/megatron/training/utils.py b/megatron/training/utils.py index 669f1972f02..33983ca372d 100644 --- a/megatron/training/utils.py +++ b/megatron/training/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """General utilities.""" import json @@ -511,7 +511,7 @@ def get_blend_and_blend_per_split(args): return blend, blend_per_split -def get_batch_on_this_tp_rank(data_iterator): +def get_batch_on_this_tp_rank(data_iterator, mtp_on_this_rank: bool = False): args = get_args() @@ -539,7 +539,7 @@ def _broadcast(item): 'position_ids': data["position_ids"].cuda(non_blocking=True), } - if args.pipeline_model_parallel_size == 1: + if args.pipeline_model_parallel_size == 1 or mtp_on_this_rank: _broadcast(batch['tokens']) _broadcast(batch['labels']) _broadcast(batch['loss_mask']) @@ -555,9 +555,6 @@ def _broadcast(item): # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - if args.mtp_num_layers is not None: - _broadcast(batch['tokens']) - _broadcast(batch['position_ids']) _broadcast(batch['labels']) _broadcast(batch['loss_mask']) _broadcast(batch['attention_mask']) @@ -593,7 +590,7 @@ def _broadcast(item): device=torch.cuda.current_device(), ) - if args.pipeline_model_parallel_size == 1: + if args.pipeline_model_parallel_size == 1 or mtp_on_this_rank: _broadcast(tokens) _broadcast(labels) _broadcast(loss_mask) @@ -612,12 +609,8 @@ def _broadcast(item): # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - if args.mtp_num_layers is not None: - _broadcast(tokens) - _broadcast(position_ids) - else: - tokens = None - position_ids = None + tokens = None + position_ids = None _broadcast(labels) _broadcast(loss_mask) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 142261e7eee..f7d918f03a7 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain and SFT GPT.""" @@ -19,6 +19,8 @@ from megatron.core.utils import StragglerDetector, get_attr_wrapped_model from megatron.training import get_args, get_timers, get_tokenizer, inprocess_restart, pretrain, print_rank_0 from megatron.training.datasets.sft_dataset import SFTDataset +from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank, get_mtp_ranks +from megatron.training.arguments import core_transformer_config_from_args from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig from megatron.training.utils import ( get_batch_on_this_cp_rank, @@ -39,14 +41,20 @@ stimer = StragglerDetector() -def get_batch(data_iterator, vp_stage=None): +def get_batch(data_iterator, vp_stage: Optional[int] = None): """Generate a batch.""" + args = get_args() + config = core_transformer_config_from_args(args) # TODO: this is pretty hacky, find a better way - if not is_first_or_last_pipeline_stage(vp_stage): + if not is_first_or_last_pipeline_stage(vp_stage) and ( + (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))): return None, None, None, None, None # get batches based on the TP rank you are on - batch = get_batch_on_this_tp_rank(data_iterator) + batch = get_batch_on_this_tp_rank( + data_iterator, + mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) + ) # slice batch along sequence dimension for context parallelism batch = get_batch_on_this_cp_rank(batch) @@ -160,7 +168,12 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa def is_dataset_built_on_rank(vp_stage=None): - return is_first_or_last_pipeline_stage(vp_stage) and parallel_state.get_tensor_model_parallel_rank() == 0 + args = get_args() + config = core_transformer_config_from_args(args) + return ( + is_first_or_last_pipeline_stage(vp_stage) + or mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) + ) and parallel_state.get_tensor_model_parallel_rank() == 0 def core_gpt_dataset_config_from_args(args): @@ -249,6 +262,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") + is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage) train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( dataset_type, train_val_test_num_samples, partial(is_dataset_built_on_rank, vp_stage=vp_stage), config ).build() @@ -258,6 +272,21 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None return train_ds, valid_ds, test_ds +def get_embedding_ranks(pp_ranks: List[int]): + """Get the embedding ranks.""" + embedding_ranks = [pp_ranks[0]] + if len(pp_ranks) > 1: + args = get_args() + if not args.untie_embeddings_and_output_weights: + embedding_ranks.append(pp_ranks[-1]) + config = core_transformer_config_from_args(args) + mtp_ranks = get_mtp_ranks(pp_ranks, config) + embedding_ranks.extend(mtp_ranks) + embedding_ranks = list(set(embedding_ranks)) + embedding_ranks = sorted(embedding_ranks) + return embedding_ranks + + if __name__ == "__main__": # Temporary for transition to core datasets @@ -274,4 +303,5 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, store=store, + get_embedding_ranks=get_embedding_ranks, ) diff --git a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py index faf83837a23..30bed32bf0b 100644 --- a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py +++ b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os from pathlib import Path @@ -21,6 +21,7 @@ ) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.checkpointing import load_checkpoint, save_checkpoint from megatron.training.global_vars import set_args @@ -53,6 +54,8 @@ def initialize_gpt_model( virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, hidden_dropout=0.0, attention_dropout=0.0, + mtp_num_layers=1 if with_mtp else None, + mtp_loss_scaling_factor=1.0 if with_mtp else None, ) default_config_kwargs.update(**config_kwargs) transformer_config = TransformerConfig(**default_config_kwargs) @@ -61,9 +64,6 @@ def initialize_gpt_model( transformer_config.moe_ffn_hidden_size = 128 transformer_config.num_moe_experts = 4 transformer_config.add_bias_linear = False - if with_mtp: - transformer_config.mtp_num_layers = 1 - transformer_config.mtp_loss_scaling_factor = 1.0 model = [] for i in range(virtual_pipeline_model_parallel_size or 1): if is_moe: @@ -71,8 +71,11 @@ def initialize_gpt_model( else: layer_spec = layer_spec_fn() - if is_moe and with_mtp and mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=i): - transformer_layer_spec_for_mtp = gpt_te_spec(transformer_config) + if with_mtp and mtp_on_this_rank(transformer_config, ignore_virtual=False, vp_stage=i): + if is_moe: + transformer_layer_spec_for_mtp = gpt_te_spec(transformer_config) + else: + transformer_layer_spec_for_mtp = layer_spec mtp_block_spec = get_gpt_mtp_block_spec( transformer_config, transformer_layer_spec_for_mtp, @@ -81,6 +84,7 @@ def initialize_gpt_model( ) else: mtp_block_spec = None + pre_process = mpu.is_pipeline_first_stage(ignore_virtual=False, vp_stage=i) post_process = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=i) this_model = ( @@ -162,7 +166,7 @@ def create_args(): [], ["decoder"], ["decoder"], - ["decoder"] * 2 + ["loss"], + ["decoder"] * 2 + ["mtp"] + ["loss"], ], False, True, @@ -184,7 +188,19 @@ def create_args(): False, ), ((1, 2, None), [["embedding"] + ["decoder"] * 4, ["decoder"] * 4 + ["loss"]], True, False), - ((1, 4, 2), "E|t*3|(t|)*5L", True, True), + ((1, 4, 2), "E|t*3|(t|)*5mL", True, True), # mtp in the last stage + ( + (1, 4, 2), + "E|t*3|(t|)*4tm|L", + True, + True, + ), # mtp in the second last stage with a decoder layer + ( + (1, 4, 2), + "E|t*3|(t|)*3tt|m|L", + True, + True, + ), # mtp in the second last stage with no other layers ], ) def test_forward_vpp(create_args, tmp_path_dist_ckpt, tp_pp_vpp, pp_layout, is_moe, with_mtp):