From 2456ea23a2db5db17ebbc099f6f2f548072e792c Mon Sep 17 00:00:00 2001 From: Haoran Zhang Date: Thu, 10 Sep 2026 21:17:30 -0700 Subject: [PATCH 1/2] Support CP in MIMO (cherry picked from commit e597a976901c3f5c2c71061f30f3c1dfeadfc3d5) Signed-off-by: Haoran Zhang --- examples/mimo/model_providers/llava_avlm.py | 7 +- examples/mimo/model_providers/llava_vlm.py | 5 +- examples/mimo/model_providers/mock.py | 5 +- examples/mimo/train.py | 29 +- examples/mimo/training/builder.py | 1 + examples/mimo/training/data.py | 11 +- examples/mimo/utils/hetero.py | 9 + megatron/core/context_parallel/utils.py | 44 ++- megatron/core/models/mimo/model/base.py | 254 ++++++--------- megatron/core/models/mimo/partition/utils.py | 273 +++++++---------- .../transformer/multi_token_prediction.py | 4 +- megatron/core/utils.py | 21 +- tests/unit_tests/data/test_get_batch.py | 43 ++- .../models/mimo/test_mimo_1f1b_schedule.py | 6 +- .../mimo/test_mimo_embedding_alignment.py | 5 +- .../models/mimo/test_mimo_mock_data.py | 11 +- .../unit_tests/models/mimo/test_mimo_model.py | 111 ++++--- .../models/mimo/test_mimo_partition.py | 289 ++++++++++-------- .../test_bridge_communicator.py | 21 ++ .../test_multi_token_prediction.py | 6 + 20 files changed, 617 insertions(+), 538 deletions(-) diff --git a/examples/mimo/model_providers/llava_avlm.py b/examples/mimo/model_providers/llava_avlm.py index bc1d554f9e1..a1c207f96ff 100644 --- a/examples/mimo/model_providers/llava_avlm.py +++ b/examples/mimo/model_providers/llava_avlm.py @@ -157,7 +157,10 @@ def model_provider_llava_avlm( # Create MIMO model cp_group = pg_collection.cp if pg_collection is not None else None tp_group = pg_collection.tp if pg_collection is not None else None - mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group) + tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None + mimo_model = MimoModel( + mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group + ) print("*"*100) print_mimo_structure(mimo_model) print("*"*100) @@ -184,4 +187,4 @@ def model_provider_llava_avlm( for param in module.parameters(): param.requires_grad = False - return mimo_model \ No newline at end of file + return mimo_model diff --git a/examples/mimo/model_providers/llava_vlm.py b/examples/mimo/model_providers/llava_vlm.py index e2de7e3f42c..6fd0678219d 100644 --- a/examples/mimo/model_providers/llava_vlm.py +++ b/examples/mimo/model_providers/llava_vlm.py @@ -129,7 +129,10 @@ def model_provider_llava_vlm( # Create MIMO model cp_group = pg_collection.cp if pg_collection is not None else None tp_group = pg_collection.tp if pg_collection is not None else None - mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group) + tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None + mimo_model = MimoModel( + mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group + ) print("*"*100) print_mimo_structure(mimo_model) print("*"*100) diff --git a/examples/mimo/model_providers/mock.py b/examples/mimo/model_providers/mock.py index c4a71e2bf13..ef9704d560e 100644 --- a/examples/mimo/model_providers/mock.py +++ b/examples/mimo/model_providers/mock.py @@ -97,6 +97,9 @@ def model_provider_mock_vlm_single_encoder( # Create MIMO model cp_group = pg_collection.cp if pg_collection is not None else None tp_group = pg_collection.tp if pg_collection is not None else None - mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group) + tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None + mimo_model = MimoModel( + mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group + ) return mimo_model diff --git a/examples/mimo/train.py b/examples/mimo/train.py index c1f9135d178..f8c30370e16 100644 --- a/examples/mimo/train.py +++ b/examples/mimo/train.py @@ -53,6 +53,29 @@ "llava_avlm": llava_avlm_dataloader_provider, } + +def _canonicalize_packed_sequence_metadata(batch: Dict[str, Any]) -> Dict[str, Any]: + """Convert the legacy Energon packing dictionary to canonical MIMO batch fields.""" + packing_kwargs = batch.pop("packing_kwargs", None) + if packing_kwargs is None: + return batch + + cu_seqlens = packing_kwargs["cu_seqlens_q"].to(dtype=torch.int32) + batch["cu_seqlens"] = cu_seqlens.unsqueeze(0) if cu_seqlens.dim() == 1 else cu_seqlens + + cu_seqlens_padded = packing_kwargs.get("cu_seqlens_q_padded") + if cu_seqlens_padded is not None: + cu_seqlens_padded = cu_seqlens_padded.to(dtype=torch.int32) + batch["cu_seqlens_padded"] = ( + cu_seqlens_padded.unsqueeze(0) + if cu_seqlens_padded.dim() == 1 + else cu_seqlens_padded + ) + + batch["max_seqlen"] = packing_kwargs["max_seqlen_q"].to(dtype=torch.int32).reshape(1) + return batch + + def add_mimo_args(parser): """Add MIMO-specific arguments to the parser.""" group = parser.add_argument_group('MIMO', 'MIMO specific arguments') @@ -131,12 +154,14 @@ def get_batch(data_iterator: Iterator[Dict[str, Any]]): # loss_mask: Optional[torch.Tensor] = None, # labels: Optional[torch.Tensor] = None, # modality_inputs: Optional[Dict[str, Dict[str, Any]]] = None, - # packing_kwargs: Optional[dict] = None, + # cu_seqlens: Optional[torch.Tensor] = None, + # cu_seqlens_padded: Optional[torch.Tensor] = None, + # max_seqlen: Optional[torch.Tensor] = None, # For the modality inputs, the keys can be arbitrary # so we do a broadcast of the schema followed by a broadcast of the actual data # check broadcast_nested_data_batch for more details - batch = broadcast_nested_data_batch(data) + batch = _canonicalize_packed_sequence_metadata(broadcast_nested_data_batch(data)) return batch diff --git a/examples/mimo/training/builder.py b/examples/mimo/training/builder.py index f7ea4c62efd..f07d518a267 100644 --- a/examples/mimo/training/builder.py +++ b/examples/mimo/training/builder.py @@ -120,6 +120,7 @@ def build_model( mimo_config, cp_group=active_pg.cp if is_language else None, tp_group=active_pg.tp if is_language else None, + tp_cp_group=active_pg.tp_cp if is_language else None, ) def build_distributed_models( diff --git a/examples/mimo/training/data.py b/examples/mimo/training/data.py index b0d7d5349d0..6efd3ceec4a 100644 --- a/examples/mimo/training/data.py +++ b/examples/mimo/training/data.py @@ -12,10 +12,10 @@ from torch.utils.data import DataLoader, Dataset from examples.mimo.training.topology import HeteroTopology +from examples.mimo.utils.hetero import get_data_lane_rank from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage -from megatron.core.utils import get_pg_rank from megatron.training.global_vars import get_run_config _ENCODER_SEED_OFFSET = 10_000 @@ -348,14 +348,7 @@ def _build_split_loaders( ) -> tuple[DataLoader, DataLoader, DataLoader]: """Build split-local datasets with deterministic module/DP/split seeds.""" cfg = get_run_config() - data_group = pg_collection.dp_cp_gtp_remat or pg_collection.dp - lane_rank = get_pg_rank(data_group) - if pg_collection.dp_cp_gtp_remat is not None: - # The combined group orders CP first (fastest), then GTP and DP. - # CP replicas consume the same full batch before the model shards it; - # GTP and DP remain distinct data lanes, matching the bridge topology. - lane_rank //= pg_collection.cp.size() - base_seed = cfg.rng.seed + module_seed_offset + lane_rank + base_seed = cfg.rng.seed + module_seed_offset + get_data_lane_rank(pg_collection) common = _mock_loader_kwargs(args, encoder_name) return tuple( _build_mock_vlm_dataloader( diff --git a/examples/mimo/utils/hetero.py b/examples/mimo/utils/hetero.py index 6c67d6da9bc..9e33946b48a 100644 --- a/examples/mimo/utils/hetero.py +++ b/examples/mimo/utils/hetero.py @@ -5,6 +5,7 @@ from __future__ import annotations from megatron.core.hyper_comm_grid import HyperCommGrid +from megatron.core.process_groups_config import ProcessGroupCollection def get_grid_dim_size(grid: HyperCommGrid, dim: str) -> int: @@ -13,3 +14,11 @@ def get_grid_dim_size(grid: HyperCommGrid, dim: str) -> int: return int(grid.shape[grid.dim_names.index(dim)]) except (ValueError, AttributeError): return 1 + + +def get_data_lane_rank(pg_collection: ProcessGroupCollection) -> int: + """Return the DP x GTP data-lane rank, excluding model-parallel dimensions.""" + gtp_group = getattr(pg_collection, "gtp_remat", None) + gtp_size = gtp_group.size() if gtp_group is not None else 1 + gtp_rank = gtp_group.rank() if gtp_group is not None else 0 + return pg_collection.dp.rank() * gtp_size + gtp_rank diff --git a/megatron/core/context_parallel/utils.py b/megatron/core/context_parallel/utils.py index 16b3352d0d0..dfff70389c4 100644 --- a/megatron/core/context_parallel/utils.py +++ b/megatron/core/context_parallel/utils.py @@ -59,7 +59,14 @@ def _get_batch_on_this_cp_rank_contiguous( cp_size = torch.distributed.get_world_size(cp_group) cp_rank = torch.distributed.get_rank(cp_group) - sequence_keys = ('tokens', 'labels', 'loss_mask', 'position_ids') + sequence_keys = ( + 'tokens', + 'labels', + 'loss_mask', + 'position_ids', + 'decoder_input', + 'mtp_input_mask', + ) if cp_size == 1: return batch @@ -108,7 +115,15 @@ def _get_batch_on_this_cp_rank_padded_zigzag( return batch sequence_tensor = None - for key in ('tokens', 'labels', 'loss_mask', 'position_ids'): + sequence_keys = ( + 'tokens', + 'labels', + 'loss_mask', + 'position_ids', + 'decoder_input', + 'mtp_input_mask', + ) + for key in sequence_keys: sequence_tensor = batch.get(key) if sequence_tensor is not None: break @@ -127,11 +142,12 @@ def _get_batch_on_this_cp_rank_padded_zigzag( index = rank_order_indices.view(cp_size, -1)[cp_rank] valid_index = index.clamp_min(0) padding = (index < 0) | ~source_valid.index_select(0, valid_index) - for key in ('tokens', 'labels', 'loss_mask', 'position_ids'): + for key in sequence_keys: tensor = batch.get(key) if tensor is not None: local_tensor = tensor.index_select(1, valid_index) - batch[key] = local_tensor.masked_fill(padding.view(1, -1), 0) + padding_shape = (1, -1, *((1,) * (tensor.dim() - 2))) + batch[key] = local_tensor.masked_fill(padding.view(padding_shape), 0) batch['cu_seqlens_padded'] = target_cu_seqlens_padded.unsqueeze(0) if batch.get('max_seqlen') is not None: max_seqlen = (target_cu_seqlens_padded[1:] - target_cu_seqlens_padded[:-1]).max() @@ -185,7 +201,7 @@ def get_batches_on_this_cp_rank( batch: Dict[str, Any], boundary_layout: CPLayout, is_hybrid_cp: bool, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, additional_layouts: Iterable[CPLayout] = (), hybrid_cp_group_func: Callable[[int], torch.distributed.ProcessGroup] | None = None, use_per_sequence_balancing: bool = False, @@ -205,12 +221,14 @@ def get_batches_on_this_cp_rank( sequence. The same rank ordering is used to build the zigzag batch tensors and their ``PackedSeqParams``. When both layouts are requested, it also defines the activation-conversion plan. All other cases use the standard batch sharder. + + A ``None`` CP group represents a size-one CP domain and leaves sequence tensors unchanged. """ from megatron.core.utils import get_batch_on_this_cp_rank requested_layouts = set(additional_layouts) requested_layouts.add(boundary_layout) - cp_size = torch.distributed.get_world_size(cp_group) + cp_size = 1 if cp_group is None else torch.distributed.get_world_size(cp_group) build_packed_zigzag_view = ( not is_hybrid_cp @@ -279,9 +297,17 @@ def get_batches_on_this_cp_rank( has_sequence_data = any( batch.get(key) is not None - for key in ('tokens', 'labels', 'loss_mask', 'position_ids', 'attention_mask') + for key in ( + 'tokens', + 'labels', + 'loss_mask', + 'position_ids', + 'attention_mask', + 'decoder_input', + 'mtp_input_mask', + ) ) - if has_sequence_data: + if has_sequence_data and cp_size > 1: # Copy the dictionary because the CP sharder replaces sequence-valued entries in place. batches_by_layout = { layout: get_batch_on_this_cp_rank( @@ -295,7 +321,7 @@ def get_batches_on_this_cp_rank( for layout in requested_layouts } else: - # Intermediate PP stages receive activations rather than token-aligned tensors. + # CP size one needs no sharding; intermediate PP stages may also have metadata only. batches_by_layout = {layout: dict(batch) for layout in requested_layouts} # No shared conversion route exists in this path, so build metadata from each physical batch diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py index b95a58d7798..b33da274867 100644 --- a/megatron/core/models/mimo/model/base.py +++ b/megatron/core/models/mimo/model/base.py @@ -7,13 +7,13 @@ import torch from megatron.core._rank_utils import warn_single_rank +from megatron.core.context_parallel import ContextParallelBatch from megatron.core.dist_checkpointing.utils import apply_prefix_mapping from megatron.core.distributed import DistributedDataParallel from megatron.core.models.mimo.comm.colocated_communicator import ColocatedBridgeCommunicator from megatron.core.models.mimo.config import MimoModelConfig from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY, ModuleLayout, RankRole from megatron.core.models.mimo.partition.utils import PartitionAdapter, PartitionConfig -from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.quantization.utils import get_quant_config_or_none from megatron.core.transformer import MegatronModule from megatron.core.transformer.module import Float16Module @@ -56,7 +56,9 @@ class MimoModel(MegatronModule): Configuration for the model, including language model and modality submodules """ - def __init__(self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None) -> None: + def __init__( + self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None, tp_cp_group=None + ) -> None: """Initialize the multimodal model. Example: @@ -94,15 +96,14 @@ def __init__(self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None) - self.partition_adapter: Optional[PartitionAdapter] = None # Only on language-module ranks: encoder-only ranks never shard and would read # process groups they do not own. - if self.role.has_language_module and ( - language_config.context_parallel_size > 1 or language_config.sequence_parallel - ): + if self.role.has_language_module: partition_config = PartitionConfig.from_mp_config( mp=language_config, max_seq_len=max_seq_len, kv_format=mimo_config.kv_format, cp_group=cp_group, tp_group=tp_group, + tp_cp_group=tp_cp_group, ) self.partition_adapter = PartitionAdapter(partition_config) @@ -629,7 +630,9 @@ def forward( loss_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, modality_inputs: Optional[Dict[str, Dict[str, Any]]] = None, - packing_kwargs: Optional[dict] = None, + cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens_padded: Optional[torch.Tensor] = None, + max_seqlen: Optional[torch.Tensor] = None, modality_token_indices: Optional[Dict[str, torch.Tensor]] = None, ): """Forward pass through the multimodal model. @@ -650,20 +653,9 @@ def forward( "whisper_encoder": {"input_features": whisper_features} } } - packing_kwargs: Optional dictionary of kwargs to construct PackedSeqParams - if packed_seq_params is not provided. For example: - { - "cu_seqlens_q": cu_seqlens, - "cu_seqlens_kv": cu_seqlens, - "cu_seqlens_q_padded": cu_seqlens_padded, - "cu_seqlens_kv_padded": cu_seqlens_padded, - "max_seqlen_q": torch.tensor( - max(seqlens_padded), dtype=torch.int32 - ), - "max_seqlen_kv": torch.tensor( - max(seqlens_padded), dtype=torch.int32 - ), - } + cu_seqlens: Logical cumulative sequence lengths for a packed language batch. + cu_seqlens_padded: Physical cumulative sequence lengths including padding. + max_seqlen: Maximum physical sequence length in the packed language batch. modality_token_indices: Optional complete mapping from every active modality name, including ``text``, to flat logical row-major ``torch.long`` indices in the ``[B, S]`` input grid. Encoder-only ranks ignore this argument. Pass ``None`` to @@ -687,7 +679,9 @@ def forward( loss_mask, labels, modality_inputs, - packing_kwargs=packing_kwargs, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, modality_token_indices=modality_token_indices, ) @@ -704,7 +698,9 @@ def forward( loss_mask, labels, input_tensors, - packing_kwargs=packing_kwargs, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, modality_token_indices=modality_token_indices, ) @@ -820,44 +816,6 @@ def _empty_encoder_output(self, encoder_name: str) -> torch.Tensor: (0, hidden_size), device=torch.cuda.current_device(), dtype=language_config.params_dtype ) - def _build_packed_seq_params(self, packing_kwargs: Optional[dict]) -> Optional[PackedSeqParams]: - """Build THD ``PackedSeqParams`` from ``packing_kwargs`` (None if not packing).""" - if packing_kwargs is None: - return None - for key in packing_kwargs: - if 'cu_seqlens' in key and packing_kwargs[key] is not None: - packing_kwargs[key] = packing_kwargs[key].to(dtype=torch.int32) - packed_seq_params = PackedSeqParams(**packing_kwargs) - packed_seq_params.qkv_format = 'thd' - return packed_seq_params - - def _shard_language_inputs( - self, - embeddings: Optional[torch.Tensor], - labels: Optional[torch.Tensor], - loss_mask: Optional[torch.Tensor], - packed_seq_params: Optional[PackedSeqParams] = None, - ) -> Tuple[ - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[PackedSeqParams], - ]: - """Apply CP/SP sharding via the partition adapter, or pass through if inactive. - - ``embeddings`` are sequence-first ``(S, B, H)`` (``None`` on non-first PP stages) - and come back in ``(S/(cp*tp), B, H)``; labels/loss_mask are ``(B, S)``. - """ - if self.partition_adapter is None: - return embeddings, labels, loss_mask, packed_seq_params - - return self.partition_adapter.shard( - embeddings=embeddings, - labels=labels, - loss_mask=loss_mask, - packed_seq_params=packed_seq_params, - ) - def _language_model_owns_mtp(self) -> bool: """Return whether this rank executes the language model's MTP block.""" if self.language_model is None: @@ -881,14 +839,19 @@ def _materialize_mtp_input_mask( mtp_input_mask &= input_ids != special_token_id return mtp_input_mask - def _prepare_mtp_inputs( + def _prepare_language_inputs( self, + embeddings: Optional[torch.Tensor], input_ids: Optional[torch.Tensor], position_ids: Optional[torch.Tensor], - packed_seq_params: Optional[PackedSeqParams], + labels: Optional[torch.Tensor], + loss_mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + cu_seqlens_padded: Optional[torch.Tensor], + max_seqlen: Optional[torch.Tensor], owns_mtp: bool, text_token_indices: Optional[torch.Tensor] = None, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + ) -> ContextParallelBatch: """Prepare CP-local position IDs and optional MTP token metadata. MTP consumes token IDs only on the stage that owns its prediction block. Under @@ -906,44 +869,37 @@ def _prepare_mtp_inputs( input_ids, self.special_token_ids, text_token_indices=text_token_indices ) - if self.partition_adapter is None or not self.partition_adapter.cfg.use_cp: - return mtp_input_ids, position_ids, mtp_input_mask - - # PartitionAdapter shards batch-first [B, S, ...] metadata along dimension 1. - # Multidimensional RoPE positions arrive as [rope_dim, B, S], so expose their - # sequence dimension in the adapter's expected layout and restore it afterward. - is_multiaxis_position_ids = position_ids is not None and position_ids.dim() == 3 - position_metadata = ( - position_ids.movedim(0, -1).contiguous() if is_multiaxis_position_ids else position_ids + assert self.partition_adapter is not None + return self.partition_adapter.partition( + embeddings=embeddings, + input_ids=mtp_input_ids, + position_ids=position_ids, + labels=labels, + loss_mask=loss_mask, + mtp_input_mask=mtp_input_mask, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, ) - packed_mtp_metadata = mtp_input_ids - if mtp_input_mask is not None: - assert mtp_input_ids is not None - packed_mtp_metadata = torch.cat( - (mtp_input_ids, mtp_input_mask.to(dtype=mtp_input_ids.dtype)), dim=0 - ) - - _, local_position_ids, packed_mtp_metadata, _ = self.partition_adapter.shard( - embeddings=None, - labels=position_metadata, - loss_mask=packed_mtp_metadata, - packed_seq_params=packed_seq_params, + def _forward_prepared_language_inputs( + self, cp_batch: ContextParallelBatch, attention_mask: Optional[torch.Tensor] + ) -> Any: + """Call the language model, including the dual-layout batch when present.""" + batch = cp_batch.get_batch() + cp_kwargs = {"cp_batch": cp_batch} if len(cp_batch.batches_by_layout) > 1 else {} + return self.language_model( + input_ids=batch["tokens"], + position_ids=batch["position_ids"], + attention_mask=attention_mask, + decoder_input=batch["decoder_input"], + labels=batch["labels"], + loss_mask=batch["loss_mask"], + mtp_input_mask=batch["mtp_input_mask"], + packed_seq_params=cp_batch.get_packed_seq_params(), + **cp_kwargs, ) - if mtp_input_mask is not None: - assert packed_mtp_metadata is not None - mtp_input_ids, mtp_input_mask = packed_mtp_metadata.chunk(2, dim=0) - mtp_input_mask = mtp_input_mask.to(dtype=torch.bool) - else: - mtp_input_ids = packed_mtp_metadata - - if is_multiaxis_position_ids: - assert local_position_ids is not None - local_position_ids = local_position_ids.movedim(-1, 0).contiguous() - - return mtp_input_ids, local_position_ids, mtp_input_mask - def _forward_language_module( self, input_ids: torch.Tensor, @@ -952,7 +908,9 @@ def _forward_language_module( loss_mask: Optional[torch.Tensor], labels: Optional[torch.Tensor], input_tensors: Optional[Dict[str, torch.Tensor]], - packing_kwargs: Optional[dict] = None, + cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens_padded: Optional[torch.Tensor] = None, + max_seqlen: Optional[torch.Tensor] = None, modality_token_indices: Optional[Dict[str, torch.Tensor]] = None, ) -> Tuple[Any, Optional[torch.Tensor]]: """Forward pass for language module on this rank. @@ -966,7 +924,9 @@ def _forward_language_module( loss_mask: Loss mask for per-token loss normalization labels: Labels for loss computation input_tensors: Hidden states or embeddings from previous stage - packing_kwargs: Optional kwargs to construct packed (THD) sequence params. + cu_seqlens: Logical cumulative sequence lengths for a packed language batch. + cu_seqlens_padded: Physical cumulative sequence lengths including padding. + max_seqlen: Maximum physical sequence length in the packed language batch. modality_token_indices: Optional complete mapping of trusted flat logical row-major token indices. See ``align_embeddings_by_token_positions``. @@ -987,7 +947,6 @@ def _forward_language_module( "CP-sharded sequence)." ) - packed_seq_params = self._build_packed_seq_params(packing_kwargs) owns_mtp = self._language_model_owns_mtp() if self.role.is_first_stage(lang_name): @@ -1025,47 +984,36 @@ def _forward_language_module( modality_token_indices=modality_token_indices, ) - # Apply CP/SP sharding; combined_embeddings returns in [S/(cp*tp), B, H]. - combined_embeddings, labels, loss_mask, packed_seq_params = self._shard_language_inputs( + # Partition every LM-aligned tensor together. Decoder embeddings return in + # [S/(cp*tp), B, H]; labels/loss metadata remain [B, S/cp]. + cp_batch = self._prepare_language_inputs( embeddings=combined_embeddings, - labels=labels, - loss_mask=loss_mask, - packed_seq_params=packed_seq_params, - ) - mtp_input_ids, position_ids, mtp_input_mask = self._prepare_mtp_inputs( input_ids=input_ids, position_ids=position_ids, - packed_seq_params=packed_seq_params, + labels=labels, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, owns_mtp=owns_mtp, text_token_indices=(modality_token_indices or {}).get("text"), ) - - lm_output = self.language_model( - # decoder_input replaces the main embedding lookup, but MTP still - # needs token IDs to construct its shifted-token embeddings. - input_ids=mtp_input_ids, - position_ids=position_ids, - decoder_input=combined_embeddings, - labels=labels, - loss_mask=loss_mask, - mtp_input_mask=mtp_input_mask, - attention_mask=attention_mask, - packed_seq_params=packed_seq_params, + lm_output = self._forward_prepared_language_inputs( + cp_batch, attention_mask=attention_mask ) else: # Non-first stage: receive hidden states from previous LM stage. # Labels/loss_mask still need CP sharding so the loss on the last stage # lines up with the CP-local hidden states. - _, labels, loss_mask, packed_seq_params = self._shard_language_inputs( + cp_batch = self._prepare_language_inputs( embeddings=None, - labels=labels, - loss_mask=loss_mask, - packed_seq_params=packed_seq_params, - ) - mtp_input_ids, position_ids, mtp_input_mask = self._prepare_mtp_inputs( input_ids=input_ids, position_ids=position_ids, - packed_seq_params=packed_seq_params, + labels=labels, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, owns_mtp=owns_mtp, text_token_indices=(modality_token_indices or {}).get("text"), ) @@ -1078,19 +1026,12 @@ def _forward_language_module( if hasattr(underlying_lm, 'set_input_tensor'): underlying_lm.set_input_tensor(hidden_states) - lm_output = self.language_model( - # Hidden states arrive via set_input_tensor; position_ids is - # still consumed by mRoPE on non-first PP stages. - input_ids=mtp_input_ids, - position_ids=position_ids, - decoder_input=None, - labels=labels, - loss_mask=loss_mask, - mtp_input_mask=mtp_input_mask, - attention_mask=attention_mask, - packed_seq_params=packed_seq_params, + lm_output = self._forward_prepared_language_inputs( + cp_batch, attention_mask=attention_mask ) + loss_mask = cp_batch.get_batch().get("loss_mask") + # Key output for non-last stages so schedule can route to next LM stage if not self.role.is_last_stage(lang_name): return {lang_name: lm_output}, loss_mask @@ -1146,14 +1087,15 @@ def _forward_all_modules( loss_mask: Optional[torch.Tensor], labels: Optional[torch.Tensor], modality_inputs: Optional[Dict[str, Dict[str, Any]]], - packing_kwargs: Optional[dict] = None, + cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens_padded: Optional[torch.Tensor] = None, + max_seqlen: Optional[torch.Tensor] = None, modality_token_indices: Optional[Dict[str, torch.Tensor]] = None, ): """Forward pass when all modules are on all ranks (no multi-module PP). This is the original behavior, preserved for backward compatibility. """ - packed_seq_params = self._build_packed_seq_params(packing_kwargs) owns_mtp = self._language_model_owns_mtp() # 1. Process each modality to get embeddings @@ -1198,36 +1140,24 @@ def _forward_all_modules( ) logger.debug(f"Combined embeddings shape: {combined_embeddings.shape}") - # 3. Apply CP/SP sharding. combined_embeddings is [S, B, H] and returns - # [S/(cp*tp), B, H] for the LM (the adapter handles the CP batch-first transpose). - combined_embeddings, labels, loss_mask, packed_seq_params = self._shard_language_inputs( + # 3. Partition all LM-aligned inputs together. combined_embeddings is [S, B, H] + # and returns [S/(cp*tp), B, H] for the LM. + cp_batch = self._prepare_language_inputs( embeddings=combined_embeddings, - labels=labels, - loss_mask=loss_mask, - packed_seq_params=packed_seq_params, - ) - mtp_input_ids, position_ids, mtp_input_mask = self._prepare_mtp_inputs( input_ids=input_ids, position_ids=position_ids, - packed_seq_params=packed_seq_params, + labels=labels, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, owns_mtp=owns_mtp, text_token_indices=(modality_token_indices or {}).get("text"), ) # 5. Forward pass through language model - lm_output = self.language_model( - # decoder_input replaces the main embedding lookup, but MTP still - # needs token IDs to construct its shifted-token embeddings. - input_ids=mtp_input_ids, - position_ids=position_ids, - decoder_input=combined_embeddings, - labels=labels, - loss_mask=loss_mask, - mtp_input_mask=mtp_input_mask, - attention_mask=None, - packed_seq_params=packed_seq_params, - ) + lm_output = self._forward_prepared_language_inputs(cp_batch, attention_mask=None) logger.debug(f"Language model output shape: {lm_output.shape}") - return lm_output, loss_mask + return lm_output, cp_batch.get_batch().get("loss_mask") diff --git a/megatron/core/models/mimo/partition/utils.py b/megatron/core/models/mimo/partition/utils.py index 70e1753e500..e6043d63058 100644 --- a/megatron/core/models/mimo/partition/utils.py +++ b/megatron/core/models/mimo/partition/utils.py @@ -1,54 +1,39 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -"""Token and weight partitioning helper (CP, TP, SP). +"""Token and activation partitioning helper for MIMO language models.""" -The adapter slices sequences across *context-parallel* ranks and can further -scatter them across *sequence-parallel* ranks when sequence-parallelism is -enabled. -""" from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional import torch # type: ignore[import-not-found] from torch.distributed import ProcessGroup # type: ignore[import-not-found] from megatron.core import tensor_parallel +from megatron.core.context_parallel import ( + ContextParallelBatch, + CPLayout, + get_batches_on_this_cp_rank, +) from megatron.core.model_parallel_config import ModelParallelConfig -from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import get_context_parallel_group, get_tensor_model_parallel_group -from megatron.core.utils import ( - get_batch_on_this_cp_rank, - get_pg_rank, - get_pg_size, - is_te_min_version, -) - -try: - import transformer_engine_torch as tex # type: ignore - - _HAVE_TEX = True -except ModuleNotFoundError: # pragma: no cover - tex = None # type: ignore - _HAVE_TEX = False +from megatron.core.utils import get_pg_size @dataclass(frozen=True) class PartitionConfig: - """Minimal runtime information needed to shard inputs. - - NOTE: Always construct PartitionConfig using the provided classmethod - (from_mp_config) to ensure all fields, including cp_group and tp_group, - are set correctly. - """ + """Minimal runtime information needed to partition language-model inputs.""" seq_parallel: bool use_cp: bool tp_comm_overlap: bool max_seq_len: int kv_format: str = "sbhd" # "sbhd" | "thd" + linear_cp_layout: CPLayout = "zigzag" + attention_cp_layout: CPLayout = "zigzag" cp_group: Optional[ProcessGroup] = None tp_group: Optional[ProcessGroup] = None + tp_cp_group: Optional[ProcessGroup] = None @classmethod def from_mp_config( @@ -59,6 +44,7 @@ def from_mp_config( kv_format: str = "sbhd", cp_group: Optional[ProcessGroup] = None, tp_group: Optional[ProcessGroup] = None, + tp_cp_group: Optional[ProcessGroup] = None, ) -> "PartitionConfig": """ Creates a PartitionConfig from a ModelParallelConfig. @@ -78,157 +64,122 @@ def from_mp_config( tp_comm_overlap=mp.tp_comm_overlap, max_seq_len=max_seq_len, kv_format=kv_format, + linear_cp_layout=getattr(mp, "linear_cp_layout", "zigzag"), + attention_cp_layout=getattr(mp, "attention_cp_layout", "zigzag"), cp_group=cp_group, tp_group=tp_group, + tp_cp_group=tp_cp_group, ) class PartitionAdapter: - """Shard MIMO sequence inputs (CP/SP) and return language-model-ready embeddings.""" + """Partition all token-aligned MIMO language inputs with one shared CP layout.""" def __init__(self, cfg: PartitionConfig): - """Initialize the partition adapter. - Args: - cfg: PartitionConfig, the configuration for the partition adapter. - """ self.cfg = cfg - def shard( + def partition( self, embeddings: Optional[torch.Tensor], + input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], labels: Optional[torch.Tensor], loss_mask: Optional[torch.Tensor], - packed_seq_params: Optional[PackedSeqParams] = None, - ) -> Tuple[ - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[PackedSeqParams], - ]: - """Apply context- (CP) and sequence-parallel (SP) sharding along the sequence. - - Inputs: - - embeddings: sequence-first ``(S, B, H)``, or ``None`` on non-first PP stages. - - labels / loss_mask: batch-first ``(B, S)``. - - Returns ``(embeddings, labels, loss_mask, packed_seq_params)`` with embeddings in - ``(S/(cp*tp), B, H)`` language-model layout. CP shards every tensor along the - sequence; SP additionally scatters only the embeddings -- labels/loss_mask are not - SP-scattered because the loss runs after the TP all-gather on the full CP-local - sequence. A dense ``[B, S]`` attention mask is intentionally not handled: under CP - it cannot line up with the sharded sequence, so MIMO masks via a causal - ``attn_mask_type`` or ``packed_seq_params`` (THD) instead. + mtp_input_mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor] = None, + cu_seqlens_padded: Optional[torch.Tensor] = None, + max_seqlen: Optional[torch.Tensor] = None, + ) -> ContextParallelBatch: + """Apply CP to all inputs and SP to decoder embeddings. + + Decoder embeddings arrive sequence-first as ``[S, B, H]``. Token-aligned metadata is + batch-first as ``[B, S, ...]`` except multi-axis position IDs, which arrive as + ``[rope_dim, B, S]``. Under dual-layout CP, the boundary view follows the linear layer + layout and the second view follows the attention/MTP layout. """ - # Sanity-check the sequence length before sharding. Embeddings are sequence-first, - # so the token sequence is dim 0. - if embeddings is not None: - shard_factor = None - - if self.cfg.use_cp and self.cfg.seq_parallel: - shard_factor = get_pg_size(self.cfg.tp_group) * get_pg_size(self.cfg.cp_group) * 2 - elif self.cfg.use_cp: - shard_factor = get_pg_size(self.cfg.cp_group) * 2 - elif self.cfg.seq_parallel: - shard_factor = get_pg_size(self.cfg.tp_group) - - if shard_factor is not None and ( - packed_seq_params is None - or getattr(packed_seq_params, 'qkv_format', 'sbhd') == 'sbhd' - ): - assert embeddings.shape[0] % shard_factor == 0, ( - f"Sequence length should be divisible by {shard_factor} " - "for Sequence/Context parallelism" - ) + self._validate_sequence_length(embeddings, is_packed=cu_seqlens is not None) - if self.cfg.seq_parallel and self.cfg.tp_comm_overlap: - assert embeddings.shape[0] == self.cfg.max_seq_len, ( - "TP Comm overlap requires Vision+Text token length " - "== language_max_sequence_length" - ) - - if self.cfg.use_cp: - # CP shards batch-first (get_batch_on_this_cp_rank requirement): transpose - # (S, B, H) -> (B, S, H) in, then the CP-local result back to (S/cp, B, H). - if embeddings is not None: - embeddings = embeddings.transpose(0, 1).contiguous() - embeddings, labels, loss_mask, packed_seq_params = self._apply_context_parallel( - embeddings, labels, loss_mask, packed_seq_params - ) - if embeddings is not None: - embeddings = embeddings.transpose(0, 1).contiguous() + is_multiaxis_position_ids = ( + self.cfg.use_cp and position_ids is not None and position_ids.dim() == 3 + ) + position_metadata = ( + position_ids.movedim(0, -1).contiguous() if is_multiaxis_position_ids else position_ids + ) + batch = { + "tokens": input_ids, + "position_ids": position_metadata, + "labels": labels, + "loss_mask": loss_mask, + "mtp_input_mask": mtp_input_mask, + "decoder_input": ( + embeddings.transpose(0, 1).contiguous() + if self.cfg.use_cp and embeddings is not None + else embeddings + ), + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, + "max_seqlen": max_seqlen, + } + + additional_layouts = set() + if self.cfg.use_cp and self.cfg.linear_cp_layout != self.cfg.attention_cp_layout: + additional_layouts.add(self.cfg.attention_cp_layout) + cp_batch = get_batches_on_this_cp_rank( + batch, + boundary_layout=self.cfg.linear_cp_layout, + is_hybrid_cp=False, + cp_group=self.cfg.cp_group, + additional_layouts=additional_layouts, + sequence_parallel=self.cfg.seq_parallel, + tp_group=self.cfg.tp_group, + tp_cp_group=self.cfg.tp_cp_group, + tokens_per_sample=self.cfg.max_seq_len, + ) - # SP scatters the sequence (dim 0); the SP-only path needs no transpose. - if self.cfg.seq_parallel and embeddings is not None: - embeddings = tensor_parallel.scatter_to_sequence_parallel_region( - embeddings, group=self.cfg.tp_group + if is_multiaxis_position_ids: + for layout_batch in cp_batch.batches_by_layout.values(): + local_position_ids = layout_batch.get("position_ids") + if local_position_ids is not None: + layout_batch["position_ids"] = local_position_ids.movedim(-1, 0).contiguous() + + boundary_batch = cp_batch.get_batch() + local_embeddings = boundary_batch.get("decoder_input") + # The alternate decoder-input view is never consumed: decoder activations are converted + # between layouts after each layer instead. Drop these references before model forward. + for layout_batch in cp_batch.batches_by_layout.values(): + layout_batch.pop("decoder_input", None) + if self.cfg.use_cp and local_embeddings is not None: + local_embeddings = local_embeddings.transpose(0, 1).contiguous() + if self.cfg.seq_parallel and local_embeddings is not None: + local_embeddings = tensor_parallel.scatter_to_sequence_parallel_region( + local_embeddings, group=self.cfg.tp_group ) - - return embeddings, labels, loss_mask, packed_seq_params - - def _apply_context_parallel( - self, - embeddings: Optional[torch.Tensor], - labels: Optional[torch.Tensor], - loss_mask: Optional[torch.Tensor], - packed_seq_params: Optional[PackedSeqParams], - ) -> Tuple[ - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[PackedSeqParams], - ]: - """ - Apply context parallel (CP) sharding to input tensors. - - Args: - embeddings (Optional[torch.Tensor]): - Input embeddings tensor. Shape: (B, S, H) - labels (Optional[torch.Tensor]): - Labels tensor. Shape: (B, S) - loss_mask (Optional[torch.Tensor]): - Loss mask tensor. Shape: (B, S) - packed_seq_params (PackedSeqParams, optional): - Packed sequence parameters. Defaults to None. - - Returns: - Tuple containing: - - embeddings (Optional[torch.Tensor]): Sharded embeddings. Shape: (B, S/cp, H) - - labels (Optional[torch.Tensor]): Possibly sharded labels. Shape: (B, S/cp) - - loss_mask (Optional[torch.Tensor]): Possibly sharded loss mask. Shape: (B, S/cp) - - packed_seq_params (PackedSeqParams, optional): Updated packed sequence parameters. - """ - if not self.cfg.use_cp: - return embeddings, labels, loss_mask, packed_seq_params - - # Distribute sequence across CP ranks - batch = dict() - if embeddings is not None: - batch["embeddings"] = embeddings - if labels is not None: - batch["labels"] = labels - if loss_mask is not None: - batch["loss_mask"] = loss_mask - - if packed_seq_params is None or getattr(packed_seq_params, 'qkv_format', 'sbhd') == 'sbhd': - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=self.cfg.cp_group) - else: - assert _HAVE_TEX and is_te_min_version("1.10.0"), ( - "Please update Transformer Engine to >= 1.10 " - "to use Context Parallel with THD format data" + boundary_batch["decoder_input"] = local_embeddings + return cp_batch + + def _validate_sequence_length( + self, embeddings: Optional[torch.Tensor], is_packed: bool + ) -> None: + """Validate dense sequence divisibility and TP-overlap constraints.""" + if embeddings is None: + return + + shard_factor = None + if self.cfg.use_cp and self.cfg.seq_parallel: + shard_factor = get_pg_size(self.cfg.tp_group) * get_pg_size(self.cfg.cp_group) * 2 + elif self.cfg.use_cp: + shard_factor = get_pg_size(self.cfg.cp_group) * 2 + elif self.cfg.seq_parallel: + shard_factor = get_pg_size(self.cfg.tp_group) + + if shard_factor is not None and not is_packed: + assert embeddings.shape[0] % shard_factor == 0, ( + f"Sequence length should be divisible by {shard_factor} " + "for Sequence/Context parallelism" ) - assert self.cfg.cp_group is not None - cp_size = get_pg_size(self.cfg.cp_group) - cp_rank = get_pg_rank(self.cfg.cp_group) - for key, data in batch.items(): - index = tex.thd_get_partitioned_indices( - packed_seq_params.cu_seqlens_q_padded, data.size(1), cp_size, cp_rank - ) - batch[key] = data.index_select(1, index) - # Extract sharded tensors; embeddings stay in [B, S/cp, H]. shard() transposes - # them to language-model layout [S/cp, B, H] after CP and before optional SP. - embeddings = batch.get("embeddings", None) - labels = batch.get("labels", None) - loss_mask = batch.get("loss_mask", None) - - return embeddings, labels, loss_mask, packed_seq_params + if self.cfg.seq_parallel and self.cfg.tp_comm_overlap: + assert embeddings.shape[0] == self.cfg.max_seq_len, ( + "TP Comm overlap requires Vision+Text token length " + "== language_max_sequence_length" + ) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 6a46cb6202c..176f6290334 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -2436,8 +2436,6 @@ def prepare_cp_layout( requires_conversion = self.cp_group.size() > 1 and source_layout != target_layout if requires_conversion: - if mtp_input_mask is not None: - raise ValueError("mtp_input_mask is not supported with CP layout conversion") if cp_batch is None: raise ValueError("cp_batch is required when MTP uses a different CP layout") hidden_states = convert_cp_layout( @@ -2495,6 +2493,8 @@ def prepare_cp_layout( position_ids = layout_batch["position_ids"] labels = layout_batch["labels"] loss_mask = layout_batch["loss_mask"] + if mtp_input_mask is not None: + mtp_input_mask = layout_batch["mtp_input_mask"] return MultiTokenPredictionInputs( input_ids=input_ids, diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 4d47cb24b37..a55dc9e6ee2 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2609,15 +2609,22 @@ def _get_batch_on_this_cp_rank_per_document_balancing( if batch["cu_seqlens_padded"] is not None else batch["cu_seqlens"] )[0] + SEQUENCE_KEYS = ( + 'tokens', + 'labels', + 'loss_mask', + 'position_ids', + 'decoder_input', + 'mtp_input_mask', + ) + sequence_tensor = next( + (batch.get(key) for key in SEQUENCE_KEYS if batch.get(key) is not None), None + ) + if sequence_tensor is None: + return batch index = tex.thd_get_partitioned_indices( - cu_seqlens_for_te, - ( - batch["tokens"].size(1) if batch["tokens"] is not None else batch["labels"].size(1) - ), # NOTE(asolergi-nv): Labels to enable PP! - cp_size, - cp_rank, + cu_seqlens_for_te, sequence_tensor.size(1), cp_size, cp_rank ) - SEQUENCE_KEYS = ('tokens', 'labels', 'loss_mask', 'position_ids') for key in SEQUENCE_KEYS: if batch.get(key) is not None: batch[key] = batch[key].index_select(1, index) diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py index 11f9e0cf59b..4ece50b79d5 100644 --- a/tests/unit_tests/data/test_get_batch.py +++ b/tests/unit_tests/data/test_get_batch.py @@ -787,13 +787,19 @@ def test_get_batch_on_this_cp_rank_contiguous_keeps_attention_mask_zigzag(cp_siz def test_get_batch_on_this_cp_rank_zigzag_packed(cp_rank, expected_indices, with_contiguous_layout): tokens = torch.arange(1, 17).view(1, 16) loss_mask = torch.tensor([[1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) + decoder_input = torch.stack((tokens, tokens + 20), dim=-1) + mtp_input_mask = tokens.remainder(2).to(dtype=torch.bool) + cu_seqlens = torch.tensor([[0, 3, 13]], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([[0, 4, 16]], dtype=torch.int32) batch = { "tokens": tokens, "labels": tokens + 100, "loss_mask": loss_mask, "position_ids": tokens - 1, - "cu_seqlens": torch.tensor([[0, 3, 13]], dtype=torch.int32), - "cu_seqlens_padded": torch.tensor([[0, 4, 16]], dtype=torch.int32), + "decoder_input": decoder_input, + "mtp_input_mask": mtp_input_mask, + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, "max_seqlen": torch.tensor([12], dtype=torch.int32), } parallel_context = SimpleNamespace( @@ -840,8 +846,16 @@ def test_get_batch_on_this_cp_rank_zigzag_packed(cp_rank, expected_indices, with expected_loss_mask = loss_mask.index_select(1, gathered_indices).masked_fill( padding.view(1, -1), 0 ) + expected_decoder_input = decoder_input.index_select(1, gathered_indices).masked_fill( + padding.view(1, -1, 1), 0 + ) + expected_mtp_input_mask = mtp_input_mask.index_select(1, gathered_indices).masked_fill( + padding.view(1, -1), 0 + ) torch.testing.assert_close(result["tokens"], expected_tokens) torch.testing.assert_close(result["loss_mask"], expected_loss_mask) + torch.testing.assert_close(result["decoder_input"], expected_decoder_input) + torch.testing.assert_close(result["mtp_input_mask"], expected_mtp_input_mask) torch.testing.assert_close( result["cu_seqlens_padded"], torch.tensor([[0, 8, 32]], dtype=torch.int32) ) @@ -855,9 +869,7 @@ def test_get_batch_on_this_cp_rank_zigzag_packed(cp_rank, expected_indices, with torch.testing.assert_close( zigzag_packed_seq_params.cu_seqlens_q_padded, result["cu_seqlens_padded"].squeeze(0) ) - torch.testing.assert_close( - zigzag_packed_seq_params.cu_seqlens_q, batch["cu_seqlens"].squeeze(0) - ) + torch.testing.assert_close(zigzag_packed_seq_params.cu_seqlens_q, cu_seqlens.squeeze(0)) assert zigzag_packed_seq_params.max_seqlen_q == 24 assert zigzag_packed_seq_params.pad_between_seqs assert zigzag_packed_seq_params.total_tokens == 32 @@ -889,6 +901,27 @@ def test_metadata_only_cp_batch_skips_sharding(): assert cp_batch.get_packed_seq_params("zigzag") is not None +def test_size_one_cp_batch_without_group_skips_sharding(): + tokens = torch.arange(8).view(1, 8) + batch = { + "tokens": tokens, + "cu_seqlens": torch.tensor([[0, 3, 8]], dtype=torch.int32), + "cu_seqlens_padded": None, + "max_seqlen": torch.tensor([5], dtype=torch.int32), + } + + with patch("megatron.core.utils.get_batch_on_this_cp_rank") as shard_batch: + cp_batch = get_batches_on_this_cp_rank( + batch, boundary_layout="zigzag", is_hybrid_cp=False, cp_group=None, tokens_per_sample=8 + ) + + shard_batch.assert_not_called() + assert cp_batch.get_batch()["tokens"] is tokens + packed_seq_params = cp_batch.get_packed_seq_params() + assert packed_seq_params is not None + torch.testing.assert_close(packed_seq_params.cu_seqlens_q, batch["cu_seqlens"].squeeze(0)) + + def test_get_batch_builds_required_cp_layouts(): cp_size = 4 seq_length = 16 diff --git a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index f5e676fe632..e88b1f68294 100644 --- a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -203,7 +203,7 @@ def test_backward_step_skips_detached_output_with_explicit_none_grad(monkeypatch def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1): """Create a HyperCommGrid (base view) plus a dense expert view, matching the topology builder. - These tests are dense (ep=1); the expert view relabels the base axes over the same ranks + These tests are dense; with CP=1, the expert view relabels the base axes over the same ranks (expt_tp=tp, ep=cp=1, expt_dp=dp), so the optimizer's expert groups resolve to the dense collapse (tp_ep_pp = tp x pp, expt_dp = dp) without changing the base rank layout. """ @@ -649,7 +649,9 @@ def get_mimo_model( module_to_grid_map=module_to_grid_map, ) - mimo_model = MimoModel(mimo_config, cp_group=language_pg.cp, tp_group=language_pg.tp) + mimo_model = MimoModel( + mimo_config, cp_group=language_pg.cp, tp_group=language_pg.tp, tp_cp_group=language_pg.tp_cp + ) mimo_model.to(torch.device("cuda")) if bf16: mimo_model.to(torch.bfloat16) diff --git a/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py index 6b08208eaaa..eebd23bd537 100644 --- a/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py +++ b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py @@ -12,6 +12,7 @@ from megatron.core.models.mimo.config import MimoModelConfig from megatron.core.models.mimo.model.base import MimoModel from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig class TestEmbeddingAlignment: @@ -23,9 +24,7 @@ def setup_method(self): language_model_spec = ModuleSpec( module=MagicMock, params={ - 'config': MagicMock( - context_parallel_size=1, sequence_parallel=False, pipeline_model_parallel_size=1 - ) + 'config': TransformerConfig(num_layers=1, hidden_size=64, num_attention_heads=4) }, ) self.mimo_config = MimoModelConfig( diff --git a/tests/unit_tests/models/mimo/test_mimo_mock_data.py b/tests/unit_tests/models/mimo/test_mimo_mock_data.py index 2e0ed2946ab..7beff32d8b8 100644 --- a/tests/unit_tests/models/mimo/test_mimo_mock_data.py +++ b/tests/unit_tests/models/mimo/test_mimo_mock_data.py @@ -44,12 +44,16 @@ def _args(): ) -def _topology(*, language_rank, encoder_rank=None): +def _topology(*, language_rank, encoder_rank=None, cp_rank=0, cp_size=1, gtp_rank=0, gtp_size=1): encoder = RADIO_ENCODER_MODULE_NAME grids = {"language": _grid(language_rank)} pgs = { "language": SimpleNamespace( - pp=_group(size=3), dp=_group(rank=0, size=2), dp_cp_gtp_remat=None + pp=_group(size=3), + cp=_group(rank=cp_rank, size=cp_size), + dp=_group(rank=0, size=2), + gtp_remat=_group(rank=gtp_rank, size=gtp_size), + dp_cp_gtp_remat=None, ) } if encoder_rank is not None: @@ -63,7 +67,6 @@ def adapter(monkeypatch, run_config): from examples.mimo.training import data run_config.rng.seed = 123 - monkeypatch.setattr(data, "get_pg_rank", lambda pg: pg.rank()) monkeypatch.setattr(data, "is_pp_first_stage", lambda pg: pg.rank() == 0) monkeypatch.setattr(data, "is_pp_last_stage", lambda pg: pg.rank() == pg.size() - 1) return data @@ -143,6 +146,8 @@ def test_cp_replicas_share_batches_without_merging_data_lanes(adapter, cp_size, topology = _topology(language_rank=True) pg = topology.module_pgs["language"] pg.cp = _group(rank=cp_rank, size=cp_size) + pg.dp = _group(rank=lane // gtp_size, size=args.mimo_llm_dp) + pg.gtp_remat = _group(rank=lane % gtp_size, size=gtp_size) pg.dp_cp_gtp_remat = _group( rank=lane * cp_size + cp_rank, size=args.mimo_llm_dp * gtp_size * cp_size ) diff --git a/tests/unit_tests/models/mimo/test_mimo_model.py b/tests/unit_tests/models/mimo/test_mimo_model.py index ad7c802ac07..24e4200a65c 100644 --- a/tests/unit_tests/models/mimo/test_mimo_model.py +++ b/tests/unit_tests/models/mimo/test_mimo_model.py @@ -16,6 +16,7 @@ from transformers import WhisperConfig, WhisperModel from megatron.core import parallel_state +from megatron.core.context_parallel import ContextParallelBatch from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, @@ -331,7 +332,7 @@ def capture_lm_forward(*args, **kwargs): expected_mtp_input_mask &= input_ids != special_token_id torch.testing.assert_close(captured['mtp_input_mask'], expected_mtp_input_mask) - def test_prepare_mtp_inputs_uses_optional_text_indices(self): + def test_prepare_language_inputs_uses_optional_text_indices(self): """Precomputed text positions and the fallback scan must produce the same MTP mask.""" mimo_model = self._make_vlm() input_ids = self._make_input_ids() @@ -344,21 +345,38 @@ def test_prepare_mtp_inputs_uses_optional_text_indices(self): expected_mask &= input_ids != special_token_id text_token_indices = expected_mask.reshape(-1).nonzero(as_tuple=False).flatten() - indexed_ids, indexed_positions, indexed_mask = mimo_model._prepare_mtp_inputs( + indexed = mimo_model._prepare_language_inputs( + embeddings=None, input_ids=input_ids, position_ids=position_ids, - packed_seq_params=None, + labels=None, + loss_mask=None, + cu_seqlens=None, + cu_seqlens_padded=None, + max_seqlen=None, owns_mtp=True, text_token_indices=text_token_indices, ) - fallback_ids, fallback_positions, fallback_mask = mimo_model._prepare_mtp_inputs( - input_ids=input_ids, position_ids=position_ids, packed_seq_params=None, owns_mtp=True + fallback = mimo_model._prepare_language_inputs( + embeddings=None, + input_ids=input_ids, + position_ids=position_ids, + labels=None, + loss_mask=None, + cu_seqlens=None, + cu_seqlens_padded=None, + max_seqlen=None, + owns_mtp=True, ) - assert indexed_ids is fallback_ids is input_ids - assert indexed_positions is fallback_positions is position_ids - torch.testing.assert_close(indexed_mask, expected_mask) - torch.testing.assert_close(indexed_mask, fallback_mask) + indexed_batch = indexed.get_batch() + fallback_batch = fallback.get_batch() + assert indexed_batch["tokens"] is fallback_batch["tokens"] is input_ids + assert indexed_batch["position_ids"] is fallback_batch["position_ids"] is position_ids + torch.testing.assert_close(indexed_batch["mtp_input_mask"], expected_mask) + torch.testing.assert_close( + indexed_batch["mtp_input_mask"], fallback_batch["mtp_input_mask"] + ) def test_forward_with_image_modality(self): """Test forward pass with text and image input.""" @@ -466,13 +484,13 @@ def test_pipeline_model_parallel_accepted(self): model = MimoModel(mimo_config) assert model is not None - def test_partition_adapter_none_by_default(self): - """Test that partition_adapter is None with default config (no CP/SP).""" + def test_partition_adapter_created_by_default(self): + """Language inputs use the same batch preparation path when CP and SP are disabled.""" mimo_model = self._make_vlm() - assert mimo_model.partition_adapter is None + assert mimo_model.partition_adapter is not None - def test_forward_with_packing_kwargs(self): - """Test that packing_kwargs builds PackedSeqParams with qkv_format='thd' and int32 seqlens.""" + def test_forward_with_packed_sequence_metadata(self): + """Canonical MIMO metadata builds the language model's THD parameters.""" from megatron.core.packed_seq_params import PackedSeqParams mimo_model = self._make_vlm() @@ -480,10 +498,9 @@ def test_forward_with_packing_kwargs(self): position_ids = self._make_position_ids() cu_seqlens = torch.tensor( - [0, self.seq_len, 2 * self.seq_len], dtype=torch.int64, device=self.device + [[0, self.seq_len, 2 * self.seq_len]], dtype=torch.int32, device=self.device ) - packing_kwargs = {"cu_seqlens_q": cu_seqlens.clone(), "cu_seqlens_kv": cu_seqlens.clone()} - + max_seqlen = torch.tensor([self.seq_len], dtype=torch.int32, device=self.device) text_emb = torch.zeros(self.batch_size * self.seq_len, self.hidden_size, device=self.device) combined_emb = torch.zeros( self.seq_len, self.batch_size, self.hidden_size, device=self.device @@ -506,7 +523,8 @@ def capture_lm_forward(*args, **kwargs): input_ids=input_ids, position_ids=position_ids, modality_inputs=None, - packing_kwargs=packing_kwargs, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) packed_seq_params = captured['packed_seq_params'] @@ -514,12 +532,14 @@ def capture_lm_forward(*args, **kwargs): assert packed_seq_params.qkv_format == 'thd' assert packed_seq_params.cu_seqlens_q.dtype == torch.int32 assert packed_seq_params.cu_seqlens_kv.dtype == torch.int32 + torch.testing.assert_close(packed_seq_params.cu_seqlens_q, cu_seqlens.squeeze(0)) + assert packed_seq_params.max_seqlen_q == self.seq_len def test_forward_with_partition_adapter(self): """MTP token metadata must use the same CP-local sequence as hidden states. - The caller no longer transposes around shard(): it passes the sequence-first - ``(S, B, H)`` combined embeddings straight in, and shard()'s LM-layout output + The caller passes the sequence-first ``(S, B, H)`` combined embeddings into one + partition operation, and the adapter's LM-layout output flows straight into the language model as ``decoder_input``. Token and position IDs used by MTP must be partitioned to that same local sequence. """ @@ -536,21 +556,28 @@ def test_forward_with_partition_adapter(self): modality_token_indices = {"text": text_token_indices} sharded_seq_len = self.seq_len // 2 - # shard() returns LM-layout [S/cp, B, H] and a (CP-sharded) loss mask. + # The adapter returns LM-layout [S/cp, B, H] and a CP-sharded loss mask. sharded_emb = torch.zeros( sharded_seq_len, self.batch_size, self.hidden_size, device=self.device ) sharded_loss_mask = torch.ones(self.batch_size, sharded_seq_len, device=self.device) sharded_input_ids = input_ids[:, :sharded_seq_len].clone() sharded_position_ids = position_ids[:, :sharded_seq_len].clone() - sharded_token_metadata = torch.cat( - (sharded_input_ids, text_mask[:, :sharded_seq_len].to(dtype=input_ids.dtype)), dim=0 + boundary_batch = { + "tokens": sharded_input_ids, + "position_ids": sharded_position_ids, + "decoder_input": sharded_emb, + "labels": None, + "loss_mask": sharded_loss_mask, + "mtp_input_mask": text_mask[:, :sharded_seq_len], + } + cp_batch = ContextParallelBatch( + boundary_layout="contiguous", + batches_by_layout={"contiguous": boundary_batch, "zigzag": {}}, + packed_seq_params_by_layout={"contiguous": None, "zigzag": None}, ) mock_adapter = MagicMock() - mock_adapter.shard.side_effect = [ - (sharded_emb, None, sharded_loss_mask, None), - (None, sharded_position_ids, sharded_token_metadata, None), - ] + mock_adapter.partition.return_value = cp_batch mimo_model.partition_adapter = mock_adapter text_emb = torch.zeros(self.batch_size * self.seq_len, self.hidden_size, device=self.device) @@ -566,6 +593,7 @@ def capture_lm_forward(*args, **kwargs): captured['decoder_input'] = kwargs.get('decoder_input') captured['loss_mask'] = kwargs.get('loss_mask') captured['mtp_input_mask'] = kwargs.get('mtp_input_mask') + captured['cp_batch'] = kwargs.get('cp_batch') return torch.zeros( self.batch_size, sharded_seq_len, self.vocab_size, device=self.device ) @@ -585,17 +613,21 @@ def capture_lm_forward(*args, **kwargs): modality_token_indices=modality_token_indices, ) - assert mock_adapter.shard.call_count == 2 - shard_kwargs = mock_adapter.shard.call_args_list[0].kwargs - # The helper passes sequence-first [S, B, H] embeddings straight to shard(). - assert shard_kwargs['embeddings'].shape == (self.seq_len, self.batch_size, self.hidden_size) - assert shard_kwargs['loss_mask'] is loss_mask - mtp_shard_kwargs = mock_adapter.shard.call_args_list[1].kwargs - assert mtp_shard_kwargs['embeddings'] is None - assert mtp_shard_kwargs['labels'] is position_ids - expected_token_metadata = torch.cat((input_ids, text_mask.to(dtype=input_ids.dtype)), dim=0) - torch.testing.assert_close(mtp_shard_kwargs['loss_mask'], expected_token_metadata) - # shard()'s LM-layout output flows straight into the LM (no extra transpose). + mock_adapter.partition.assert_called_once() + partition_kwargs = mock_adapter.partition.call_args.kwargs + assert partition_kwargs['embeddings'].shape == ( + self.seq_len, + self.batch_size, + self.hidden_size, + ) + assert partition_kwargs['input_ids'] is input_ids + assert partition_kwargs['position_ids'] is position_ids + assert partition_kwargs['loss_mask'] is loss_mask + assert partition_kwargs['cu_seqlens'] is None + assert partition_kwargs['cu_seqlens_padded'] is None + assert partition_kwargs['max_seqlen'] is None + torch.testing.assert_close(partition_kwargs['mtp_input_mask'], text_mask) + # The adapter's LM-layout output and dual-layout CP batch flow straight into the LM. assert captured['decoder_input'].shape == ( sharded_seq_len, self.batch_size, @@ -606,7 +638,8 @@ def capture_lm_forward(*args, **kwargs): assert captured['position_ids'] is sharded_position_ids assert captured['mtp_input_mask'].dtype == torch.bool assert torch.equal(captured['mtp_input_mask'], text_mask[:, :sharded_seq_len]) - # forward() returns the (possibly sharded) loss mask from shard(). + assert captured['cp_batch'] is cp_batch + # forward() returns the adapter's possibly sharded loss mask. assert out_loss_mask is sharded_loss_mask def test_get_text_embeddings_raises_when_sp_and_embedding_scatter_enabled(self): diff --git a/tests/unit_tests/models/mimo/test_mimo_partition.py b/tests/unit_tests/models/mimo/test_mimo_partition.py index 72071def5d7..c86a8f605e3 100644 --- a/tests/unit_tests/models/mimo/test_mimo_partition.py +++ b/tests/unit_tests/models/mimo/test_mimo_partition.py @@ -12,8 +12,10 @@ import pytest import torch +from megatron.core.context_parallel import ContextParallelBatch from megatron.core.hyper_comm_grid import HyperCommGrid from megatron.core.models.mimo.partition.utils import PartitionAdapter, PartitionConfig +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -60,6 +62,7 @@ def test_from_mp_config_explicit_cp_group(self): def test_from_mp_config_explicit_tp_group(self): mock_tp_group = MagicMock() + mock_tp_cp_group = MagicMock() mp = TransformerConfig( num_layers=1, hidden_size=64, @@ -68,9 +71,12 @@ def test_from_mp_config_explicit_tp_group(self): sequence_parallel=True, ) with patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=1): - cfg = PartitionConfig.from_mp_config(mp, max_seq_len=512, tp_group=mock_tp_group) + cfg = PartitionConfig.from_mp_config( + mp, max_seq_len=512, tp_group=mock_tp_group, tp_cp_group=mock_tp_cp_group + ) assert cfg.seq_parallel is True assert cfg.tp_group is mock_tp_group + assert cfg.tp_cp_group is mock_tp_cp_group def test_from_mp_config_auto_fetch_cp_group(self): mock_group = MagicMock() @@ -108,8 +114,8 @@ def test_from_mp_config_auto_fetch_tp_group(self): @pytest.mark.experimental -class TestPartitionAdapterShard: - """Tests for PartitionAdapter.shard().""" +class TestPartitionAdapter: + """Tests for PartitionAdapter.partition().""" def _make_cfg( self, @@ -137,20 +143,20 @@ def _make_tensors(self, B=2, S=8, H=16): return embeddings, labels, loss_mask def test_noop_when_both_disabled(self): - """With neither CP nor SP active, shard() is a pure passthrough. + """With neither CP nor SP active, partition() is a pure passthrough. - Production never constructs a PartitionAdapter unless CP or SP is enabled. - Embeddings are already sequence-first (S, B, H), so with no collectives the - inputs are returned untouched (no transpose, no sharding). + Embeddings are already sequence-first (S, B, H), so with no collectives the inputs are + returned untouched while packed metadata still follows the shared preparation path. """ cfg = self._make_cfg(use_cp=False, seq_parallel=False) adapter = PartitionAdapter(cfg) embeddings, labels, loss_mask = self._make_tensors(B=2, S=8, H=16) - out = adapter.shard(embeddings, labels, loss_mask) - assert out[0] is embeddings - assert out[1] is labels - assert out[2] is loss_mask - assert out[3] is None + cp_batch = adapter.partition(embeddings, None, None, labels, loss_mask, None) + batch = cp_batch.get_batch() + assert batch["decoder_input"] is embeddings + assert batch["labels"] is labels + assert batch["loss_mask"] is loss_mask + assert cp_batch.get_packed_seq_params() is None def test_seq_not_divisible_raises(self): mock_cp_group = MagicMock() @@ -163,7 +169,7 @@ def test_seq_not_divisible_raises(self): patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), pytest.raises(AssertionError, match="divisible"), ): - adapter.shard(embeddings, labels, loss_mask) + adapter.partition(embeddings, None, None, labels, loss_mask, None) def test_tp_comm_overlap_seq_len_assertion(self): mock_tp_group = MagicMock() @@ -179,33 +185,45 @@ def test_tp_comm_overlap_seq_len_assertion(self): patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), pytest.raises(AssertionError, match="TP Comm overlap"), ): - adapter.shard(embeddings, labels, loss_mask) + adapter.partition(embeddings, None, None, labels, loss_mask, None) def test_thd_format_skips_divisibility_check(self): - """PackedSeqParams with qkv_format='thd' bypasses the divisibility assertion.""" - from megatron.core.packed_seq_params import PackedSeqParams - + """Packed-sequence metadata bypasses the dense divisibility assertion.""" mock_cp_group = MagicMock() cfg = self._make_cfg(use_cp=True, max_seq_len=7, cp_group=mock_cp_group) adapter = PartitionAdapter(cfg) embeddings = torch.rand(7, 2, 16) # seq-first; len=7 not divisible by cp*2, THD skips check labels = torch.randint(0, 100, (2, 7)) loss_mask = torch.ones(2, 7) - packed_seq_params = MagicMock(spec=PackedSeqParams) - packed_seq_params.qkv_format = 'thd' - packed_seq_params.cu_seqlens_q_padded = torch.tensor([0, 4, 7], dtype=torch.int32) - - # THD path calls tex.thd_get_partitioned_indices — mock it to return first 4 indices - fake_index = torch.arange(4, dtype=torch.int32) + cu_seqlens = torch.tensor([[0, 4, 7]], dtype=torch.int32) + cp_batch = ContextParallelBatch.from_single_layout( + "zigzag", + { + "decoder_input": embeddings.transpose(0, 1)[:, :4], + "labels": labels[:, :4], + "loss_mask": loss_mask[:, :4], + }, + MagicMock(spec=PackedSeqParams), + ) with ( patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), - patch('megatron.core.models.mimo.partition.utils.get_pg_rank', return_value=0), - patch('megatron.core.models.mimo.partition.utils.tex') as mock_tex, + patch( + 'megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank', + return_value=cp_batch, + ), ): - mock_tex.thd_get_partitioned_indices.return_value = fake_index # Should NOT raise AssertionError about divisibility - out = adapter.shard(embeddings, labels, loss_mask, packed_seq_params) - assert out[0] is not None + result = adapter.partition( + embeddings, + None, + None, + labels, + loss_mask, + None, + cu_seqlens=cu_seqlens, + max_seqlen=torch.tensor([4], dtype=torch.int32), + ) + assert result.get_batch()["decoder_input"] is not None def test_none_embeddings_skips_shard_factor_check(self): """When embeddings is None, the divisibility check is skipped (non-first PP stage).""" @@ -214,117 +232,116 @@ def test_none_embeddings_skips_shard_factor_check(self): adapter = PartitionAdapter(cfg) labels = torch.randint(0, 100, (2, 7)) loss_mask = torch.ones(2, 7) - cp_sharded = {'labels': labels[:, :4], 'loss_mask': loss_mask[:, :4]} + cp_batch = ContextParallelBatch.from_single_layout( + "zigzag", {'labels': labels[:, :4], 'loss_mask': loss_mask[:, :4]}, None + ) with ( patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), patch( - 'megatron.core.models.mimo.partition.utils.get_batch_on_this_cp_rank', - return_value=cp_sharded, + 'megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank', + return_value=cp_batch, ), ): - out = adapter.shard(None, labels, loss_mask) - assert out[0] is None - assert out[1].shape == (2, 4) - assert out[2].shape == (2, 4) + result = adapter.partition(None, None, None, labels, loss_mask, None) + batch = result.get_batch() + assert batch.get("decoder_input") is None + assert batch["labels"].shape == (2, 4) + assert batch["loss_mask"].shape == (2, 4) @pytest.mark.experimental -class TestPartitionAdapterApplyContextParallel: - """Tests for PartitionAdapter._apply_context_parallel().""" +class TestPartitionAdapterContextParallelBatch: + """Tests for constructing MIMO's dual-layout CP batch.""" - def _make_cfg(self, use_cp=True, cp_group=None): - return PartitionConfig( - use_cp=use_cp, - seq_parallel=False, + def test_partitions_all_language_inputs_together(self): + cp_group = MagicMock() + tp_group = MagicMock() + tp_cp_group = MagicMock() + cfg = PartitionConfig( + use_cp=True, + seq_parallel=True, tp_comm_overlap=False, - max_seq_len=128, + max_seq_len=8, + linear_cp_layout="contiguous", + attention_cp_layout="zigzag", cp_group=cp_group, + tp_group=tp_group, + tp_cp_group=tp_cp_group, ) - - def test_returns_unchanged_when_cp_disabled(self): - cfg = self._make_cfg(use_cp=False) - adapter = PartitionAdapter(cfg) - embeddings = torch.rand(2, 8, 16) - labels = torch.randint(0, 100, (2, 8)) - loss_mask = torch.ones(2, 8) - out = adapter._apply_context_parallel(embeddings, labels, loss_mask, None) - assert out[0] is embeddings - assert out[1] is labels - assert out[2] is loss_mask - assert out[3] is None - - def test_sbhd_path_calls_get_batch_on_this_cp_rank(self): - mock_cp_group = MagicMock() - cfg = self._make_cfg(use_cp=True, cp_group=mock_cp_group) adapter = PartitionAdapter(cfg) - embeddings = torch.rand(2, 8, 16) - labels = torch.randint(0, 100, (2, 8)) - loss_mask = torch.ones(2, 8) - sharded = { - 'embeddings': embeddings[:, :4, :], - 'labels': labels[:, :4], - 'loss_mask': loss_mask[:, :4], + embeddings, labels, loss_mask = TestPartitionAdapter()._make_tensors(S=8) + input_ids = torch.arange(8).view(1, 8).expand(2, -1) + position_ids = torch.stack((input_ids, input_ids + 10, input_ids + 20)) + mtp_input_mask = input_ids != 3 + cu_seqlens = torch.tensor([[0, 8, 16]], dtype=torch.int32) + max_seqlen = torch.tensor([8], dtype=torch.int32) + + boundary_batch = { + "tokens": input_ids[:, :4], + "position_ids": position_ids.movedim(0, -1)[:, :4], + "labels": labels[:, :4], + "loss_mask": loss_mask[:, :4], + "mtp_input_mask": mtp_input_mask[:, :4], + "decoder_input": embeddings.transpose(0, 1)[:, :4], } - with patch( - 'megatron.core.models.mimo.partition.utils.get_batch_on_this_cp_rank', - return_value=sharded, - ) as mock_fn: - out = adapter._apply_context_parallel(embeddings, labels, loss_mask, None) - mock_fn.assert_called_once() - # _apply_context_parallel keeps batch-first [B, S/cp, H]; shard() transposes later. - assert out[0].shape == (2, 4, 16) - assert out[1].shape == (2, 4) - - def test_all_none_inputs_produces_none_outputs(self): - mock_cp_group = MagicMock() - cfg = self._make_cfg(use_cp=True, cp_group=mock_cp_group) - adapter = PartitionAdapter(cfg) - with patch( - 'megatron.core.models.mimo.partition.utils.get_batch_on_this_cp_rank', return_value={} - ): - out = adapter._apply_context_parallel(None, None, None, None) - assert all(v is None for v in out[:3]) - - def test_only_non_none_tensors_added_to_batch(self): - """None tensors must not appear in the batch dict passed to get_batch_on_this_cp_rank.""" - mock_cp_group = MagicMock() - cfg = self._make_cfg(use_cp=True, cp_group=mock_cp_group) - adapter = PartitionAdapter(cfg) - embeddings = torch.rand(2, 8, 16) - sharded = {'embeddings': embeddings[:, :4, :]} - captured = {} - - def mock_fn(batch, **kwargs): - captured.update(batch) - return sharded - - with patch( - 'megatron.core.models.mimo.partition.utils.get_batch_on_this_cp_rank', - side_effect=mock_fn, - ): - out = adapter._apply_context_parallel(embeddings, None, None, None) - - assert 'embeddings' in captured - assert 'labels' not in captured - assert 'loss_mask' not in captured - assert out[0] is not None - assert out[1] is None - - def test_thd_path_raises_when_te_unavailable(self): - """THD format must assert when Transformer Engine is not available.""" - from megatron.core.packed_seq_params import PackedSeqParams + zigzag_batch = { + "tokens": input_ids[:, 4:], + "position_ids": position_ids.movedim(0, -1)[:, 4:], + "labels": labels[:, 4:], + "loss_mask": loss_mask[:, 4:], + "mtp_input_mask": mtp_input_mask[:, 4:], + "decoder_input": embeddings.transpose(0, 1)[:, 4:], + } + local_packed = MagicMock(spec=PackedSeqParams) + cp_batch = ContextParallelBatch( + boundary_layout="contiguous", + batches_by_layout={"contiguous": boundary_batch, "zigzag": zigzag_batch}, + packed_seq_params_by_layout={ + "contiguous": local_packed, + "zigzag": MagicMock(spec=PackedSeqParams), + }, + thd_plan=MagicMock(), + ) + local_embeddings = embeddings[:2] - mock_cp_group = MagicMock() - cfg = self._make_cfg(use_cp=True, cp_group=mock_cp_group) - adapter = PartitionAdapter(cfg) - embeddings = torch.rand(2, 5, 16) - packed_seq_params = MagicMock(spec=PackedSeqParams) - packed_seq_params.qkv_format = 'thd' with ( - patch('megatron.core.models.mimo.partition.utils._HAVE_TEX', False), - pytest.raises(AssertionError, match="Transformer Engine"), + patch( + "megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank", + return_value=cp_batch, + ) as get_cp_batches, + patch( + "megatron.core.models.mimo.partition.utils." + "tensor_parallel.scatter_to_sequence_parallel_region", + return_value=local_embeddings, + ) as scatter, ): - adapter._apply_context_parallel(embeddings, None, None, packed_seq_params) + result = adapter.partition( + embeddings=embeddings, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + loss_mask=loss_mask, + mtp_input_mask=mtp_input_mask, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + (unsharded_batch,) = get_cp_batches.call_args.args + call_kwargs = get_cp_batches.call_args.kwargs + assert call_kwargs["boundary_layout"] == "contiguous" + assert call_kwargs["additional_layouts"] == {"zigzag"} + assert call_kwargs["tp_cp_group"] is tp_cp_group + assert unsharded_batch["decoder_input"].shape == (2, 8, 16) + assert unsharded_batch["position_ids"].shape == (2, 8, 3) + assert unsharded_batch["mtp_input_mask"] is mtp_input_mask + assert unsharded_batch["cu_seqlens"] is cu_seqlens + assert unsharded_batch["max_seqlen"] is max_seqlen + assert result is cp_batch + assert boundary_batch["decoder_input"] is local_embeddings + assert boundary_batch["position_ids"].shape == (3, 2, 4) + assert result.get_packed_seq_params() is local_packed + assert "decoder_input" not in zigzag_batch + scatter.assert_called_once() def _expected_cp_zigzag_shard(tensor: torch.Tensor, cp_size: int, cp_rank: int) -> torch.Tensor: @@ -346,8 +363,8 @@ def _expected_cp_zigzag_shard(tensor: torch.Tensor, cp_size: int, cp_rank: int) int(os.environ.get('WORLD_SIZE', '1')) != 8, reason="Real MIMO CP/SP sharding tests require an 8-GPU world", ) -class TestPartitionAdapterShardRealDistributed: - """Real 8-GPU tests for ``PartitionAdapter.shard()``. +class TestPartitionAdapterRealDistributed: + """Real 8-GPU tests for ``PartitionAdapter.partition()``. These exercise the genuine collectives (CP zigzag chunking via ``get_batch_on_this_cp_rank`` and the SP scatter via @@ -415,9 +432,13 @@ def test_sp_only_scatters_sequence_real(self): adapter = PartitionAdapter(cfg) embeddings, labels, loss_mask = self._make_inputs(B, S, H) - out_emb, out_labels, out_loss_mask, _ = adapter.shard( - embeddings.clone(), labels.clone(), loss_mask.clone() + cp_batch = adapter.partition( + embeddings.clone(), None, None, labels.clone(), loss_mask.clone(), None ) + batch = cp_batch.get_batch() + out_emb = batch["decoder_input"] + out_labels = batch["labels"] + out_loss_mask = batch["loss_mask"] tp_rank = tp_group.rank() shard = S // tp_size @@ -449,9 +470,13 @@ def test_cp_only_shards_sequence_real(self): adapter = PartitionAdapter(cfg) embeddings, labels, loss_mask = self._make_inputs(B, S, H) - out_emb, out_labels, out_loss_mask, _ = adapter.shard( - embeddings.clone(), labels.clone(), loss_mask.clone() + cp_batch = adapter.partition( + embeddings.clone(), None, None, labels.clone(), loss_mask.clone(), None ) + batch = cp_batch.get_batch() + out_emb = batch["decoder_input"] + out_labels = batch["labels"] + out_loss_mask = batch["loss_mask"] cp_rank = cp_group.rank() shard = S // cp_size @@ -487,9 +512,13 @@ def test_cp_and_sp_combined_real(self): adapter = PartitionAdapter(cfg) embeddings, labels, loss_mask = self._make_inputs(B, S, H) - out_emb, out_labels, out_loss_mask, _ = adapter.shard( - embeddings.clone(), labels.clone(), loss_mask.clone() + cp_batch = adapter.partition( + embeddings.clone(), None, None, labels.clone(), loss_mask.clone(), None ) + batch = cp_batch.get_batch() + out_emb = batch["decoder_input"] + out_labels = batch["labels"] + out_loss_mask = batch["loss_mask"] cp_rank = cp_group.rank() tp_rank = tp_group.rank() diff --git a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py index 10e51791e7c..2e213547a27 100644 --- a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py +++ b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py @@ -442,6 +442,27 @@ def test_destination_cp_topology(self): else: assert bridge.dest_cp_reduce_pg is None + def test_cp_forward_broadcast(self): + """Every destination TP/CP rank receives the complete fan-in activation.""" + src_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=4) + dest_grid = create_hypercomm_grid(offset=4, tp=2, cp=2, pp=1, dp=1) + bridge = BridgeCommunicator(src_grid, dest_grid, comm_dtype=torch.float32, tensor_ndim=2) + + activation_matches = True + if bridge.is_current_rank_in_grid(src_grid): + activation = torch.full((2, 3), float(dist.get_rank()), device="cuda") + bridge.send_forward(activation) + else: + received_activation = bridge.recv_forward() + expected_activation = torch.cat( + [torch.full((2, 3), float(rank), device="cuda") for rank in range(4)] + ) + activation_matches = torch.equal(received_activation, expected_activation) + + global_check = torch.tensor(activation_matches, device="cuda", dtype=torch.int32) + dist.all_reduce(global_check, op=dist.ReduceOp.MIN) + assert global_check.item() == 1 + def test_send_forward_recv_forward(self): """Test send_forward and recv_forward operations.""" diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 1f08ed373a6..e8ada757323 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -4042,22 +4042,26 @@ def mtp(**kwargs): position_ids = torch.tensor([[0, 1]]) labels = torch.tensor([[2, 3]]) loss_mask = torch.ones(1, 2) + mtp_input_mask = torch.tensor([[True, False]]) zigzag_input_ids = torch.tensor([[4, 3]]) zigzag_position_ids = torch.tensor([[3, 2]]) zigzag_labels = torch.tensor([[5, 4]]) zigzag_loss_mask = torch.tensor([[1.0, 0.0]]) + zigzag_mtp_input_mask = torch.tensor([[False, True]]) batches_by_layout = { "contiguous": { "tokens": input_ids, "position_ids": position_ids, "labels": labels, "loss_mask": loss_mask, + "mtp_input_mask": mtp_input_mask, }, "zigzag": { "tokens": zigzag_input_ids, "position_ids": zigzag_position_ids, "labels": zigzag_labels, "loss_mask": zigzag_loss_mask, + "mtp_input_mask": zigzag_mtp_input_mask, }, } cp_batch = ContextParallelBatch( @@ -4084,6 +4088,7 @@ def process_mtp_loss_spy(**kwargs): decoder_input=hidden_states, labels=labels, loss_mask=loss_mask, + mtp_input_mask=mtp_input_mask, packed_seq_params=contiguous_packed_seq_params, cp_batch=cp_batch, padding_mask=padding_mask, @@ -4093,6 +4098,7 @@ def process_mtp_loss_spy(**kwargs): assert captured["mtp"]["decoder_input"] is zigzag_hidden_states assert captured["mtp"]["input_ids"] is zigzag_input_ids assert captured["mtp"]["position_ids"] is zigzag_position_ids + assert captured["mtp"]["mtp_input_mask"] is zigzag_mtp_input_mask assert captured["mtp"]["packed_seq_params"] is zigzag_packed_seq_params assert ( captured["mtp"]["packed_seq_params_by_layout"] is expected_packed_seq_params_by_layout From 807f3bda8d8979ce017316010479f3196ff68c72 Mon Sep 17 00:00:00 2001 From: Haoran Zhang Date: Wed, 30 Sep 2026 14:53:26 -0700 Subject: [PATCH 2/2] Wire CP/GTP into MIMO Signed-off-by: Haoran Zhang --- .../mimo/model_providers/nemotron_moe_vlm.py | 7 +- .../mimo/model_providers/radio_encoder.py | 1 + examples/mimo/train.py | 4 +- examples/mimo/training/args.py | 23 ++- examples/mimo/training/data.py | 6 +- examples/mimo/training/topology.py | 111 +++++++---- examples/mimo/utils/hetero.py | 27 ++- megatron/core/context_parallel/utils.py | 40 ++-- megatron/core/models/mimo/model/base.py | 6 +- .../pipeline_parallel/bridge_communicator.py | 27 +-- megatron/core/utils.py | 12 +- .../mimo/test_mimo_embedding_alignment.py | 2 +- .../models/mimo/test_mimo_gtp_cp_layout.py | 185 ++++++++++++++++++ .../models/mimo/test_mimo_hetero_grid_args.py | 47 +++++ .../models/mimo/test_mimo_mock_data.py | 42 +++- .../unit_tests/models/mimo/test_mimo_model.py | 20 +- .../models/mimo/test_mimo_partition.py | 6 +- .../mimo/test_nemotron_moe_vlm_provider.py | 49 ++++- .../test_bridge_communicator.py | 21 +- 19 files changed, 523 insertions(+), 113 deletions(-) create mode 100644 tests/unit_tests/models/mimo/test_mimo_gtp_cp_layout.py diff --git a/examples/mimo/model_providers/nemotron_moe_vlm.py b/examples/mimo/model_providers/nemotron_moe_vlm.py index 9fd11a079e2..2d4abd57b38 100644 --- a/examples/mimo/model_providers/nemotron_moe_vlm.py +++ b/examples/mimo/model_providers/nemotron_moe_vlm.py @@ -134,6 +134,7 @@ def nemotron_language_config( ) -> TransformerConfig: """Nemotron6-MoE language config: stock from-args base + model-specific overrides.""" config = deepcopy(_base_config(args)) + gtp_remat_num_sequence_shards = config.gtp_remat_num_sequence_shards # Code-only fields + hetero parallelism pins. config.variable_seq_lengths = True config.expert_model_parallel_size = ep_size @@ -158,6 +159,7 @@ def nemotron_language_config( ) ) config.sequence_parallel = tp_size > 1 + config.tensor_parallel_num_sequence_shards = tp_size * gtp_remat_num_sequence_shards config.position_embedding_type = "none" return config @@ -202,6 +204,7 @@ def nemotron_projection_config( if base_config is None: _disable_gtp(config) config.sequence_parallel = False + config.tensor_parallel_num_sequence_shards = None return config @@ -243,7 +246,9 @@ def language_model_spec( pp_rank = 0 pp_size = get_grid_dim_size(llm_grid, "pp") tp_size = get_grid_dim_size(llm_grid, "tp") - cp_size = get_grid_dim_size(llm_grid, "cp") + cp_size = get_grid_dim_size(llm_grid, "cp") * get_grid_dim_size( + llm_grid, "gtp_remat_sequence_shards" + ) ep_size = getattr(args, "mimo_llm_ep", 1) expt_tp_size = getattr(args, "mimo_llm_expt_tp", None) or 1 else: diff --git a/examples/mimo/model_providers/radio_encoder.py b/examples/mimo/model_providers/radio_encoder.py index 535afe64c9e..e486ed234c3 100644 --- a/examples/mimo/model_providers/radio_encoder.py +++ b/examples/mimo/model_providers/radio_encoder.py @@ -109,6 +109,7 @@ def _disable_gtp(config: TransformerConfig) -> None: """Keep this module replicated across any LLM GTP axes.""" config.tensor_parallel_num_weight_shards = config.tensor_model_parallel_size config.gtp_weight_remat_size = 1 + config.tensor_parallel_num_sequence_shards = None expert_tp = config.expert_tensor_parallel_size or config.tensor_model_parallel_size config.expert_tensor_parallel_num_weight_shards = expert_tp config.expert_gtp_weight_remat_size = 1 diff --git a/examples/mimo/train.py b/examples/mimo/train.py index f8c30370e16..09f33a77dfd 100644 --- a/examples/mimo/train.py +++ b/examples/mimo/train.py @@ -67,9 +67,7 @@ def _canonicalize_packed_sequence_metadata(batch: Dict[str, Any]) -> Dict[str, A if cu_seqlens_padded is not None: cu_seqlens_padded = cu_seqlens_padded.to(dtype=torch.int32) batch["cu_seqlens_padded"] = ( - cu_seqlens_padded.unsqueeze(0) - if cu_seqlens_padded.dim() == 1 - else cu_seqlens_padded + cu_seqlens_padded.unsqueeze(0) if cu_seqlens_padded.dim() == 1 else cu_seqlens_padded ) batch["max_seqlen"] = packing_kwargs["max_seqlen_q"].to(dtype=torch.int32).reshape(1) diff --git a/examples/mimo/training/args.py b/examples/mimo/training/args.py index 4caa6f990dc..d750c5decd8 100644 --- a/examples/mimo/training/args.py +++ b/examples/mimo/training/args.py @@ -8,8 +8,12 @@ from typing import List from examples.mimo.training.topology import ModuleGridSpec +from examples.mimo.utils.hetero import get_language_batch_parallel_size from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY +from megatron.core.tensor_parallel.gtp_parallel_layout import ( + resolve_tensor_parallel_sequence_shards, +) def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -98,9 +102,9 @@ def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentPa def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tuple[int, int]: """Validate the disjoint hetero grid layout; returns ``(encoder_size, llm_size)``.""" - gtp_weight_remat_size, _ = resolve_hetero_gtp_degrees(args) if args.mimo_llm_cp < 1: raise ValueError("--mimo-llm-cp must be positive") + llm_data_parallel_size = get_language_batch_parallel_size(args) if getattr(args, "mimo_encoder_ddp_overlap", False) and not getattr( args, "overlap_grad_reduce", False @@ -120,12 +124,13 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl f"--mimo-llm-ep ({args.mimo_llm_ep})" ) + gtp_weight_remat_size, _ = resolve_hetero_gtp_degrees(args) llm_size = ( args.mimo_llm_tp - * gtp_weight_remat_size * args.mimo_llm_cp * args.mimo_llm_pp * args.mimo_llm_dp + * gtp_weight_remat_size ) if args.mimo_llm_only: @@ -147,13 +152,12 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl # Fan-out divisibility: the bridge splits every LLM data lane across # mimo_encoder_dp encoder lanes; the split must be exact. - llm_data_parallel_size = args.mimo_llm_dp * gtp_weight_remat_size if (args.micro_batch_size * llm_data_parallel_size) % args.mimo_encoder_dp != 0: raise ValueError( - "--micro-batch-size * --mimo-llm-dp * GTP must be divisible by " + "--micro-batch-size * language sample lanes must be divisible by " "--mimo-encoder-dp " - f"(got {args.micro_batch_size} * {args.mimo_llm_dp} * " - f"{gtp_weight_remat_size} % {args.mimo_encoder_dp} != 0)" + f"(got {args.micro_batch_size} * {llm_data_parallel_size} " + f"% {args.mimo_encoder_dp} != 0)" ) encoder_size = args.mimo_encoder_tp * args.mimo_encoder_dp @@ -182,6 +186,12 @@ def build_module_grid_specs( """Map grid args to the ModuleGridSpec list create_topology consumes.""" encoder_size, llm_size = validate_hetero_grid_args(args, world_size) gtp_weight_remat_size, expert_gtp_weight_remat_size = resolve_hetero_gtp_degrees(args) + _, gtp_remat_num_sequence_shards = resolve_tensor_parallel_sequence_shards( + args.mimo_llm_tp, + getattr(args, "tensor_parallel_num_sequence_shards", None), + gtp_weight_remat_size, + getattr(args, "sequence_parallel", False), + ) language_grid_spec = ModuleGridSpec( name=MIMO_LANGUAGE_MODULE_KEY, @@ -191,6 +201,7 @@ def build_module_grid_specs( pp=args.mimo_llm_pp, ep=args.mimo_llm_ep, gtp_remat=gtp_weight_remat_size, + gtp_remat_num_sequence_shards=gtp_remat_num_sequence_shards, rank_offset=args.mimo_llm_offset, expt_tp=args.mimo_llm_expt_tp or 1, expt_gtp_remat=expert_gtp_weight_remat_size, diff --git a/examples/mimo/training/data.py b/examples/mimo/training/data.py index 6efd3ceec4a..0cbcac3f01b 100644 --- a/examples/mimo/training/data.py +++ b/examples/mimo/training/data.py @@ -12,7 +12,7 @@ from torch.utils.data import DataLoader, Dataset from examples.mimo.training.topology import HeteroTopology -from examples.mimo.utils.hetero import get_data_lane_rank +from examples.mimo.utils.hetero import get_data_lane_rank, get_language_batch_parallel_size from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage @@ -292,13 +292,13 @@ def build_train_valid_test_data_loaders( raise ValueError(f"unsupported dataset provider: {args.dataset_provider}") encoder_name = _encoder_name(topology) - llm_data_parallel_size = args.mimo_llm_dp * args.gtp_weight_remat_size + llm_data_parallel_size = get_language_batch_parallel_size(args) if ( encoder_name is not None and (args.micro_batch_size * llm_data_parallel_size) % args.mimo_encoder_dp ): raise ValueError( - "micro_batch_size * mimo_llm_dp * GTP must be divisible by mimo_encoder_dp" + "micro_batch_size * language sample lanes must be divisible by mimo_encoder_dp" ) language_grid = topology.grids[MIMO_LANGUAGE_MODULE_KEY] diff --git a/examples/mimo/training/topology.py b/examples/mimo/training/topology.py index a6d8798f237..77272db535a 100644 --- a/examples/mimo/training/topology.py +++ b/examples/mimo/training/topology.py @@ -16,8 +16,10 @@ MultiModuleProcessGroupCollection, ProcessGroupCollection, ) +from megatron.core.tensor_parallel.gtp_parallel_layout import GTPParallelLayout _EXPERT_VIEW = "expert" +_WEIGHT_VIEW = "weights" @dataclass @@ -38,16 +40,14 @@ class ModuleGridSpec: # Experts default to TP=1 (set explicitly for MoE); intentionally not Megatron's etp=tp default. expt_tp: int = 1 expt_gtp_remat: int = 1 + gtp_remat_num_sequence_shards: int = 1 + """Sequence partitions within GTP_remat, excluding TP/SP; ``cp`` adds partitions.""" + dp: int = field(init=False) expt_dp: int = field(init=False) def __post_init__(self) -> None: - dense = self.tp * self.gtp_remat * self.cp * self.pp - if self.num_ranks % dense != 0: - raise ValueError( - f"num_ranks ({self.num_ranks}) must be divisible by tp*gtp_remat*cp*pp ({dense})" - ) - self.dp = self.num_ranks // dense + self.dp = self.gtp_layout.dp expert = self.expt_tp * self.ep * self.expt_gtp_remat * self.pp if self.num_ranks % expert != 0: raise ValueError( @@ -61,6 +61,19 @@ def size(self) -> int: """Total ranks spanned by this module's grid.""" return self.num_ranks + @property + def gtp_layout(self) -> GTPParallelLayout: + """Weight and token ownership with GTP sequence sharding.""" + return GTPParallelLayout( + self.num_ranks, + self.tp, + self.pp, + self.cp, + self.gtp_remat, + num_sequence_shards=self.gtp_remat_num_sequence_shards, + rank_offset=self.rank_offset, + ) + @dataclass class HeteroTopology: @@ -128,12 +141,37 @@ def _build_grid( spec: ModuleGridSpec, high_priority_stream_groups: Optional[list[str]] = None ) -> HyperCommGrid: """Create a dense grid plus its expert view and the process groups MIMO needs.""" + layout = spec.gtp_layout if spec.gtp_remat_num_sequence_shards > 1 else None grid = HyperCommGrid( - shape=[spec.tp, spec.cp, spec.gtp_remat, spec.dp, spec.pp], - dim_names=["tp", "cp", "gtp_remat", "dp", "pp"], + shape=( + [ + spec.tp, + spec.gtp_remat_num_sequence_shards, + spec.gtp_remat // spec.gtp_remat_num_sequence_shards, + spec.cp, + spec.dp, + spec.pp, + ] + if layout is not None + else [spec.tp, spec.cp, spec.gtp_remat, spec.dp, spec.pp] + ), + dim_names=( + ["tp", "gtp_remat_sequence_shards", "gtp_remat", "cp", "dp", "pp"] + if layout is not None + else ["tp", "cp", "gtp_remat", "dp", "pp"] + ), rank_offset=spec.rank_offset, backend="nccl", ) + if layout is not None: + # CP and GTP sequence shards share a sample; the remaining GTP extent + # assigns independent samples. The weight view keeps all GTP weight shards together. + grid.register_view( + _WEIGHT_VIEW, + shape=[spec.tp, spec.gtp_remat, layout.num_weight_replicas, spec.pp], + dim_names=["tp", "gtp_remat", "dp", "pp"], + shared_dims=["tp", "pp"], + ) # Expert factorization over the same rank span; pp is shared with the base view. grid.register_view( _EXPERT_VIEW, @@ -148,19 +186,13 @@ def create_pg(dims: list[str], pg_name: str, *, view: Optional[str] = None) -> N ) try: - create_pg(["tp"], "tp") - create_pg(["gtp_remat"], "gtp_remat") - create_pg(["cp"], "cp") - create_pg(["pp"], "pp") - create_pg(["dp"], "dp") - create_pg(["dp", "cp"], "dp_cp") - create_pg(["gtp_remat", "dp"], "gtp_remat_dp") - create_pg(["cp", "gtp_remat", "dp"], "gtp_remat_dp_cp") - create_pg(["tp", "cp"], "tp_cp") - create_pg(["tp", "gtp_remat", "pp"], "mp") - create_pg(["tp", "gtp_remat", "dp"], "tp_dp") - create_pg(["tp", "cp", "gtp_remat", "dp"], "tp_dp_cp") - create_pg(["tp", "cp", "gtp_remat", "dp", "pp"], "intra_dist_opt_instance") + pg_names = { + "dp_gtp_remat": "gtp_remat_dp", + "dp_cp_gtp_remat": "gtp_remat_dp_cp", + "intra_dist_opt": "intra_dist_opt_instance", + } + for field_name, (dims, view) in _dense_group_dims(layout is not None).items(): + create_pg(dims, pg_names.get(field_name, field_name), view=view) create_pg(["ep"], "ep", view=_EXPERT_VIEW) create_pg(["expt_tp"], "ep_tp", view=_EXPERT_VIEW) @@ -178,6 +210,27 @@ def create_pg(dims: list[str], pg_name: str, *, view: Optional[str] = None) -> N return grid +def _dense_group_dims(sequence_sharding: bool) -> dict[str, tuple[list[str], str | None]]: + """Map PGC fields to token or weight dimensions in the module grid.""" + cp = ["cp", "gtp_remat_sequence_shards"] if sequence_sharding else ["cp"] + weights = _WEIGHT_VIEW if sequence_sharding else None + return { + "tp": (["tp"], None), + "gtp_remat": (["gtp_remat"], weights), + "cp": (cp, None), + "pp": (["pp"], None), + "dp": (["dp"], None), + "dp_cp": (["dp"] if sequence_sharding else ["dp", "cp"], weights), + "dp_gtp_remat": (["gtp_remat", "dp"], None), + "dp_cp_gtp_remat": (cp + ["gtp_remat", "dp"], None), + "tp_cp": (["tp"] + cp, None), + "mp": (["tp", "gtp_remat", "pp"], weights), + "tp_dp": (["tp", "gtp_remat", "dp"], None), + "tp_dp_cp": (["tp"] + cp + ["gtp_remat", "dp"], None), + "intra_dist_opt": (["tp"] + cp + ["gtp_remat", "dp", "pp"], None), + } + + def _get_pg_options( group_name: str, high_priority_stream_groups: Optional[list[str]] ) -> Optional[dist.ProcessGroupNCCL.Options]: @@ -231,20 +284,10 @@ def pg_collection_from_grid( Only the language module gets embedding groups; others leave ``embd``/``pos_embd`` as ``None``. """ pgc = ProcessGroupCollection() - pgc.tp = grid.get_pg("tp") - pgc.cp = grid.get_pg("cp") - pgc.pp = grid.get_pg("pp") - pgc.dp = grid.get_pg("dp") - pgc.dp_gtp_remat = grid.get_pg(["gtp_remat", "dp"]) - pgc.dp_cp = grid.get_pg(["dp", "cp"]) - pgc.dp_cp_gtp_remat = grid.get_pg(["cp", "gtp_remat", "dp"]) + group_dims = _dense_group_dims("gtp_remat_sequence_shards" in grid.dim_names) + for field_name, (dims, view) in group_dims.items(): + setattr(pgc, field_name, grid.get_pg(dims, view=view)) pgc.intra_dp_cp = pgc.dp_cp - pgc.gtp_remat = grid.get_pg("gtp_remat") - pgc.tp_cp = grid.get_pg(["tp", "cp"]) - pgc.tp_dp = grid.get_pg(["tp", "gtp_remat", "dp"]) - pgc.tp_dp_cp = grid.get_pg(["tp", "cp", "gtp_remat", "dp"]) - pgc.mp = grid.get_pg(["tp", "gtp_remat", "pp"]) - pgc.intra_dist_opt = grid.get_pg(["tp", "cp", "gtp_remat", "dp", "pp"]) pgc.ep = grid.get_pg("ep", view=_EXPERT_VIEW) pgc.expt_tp = grid.get_pg("expt_tp", view=_EXPERT_VIEW) pgc.expt_gtp_remat = grid.get_pg("expt_gtp_remat", view=_EXPERT_VIEW) diff --git a/examples/mimo/utils/hetero.py b/examples/mimo/utils/hetero.py index 9e33946b48a..ddf3a61c6d6 100644 --- a/examples/mimo/utils/hetero.py +++ b/examples/mimo/utils/hetero.py @@ -5,7 +5,11 @@ from __future__ import annotations from megatron.core.hyper_comm_grid import HyperCommGrid +from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.gtp_parallel_layout import ( + resolve_tensor_parallel_sequence_shards, +) def get_grid_dim_size(grid: HyperCommGrid, dim: str) -> int: @@ -17,8 +21,21 @@ def get_grid_dim_size(grid: HyperCommGrid, dim: str) -> int: def get_data_lane_rank(pg_collection: ProcessGroupCollection) -> int: - """Return the DP x GTP data-lane rank, excluding model-parallel dimensions.""" - gtp_group = getattr(pg_collection, "gtp_remat", None) - gtp_size = gtp_group.size() if gtp_group is not None else 1 - gtp_rank = gtp_group.rank() if gtp_group is not None else 0 - return pg_collection.dp.rank() * gtp_size + gtp_rank + """Return the independent sample rank, excluding every context partition.""" + return pg_collection.dp_gtp_remat.rank() + + +def get_language_batch_parallel_size(args) -> int: + """Count language batch partitions before or after stock argument validation.""" + _, weight_size = resolve_tensor_parallel_weight_shards( + getattr(args, "mimo_llm_tp", 1), + getattr(args, "tensor_parallel_num_weight_shards", None), + getattr(args, "gtp_weight_remat_size", 1), + ) + _, num_sequence_shards = resolve_tensor_parallel_sequence_shards( + getattr(args, "mimo_llm_tp", 1), + getattr(args, "tensor_parallel_num_sequence_shards", None), + weight_size, + getattr(args, "sequence_parallel", False), + ) + return args.mimo_llm_dp * (weight_size // num_sequence_shards) diff --git a/megatron/core/context_parallel/utils.py b/megatron/core/context_parallel/utils.py index dfff70389c4..ea188405410 100644 --- a/megatron/core/context_parallel/utils.py +++ b/megatron/core/context_parallel/utils.py @@ -60,12 +60,12 @@ def _get_batch_on_this_cp_rank_contiguous( cp_rank = torch.distributed.get_rank(cp_group) sequence_keys = ( - 'tokens', - 'labels', - 'loss_mask', - 'position_ids', - 'decoder_input', - 'mtp_input_mask', + "tokens", + "labels", + "loss_mask", + "position_ids", + "decoder_input", + "mtp_input_mask", ) if cp_size == 1: return batch @@ -116,12 +116,12 @@ def _get_batch_on_this_cp_rank_padded_zigzag( sequence_tensor = None sequence_keys = ( - 'tokens', - 'labels', - 'loss_mask', - 'position_ids', - 'decoder_input', - 'mtp_input_mask', + "tokens", + "labels", + "loss_mask", + "position_ids", + "decoder_input", + "mtp_input_mask", ) for key in sequence_keys: sequence_tensor = batch.get(key) @@ -298,16 +298,16 @@ def get_batches_on_this_cp_rank( has_sequence_data = any( batch.get(key) is not None for key in ( - 'tokens', - 'labels', - 'loss_mask', - 'position_ids', - 'attention_mask', - 'decoder_input', - 'mtp_input_mask', + "tokens", + "labels", + "loss_mask", + "position_ids", + "attention_mask", + "decoder_input", + "mtp_input_mask", ) ) - if has_sequence_data and cp_size > 1: + if has_sequence_data and (cp_size > 1 or is_hybrid_cp): # Copy the dictionary because the CP sharder replaces sequence-valued entries in place. batches_by_layout = { layout: get_batch_on_this_cp_rank( diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py index b33da274867..cedbda90f22 100644 --- a/megatron/core/models/mimo/model/base.py +++ b/megatron/core/models/mimo/model/base.py @@ -57,7 +57,11 @@ class MimoModel(MegatronModule): """ def __init__( - self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None, tp_cp_group=None + self, + mimo_config: MimoModelConfig, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> None: """Initialize the multimodal model. diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py index 48dc5aa522d..73bbed7c0c1 100644 --- a/megatron/core/pipeline_parallel/bridge_communicator.py +++ b/megatron/core/pipeline_parallel/bridge_communicator.py @@ -4,6 +4,7 @@ from collections import deque from dataclasses import dataclass, field from enum import Enum +from math import prod from typing import Deque, Dict, List, Optional, Tuple import torch @@ -122,16 +123,15 @@ def __init__( "a bridge without backward communication requires source pipeline parallelism 1" ) - if 'cp' in self.src_grid.dim_names: - assert self.src_grid.shape[self.src_grid.dim_names.index('cp')] == 1, ( - f"Source grid CP size must be 1, got " - f"{self.src_grid.shape[self.src_grid.dim_names.index('cp')]}" - ) + src_cp_size = prod( + self.src_grid.shape[self.src_grid.dim_names.index(dim)] + for dim in self._context_parallel_dims(self.src_grid) + ) + assert src_cp_size == 1, f"Source grid CP size must be 1, got {src_cp_size}" - self.dest_cp_size = ( - self.dest_grid.shape[self.dest_grid.dim_names.index('cp')] - if 'cp' in self.dest_grid.dim_names - else 1 + self.dest_cp_size = prod( + self.dest_grid.shape[self.dest_grid.dim_names.index(dim)] + for dim in self._context_parallel_dims(self.dest_grid) ) self.current_rank = dist.get_rank() @@ -192,7 +192,7 @@ def __init__( and self.dest_local_leader_rank is not None and self.current_rank in self.dest_grid_broadcast_ranks ): - dest_cp_pg = self.dest_grid.get_pg("cp") + dest_cp_pg = self.dest_grid.get_pg(self._context_parallel_dims(self.dest_grid)) if self.dest_local_leader_rank in dist.get_process_group_ranks(dest_cp_pg): self.dest_cp_reduce_pg = dest_cp_pg @@ -289,6 +289,11 @@ def get_leader_rank(self, grid: HyperCommGrid, is_src: bool) -> List[int]: leader_ranks.extend(ranks) return leader_ranks, local_leader_rank + @staticmethod + def _context_parallel_dims(grid: HyperCommGrid) -> List[str]: + """Return every grid dimension that partitions a sample's sequence.""" + return [dim for dim in ("cp", "gtp_remat_sequence_shards") if dim in grid.dim_names] + def get_boundary_pp_stage_ranks(self, grid: HyperCommGrid, is_src: bool): """Get TP-CP ranks at boundary PP stage for each DP replica. @@ -297,7 +302,7 @@ def get_boundary_pp_stage_ranks(self, grid: HyperCommGrid, is_src: bool): """ # Get tp-cp rank enumeration (each list has same dp and pp, different tp and cp) - tpcp_rank_lists = grid._gen_rank_enum(['tp', 'cp']) + tpcp_rank_lists = grid._gen_rank_enum(["tp"] + self._context_parallel_dims(grid)) pp_size = grid.shape[grid.dim_names.index('pp')] # Determine boundary pp stage diff --git a/megatron/core/utils.py b/megatron/core/utils.py index a55dc9e6ee2..efbfdfa3d95 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2610,12 +2610,12 @@ def _get_batch_on_this_cp_rank_per_document_balancing( else batch["cu_seqlens"] )[0] SEQUENCE_KEYS = ( - 'tokens', - 'labels', - 'loss_mask', - 'position_ids', - 'decoder_input', - 'mtp_input_mask', + "tokens", + "labels", + "loss_mask", + "position_ids", + "decoder_input", + "mtp_input_mask", ) sequence_tensor = next( (batch.get(key) for key in SEQUENCE_KEYS if batch.get(key) is not None), None diff --git a/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py index eebd23bd537..798980ff0d2 100644 --- a/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py +++ b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py @@ -24,7 +24,7 @@ def setup_method(self): language_model_spec = ModuleSpec( module=MagicMock, params={ - 'config': TransformerConfig(num_layers=1, hidden_size=64, num_attention_heads=4) + "config": TransformerConfig(num_layers=1, hidden_size=64, num_attention_heads=4) }, ) self.mimo_config = MimoModelConfig( diff --git a/tests/unit_tests/models/mimo/test_mimo_gtp_cp_layout.py b/tests/unit_tests/models/mimo/test_mimo_gtp_cp_layout.py new file mode 100644 index 00000000000..f9bc260108d --- /dev/null +++ b/tests/unit_tests/models/mimo/test_mimo_gtp_cp_layout.py @@ -0,0 +1,185 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""MIMO rank ownership and bridge routing with simulated distributed groups.""" + +from types import SimpleNamespace + +import pytest +import torch.distributed as dist + +from examples.mimo.training.topology import ModuleGridSpec, _build_grid, pg_collection_from_grid +from examples.mimo.utils.hetero import get_data_lane_rank +from megatron.core.parallel_state import RankGenerator +from megatron.core.pipeline_parallel.bridge_communicator import BridgeCommunicator +from megatron.core.tensor_parallel.gtp_parallel_layout import GTPParallelLayout + + +@pytest.fixture +def rank_world(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + world = SimpleNamespace(rank=0) + monkeypatch.setenv("WORLD_SIZE", "4096") + monkeypatch.setattr( + dist, "get_rank", lambda group=None: world.rank if group is None else group.rank() + ) + monkeypatch.setattr(dist, "get_process_group_ranks", lambda group: group.ranks) + monkeypatch.setattr(dist, "barrier", lambda: None) + + def group(ranks, **kwargs): + if world.rank not in ranks: + return None + ranks = sorted(ranks) + return SimpleNamespace( + ranks=ranks, rank=lambda: ranks.index(world.rank), size=lambda: len(ranks) + ) + + monkeypatch.setattr(dist, "new_group", group) + monkeypatch.setattr( + dist, + "new_subgroups_by_enumeration", + lambda rank_lists, **kwargs: ( + next((group(ranks) for ranks in rank_lists if world.rank in ranks), None), + [], + ), + ) + # Keep the test's fake groups out of the bridge's process-wide caches. + monkeypatch.setattr(BridgeCommunicator, "_broadcast_pg_cache", {}) + monkeypatch.setattr(BridgeCommunicator, "_bridge_pg_cache", {}) + return world + + +@pytest.mark.parametrize( + "tp,pp,cp,weight,sequence_shards,dp", + [ + (1, 1, 8, 1, 1, 1), + (1, 1, 1, 4, 1, 1), + (1, 1, 1, 1, 1, 4), + (1, 1, 2, 4, 4, 1), + (1, 1, 2, 64, 64, 1), + (2, 1, 2, 32, 32, 1), + (2, 2, 1, 4, 2, 2), + (1, 2, 2, 6, 2, 2), + (1, 1, 3, 4, 1, 1), + ], +) +def test_mimo_groups_match_stock_weight_and_token_ownership( + rank_world: SimpleNamespace, + tp: int, + pp: int, + cp: int, + weight: int, + sequence_shards: int, + dp: int, +) -> None: + size = tp * pp * cp * weight * dp + spec = ModuleGridSpec( + "language", + size, + tp=tp, + pp=pp, + cp=cp, + gtp_remat=weight, + ep=2, + expt_gtp_remat=2, + rank_offset=4, + gtp_remat_num_sequence_shards=sequence_shards, + ) + if sequence_shards > 1: + layout = GTPParallelLayout( + size, tp, pp, cp, weight, num_sequence_shards=sequence_shards, rank_offset=4 + ) + data, weights = layout.data, layout.weights + else: + data = weights = RankGenerator( + tp=tp, + ep=1, + dp=dp, + pp=pp, + cp=cp, + gtp_remat=weight, + order="tp-cp-gtp_remat-dp-pp", + rank_offset=4, + ) + expected = { + "cp": data.get_ranks("cp"), + "dp_gtp_remat": data.get_ranks("dp-gtp_remat"), + "dp_cp_gtp_remat": data.get_ranks("dp-cp-gtp_remat"), + "gtp_remat": weights.get_ranks("gtp_remat"), + "dp_cp": weights.get_ranks("dp" if sequence_shards > 1 else "dp-cp"), + "mp": weights.get_ranks("tp-gtp_remat-pp"), + } + assert spec.dp == dp + for rank in range(4, 4 + size): + rank_world.rank = rank + grid = _build_grid(spec) + pgc = pg_collection_from_grid(grid) + for field, groups in expected.items(): + assert getattr(pgc, field).ranks == next(ranks for ranks in groups if rank in ranks) + assert pgc.intra_dp_cp is pgc.dp_cp + assert get_data_lane_rank(pgc) == pgc.dp_gtp_remat.ranks.index(rank) + assert pgc.ep.size() == 2 + assert pgc.expt_gtp_remat.size() == 2 + assert pgc.expt_dp.size() == size // (pp * 4) + + +@pytest.mark.parametrize( + "cp,weight,sequence_shards", [(8, 1, 1), (1, 4, 1), (2, 4, 4), (1, 4, 2), (2, 6, 2)] +) +@pytest.mark.parametrize("encoder_lanes_per_sample", [1, 2]) +def test_bridge_routes_each_sample_to_all_context_partitions( + rank_world: SimpleNamespace, + cp: int, + weight: int, + sequence_shards: int, + encoder_lanes_per_sample: int, +) -> None: + tp, pp, dp = 2, 2, 2 + size = tp * pp * cp * weight * dp + layout = GTPParallelLayout(size, tp, pp, cp, weight, num_sequence_shards=sequence_shards) + context_size = cp * sequence_shards + encoder_size = layout.batch_parallel_size * encoder_lanes_per_sample + language_spec = ModuleGridSpec( + "language", + size, + tp=tp, + pp=pp, + cp=cp, + gtp_remat=weight, + rank_offset=encoder_size, + gtp_remat_num_sequence_shards=sequence_shards, + ) + for rank in range(encoder_size, encoder_size + size // pp): + rank_world.rank = rank + # Bridge caches are rank-local in production; simulate a fresh rank each time. + BridgeCommunicator._broadcast_pg_cache.clear() + BridgeCommunicator._bridge_pg_cache.clear() + src = _build_grid(ModuleGridSpec("images", encoder_size)) + dest = _build_grid(language_spec) + pgc = pg_collection_from_grid(dest) + sample_rank = get_data_lane_rank(pgc) + bridge = BridgeCommunicator(src, dest) + assert bridge.dest_cp_size == context_size + assert len(bridge.dest_tp_leaders) == layout.batch_parallel_size + assert len(bridge.dest_grid_broadcast_ranks) == tp * context_size + assert bridge.dest_local_leader_rank == bridge.dest_tp_leaders[sample_rank] + leader_info = bridge.comm_map[bridge.dest_local_leader_rank] + assert leader_info.recv_from_ranks == list( + range( + sample_rank * encoder_lanes_per_sample, (sample_rank + 1) * encoder_lanes_per_sample + ) + ) + if context_size > 1 and pgc.tp.rank() == 0: + assert bridge.dest_cp_reduce_pg is pgc.cp + else: + assert bridge.dest_cp_reduce_pg is None + + +def test_cp_extension_keeps_mimo_weight_and_optimizer_groups(rank_world: SimpleNamespace) -> None: + memberships = [] + for cp in (1, 2, 4): + rank_world.rank = 11 + grid = _build_grid( + ModuleGridSpec("language", 16, cp=cp, gtp_remat=4, gtp_remat_num_sequence_shards=2) + ) + pgc = pg_collection_from_grid(grid) + memberships.append((pgc.gtp_remat.ranks, pgc.dp_cp.ranks, pgc.mp.ranks)) + assert memberships[0] == memberships[1] == memberships[2] diff --git a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py index 6c6702e925c..a16efaa9beb 100644 --- a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py +++ b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py @@ -13,6 +13,7 @@ build_module_grid_specs, validate_hetero_grid_args, ) +from examples.mimo.utils.hetero import get_language_batch_parallel_size from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY WORLD_SIZE_8 = 8 @@ -108,6 +109,52 @@ def test_canonical_layout_validates_and_maps_specs(): assert language_grid_spec.expt_tp == 1 +@pytest.mark.parametrize("llm_only", [False, True]) +def test_sequence_sharding_counts_only_independent_language_samples(llm_only): + args = _layout_8gpu_20l( + mimo_llm_dp=1, + mimo_llm_cp=1, + mimo_llm_ep=1, + tensor_parallel_num_weight_shards=8, + tensor_parallel_num_sequence_shards=4, + sequence_parallel=True, + mimo_llm_only=llm_only, + mimo_llm_offset=0 if llm_only else 4, + ) + world_size = 8 if llm_only else 12 + specs = build_module_grid_specs(args, world_size, encoder_module_name="images") + language = specs[-1] + assert language.num_ranks == 8 + assert language.dp == 1 + assert language.gtp_remat == 4 + assert language.gtp_remat_num_sequence_shards == 2 + assert get_language_batch_parallel_size(args) == 2 + if not llm_only: + assert specs[0].num_ranks == 4 + assert specs[0].gtp_remat_num_sequence_shards == 1 + + +@pytest.mark.parametrize("tp,weight_shards", [(1, 1), (1, 64), (2, 64)]) +def test_sequence_sharding_preserves_requested_weight_shards(tp, weight_shards): + args = _layout_8gpu_20l( + micro_batch_size=2, + mimo_llm_tp=tp, + mimo_llm_cp=128 // weight_shards, + mimo_llm_dp=1, + tensor_parallel_num_weight_shards=weight_shards, + tensor_parallel_num_sequence_shards=weight_shards, + sequence_parallel=tp > 1, + ) + encoder, language = build_module_grid_specs(args, 132, "vision") + assert encoder.num_ranks == 4 + assert language.num_ranks == 128 + assert language.cp == 128 // weight_shards + assert language.gtp_remat == weight_shards // tp + assert language.gtp_remat_num_sequence_shards == weight_shards // tp + assert language.dp == 1 + assert get_language_batch_parallel_size(args) == 1 + + def test_gtp_layout_validates_and_maps_weight_shard_axes(): args = _layout_8gpu_20l( mimo_llm_dp=1, diff --git a/tests/unit_tests/models/mimo/test_mimo_mock_data.py b/tests/unit_tests/models/mimo/test_mimo_mock_data.py index 7beff32d8b8..b14408ea090 100644 --- a/tests/unit_tests/models/mimo/test_mimo_mock_data.py +++ b/tests/unit_tests/models/mimo/test_mimo_mock_data.py @@ -53,12 +53,18 @@ def _topology(*, language_rank, encoder_rank=None, cp_rank=0, cp_size=1, gtp_ran cp=_group(rank=cp_rank, size=cp_size), dp=_group(rank=0, size=2), gtp_remat=_group(rank=gtp_rank, size=gtp_size), + dp_gtp_remat=_group(rank=gtp_rank, size=2 * gtp_size), dp_cp_gtp_remat=None, ) } if encoder_rank is not None: grids[encoder] = _grid(encoder_rank) - pgs[encoder] = SimpleNamespace(pp=_group(), dp=_group(rank=1, size=2), dp_cp_gtp_remat=None) + pgs[encoder] = SimpleNamespace( + pp=_group(), + dp=_group(rank=1, size=2), + dp_gtp_remat=_group(rank=1, size=2), + dp_cp_gtp_remat=None, + ) return SimpleNamespace(grids=grids, module_pgs=pgs) @@ -146,8 +152,7 @@ def test_cp_replicas_share_batches_without_merging_data_lanes(adapter, cp_size, topology = _topology(language_rank=True) pg = topology.module_pgs["language"] pg.cp = _group(rank=cp_rank, size=cp_size) - pg.dp = _group(rank=lane // gtp_size, size=args.mimo_llm_dp) - pg.gtp_remat = _group(rank=lane % gtp_size, size=gtp_size) + pg.dp_gtp_remat = _group(rank=lane, size=args.mimo_llm_dp * gtp_size) pg.dp_cp_gtp_remat = _group( rank=lane * cp_size + cp_rank, size=args.mimo_llm_dp * gtp_size * cp_size ) @@ -162,3 +167,34 @@ def test_cp_replicas_share_batches_without_merging_data_lanes(adapter, cp_size, for key in ("input_ids", "labels", "loss_mask", "position_ids"): assert torch.equal(batch[key], reference[key]), key assert len(set(lane_seeds)) == args.mimo_llm_dp * gtp_size + + +@pytest.mark.parametrize( + "cp_size,weight_size,sequence_shards,sample_size", + [(4, 1, 1, 2), (1, 8, 4, 4), (2, 8, 4, 4), (1, 4, 1, 8)], +) +def test_sequence_sharding_mock_data_uses_sample_rank_and_encoder_batch_size( + adapter, cp_size, weight_size, sequence_shards, sample_size +): + args = _args() + args.mimo_llm_cp = cp_size + args.gtp_weight_remat_size = weight_size + args.tensor_parallel_num_sequence_shards = sequence_shards + cp_size *= sequence_shards + seeds = [] + for lane in range(sample_size): + cp_seeds = [] + for cp_rank in range(cp_size): + topology = _topology(language_rank=True, cp_rank=cp_rank, cp_size=cp_size) + topology.module_pgs["language"].dp_gtp_remat = _group(rank=lane, size=sample_size) + loaders = adapter.build_train_valid_test_data_loaders(args, topology) + cp_seeds.append([loader.dataset.seed for loader in loaders]) + assert all(seed == cp_seeds[0] for seed in cp_seeds) + seeds.append(cp_seeds[0][0]) + assert len(set(seeds)) == sample_size + encoder_loaders = adapter.build_train_valid_test_data_loaders( + args, _topology(encoder_rank=True, language_rank=False) + ) + assert all( + loader.batch_size == args.micro_batch_size * sample_size for loader in encoder_loaders + ) diff --git a/tests/unit_tests/models/mimo/test_mimo_model.py b/tests/unit_tests/models/mimo/test_mimo_model.py index 24e4200a65c..a9104be6842 100644 --- a/tests/unit_tests/models/mimo/test_mimo_model.py +++ b/tests/unit_tests/models/mimo/test_mimo_model.py @@ -593,7 +593,7 @@ def capture_lm_forward(*args, **kwargs): captured['decoder_input'] = kwargs.get('decoder_input') captured['loss_mask'] = kwargs.get('loss_mask') captured['mtp_input_mask'] = kwargs.get('mtp_input_mask') - captured['cp_batch'] = kwargs.get('cp_batch') + captured["cp_batch"] = kwargs.get("cp_batch") return torch.zeros( self.batch_size, sharded_seq_len, self.vocab_size, device=self.device ) @@ -615,18 +615,18 @@ def capture_lm_forward(*args, **kwargs): mock_adapter.partition.assert_called_once() partition_kwargs = mock_adapter.partition.call_args.kwargs - assert partition_kwargs['embeddings'].shape == ( + assert partition_kwargs["embeddings"].shape == ( self.seq_len, self.batch_size, self.hidden_size, ) - assert partition_kwargs['input_ids'] is input_ids - assert partition_kwargs['position_ids'] is position_ids - assert partition_kwargs['loss_mask'] is loss_mask - assert partition_kwargs['cu_seqlens'] is None - assert partition_kwargs['cu_seqlens_padded'] is None - assert partition_kwargs['max_seqlen'] is None - torch.testing.assert_close(partition_kwargs['mtp_input_mask'], text_mask) + assert partition_kwargs["input_ids"] is input_ids + assert partition_kwargs["position_ids"] is position_ids + assert partition_kwargs["loss_mask"] is loss_mask + assert partition_kwargs["cu_seqlens"] is None + assert partition_kwargs["cu_seqlens_padded"] is None + assert partition_kwargs["max_seqlen"] is None + torch.testing.assert_close(partition_kwargs["mtp_input_mask"], text_mask) # The adapter's LM-layout output and dual-layout CP batch flow straight into the LM. assert captured['decoder_input'].shape == ( sharded_seq_len, @@ -638,7 +638,7 @@ def capture_lm_forward(*args, **kwargs): assert captured['position_ids'] is sharded_position_ids assert captured['mtp_input_mask'].dtype == torch.bool assert torch.equal(captured['mtp_input_mask'], text_mask[:, :sharded_seq_len]) - assert captured['cp_batch'] is cp_batch + assert captured["cp_batch"] is cp_batch # forward() returns the adapter's possibly sharded loss mask. assert out_loss_mask is sharded_loss_mask diff --git a/tests/unit_tests/models/mimo/test_mimo_partition.py b/tests/unit_tests/models/mimo/test_mimo_partition.py index c86a8f605e3..e35d21c2ec3 100644 --- a/tests/unit_tests/models/mimo/test_mimo_partition.py +++ b/tests/unit_tests/models/mimo/test_mimo_partition.py @@ -208,7 +208,7 @@ def test_thd_format_skips_divisibility_check(self): with ( patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), patch( - 'megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank', + "megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank", return_value=cp_batch, ), ): @@ -233,12 +233,12 @@ def test_none_embeddings_skips_shard_factor_check(self): labels = torch.randint(0, 100, (2, 7)) loss_mask = torch.ones(2, 7) cp_batch = ContextParallelBatch.from_single_layout( - "zigzag", {'labels': labels[:, :4], 'loss_mask': loss_mask[:, :4]}, None + "zigzag", {"labels": labels[:, :4], "loss_mask": loss_mask[:, :4]}, None ) with ( patch('megatron.core.models.mimo.partition.utils.get_pg_size', return_value=2), patch( - 'megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank', + "megatron.core.models.mimo.partition.utils.get_batches_on_this_cp_rank", return_value=cp_batch, ), ): diff --git a/tests/unit_tests/models/mimo/test_nemotron_moe_vlm_provider.py b/tests/unit_tests/models/mimo/test_nemotron_moe_vlm_provider.py index 892eab22190..bc0f4eedbdd 100644 --- a/tests/unit_tests/models/mimo/test_nemotron_moe_vlm_provider.py +++ b/tests/unit_tests/models/mimo/test_nemotron_moe_vlm_provider.py @@ -609,12 +609,14 @@ def test_language_rank_placement_uses_language_parallelism(): language_config = nemotron_language_config( args, tp_size=2, pp_size=1, ep_size=1, expt_tp_size=1 ) + assert language_config.gtp_remat_num_sequence_shards == 1 language_spec = ModuleSpec(module=object, params={"config": language_config}) projection = language_input_projection_spec(args, None, None, language_spec) assert projection.params["input_size"] == 5120 assert projection.params["config"].tensor_model_parallel_size == 2 assert projection.params["config"].gtp_weight_remat_size == 2 + assert projection.params["config"].gtp_remat_num_sequence_shards == 1 provider = resolve_provider(args) assert provider.language_input_projection_specs[RADIO_ENCODER_MODULE_NAME] is ( language_input_projection_spec @@ -627,7 +629,8 @@ def test_language_rank_placement_uses_language_parallelism(): @pytest.mark.parametrize("cp_size", [1, 2, 4]) @pytest.mark.parametrize("use_groups", [False, True]) -def test_language_cp_comes_from_its_grid(monkeypatch, cp_size, use_groups): +@pytest.mark.parametrize("sequence_shards", [1, 2]) +def test_language_cp_comes_from_its_grid(monkeypatch, cp_size, use_groups, sequence_shards): from types import SimpleNamespace from examples.mimo.model_providers import nemotron_moe_vlm as provider @@ -636,6 +639,8 @@ def test_language_cp_comes_from_its_grid(monkeypatch, cp_size, use_groups): args = SimpleNamespace( mimo_llm_ep=1, mimo_llm_expt_tp=1, + tensor_parallel_num_weight_shards=sequence_shards, + tensor_parallel_num_sequence_shards=sequence_shards, vocab_size=64, seq_length=32, hybrid_layer_pattern="*", @@ -644,9 +649,17 @@ def test_language_cp_comes_from_its_grid(monkeypatch, cp_size, use_groups): monkeypatch.setattr( provider, "_base_config", - lambda args: SimpleNamespace(context_parallel_size=8, calculate_per_token_loss=True), + lambda args: SimpleNamespace( + context_parallel_size=8, + calculate_per_token_loss=True, + gtp_remat_num_sequence_shards=sequence_shards, + ), ) grid = SimpleNamespace(shape=[1, cp_size, 1], dim_names=["tp", "cp", "pp"]) + if sequence_shards > 1: + grid.shape.append(sequence_shards) + grid.dim_names.append("gtp_remat_sequence_shards") + cp_size *= sequence_shards groups = None if use_groups: groups = SimpleNamespace( @@ -667,7 +680,39 @@ def test_encoder_and_projection_do_not_inherit_language_cp(): args = _parse_validate(_build_argv(*_PRESET_20L)) args.context_parallel_size = 2 + args.tensor_parallel_num_weight_shards = 2 + args.tensor_parallel_num_sequence_shards = 2 vision = radio_vision_config(args, tp_size=1, pp_size=1) projection = nemotron_projection_config(args, tp_size=1, projection_input_size=5120) assert vision.context_parallel_size == 1 assert projection.context_parallel_size == 1 + assert vision.gtp_remat_num_sequence_shards == 1 + assert projection.gtp_remat_num_sequence_shards == 1 + + +@pytest.mark.parametrize("tp", [1, 2]) +@pytest.mark.parametrize("sequence_shards", [1, 4]) +def test_language_sequence_shards_are_excluded_from_projection(tp, sequence_shards): + from examples.mimo.model_providers.nemotron_moe_vlm import ( + nemotron_language_config, + nemotron_projection_config, + ) + + args = _parse_validate(_build_argv(*_PRESET_20L)) + args.tensor_model_parallel_size = tp + args.sequence_parallel = tp > 1 + args.tensor_parallel_num_weight_shards = 8 + args.tensor_parallel_num_sequence_shards = tp * sequence_shards + args.context_parallel_size = 2 * sequence_shards + + language = nemotron_language_config( + args, tp, pp_size=1, ep_size=1, expt_tp_size=1, cp_size=args.context_parallel_size + ) + assert language.gtp_remat_num_sequence_shards == sequence_shards + assert language.context_parallel_size == 2 * sequence_shards + projection = nemotron_projection_config( + args, tp, projection_input_size=5120, base_config=language + ) + assert projection.gtp_remat_num_sequence_shards == 1 + assert projection.context_parallel_size == 1 + assert projection.gtp_weight_remat_size == language.gtp_weight_remat_size diff --git a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py index 2e213547a27..0a229d9f8a9 100644 --- a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py +++ b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py @@ -507,10 +507,21 @@ def test_send_backward_recv_backward(self): 512, ), f"Expected gradient shape {(4, 128, 512)}, got {received_gradient.shape}" - def test_cp_gradient_reconstruction(self): + @pytest.mark.parametrize("sequence_sharding", [False, True]) + def test_cp_gradient_reconstruction(self, sequence_sharding): """Complementary destination-CP gradients are summed before fan-in splitting.""" src_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=4) - dest_grid = create_hypercomm_grid(offset=4, tp=2, cp=2, pp=1, dp=1) + if sequence_sharding: + from examples.mimo.training.topology import ModuleGridSpec, _build_grid + + dest_grid = _build_grid( + ModuleGridSpec( + "language", 4, cp=2, gtp_remat=2, rank_offset=4, gtp_remat_num_sequence_shards=2 + ) + ) + _active_grids.append(dest_grid) + else: + dest_grid = create_hypercomm_grid(offset=4, tp=2, cp=2, pp=1, dp=1) bridge = BridgeCommunicator(src_grid, dest_grid, comm_dtype=torch.float32, tensor_ndim=2) rows_per_chunk = 2 @@ -522,8 +533,10 @@ def test_cp_gradient_reconstruction(self): source_values_match = True if bridge.is_current_rank_in_grid(dest_grid): - cp_rank = dist.get_rank(group=dest_grid.get_pg("cp")) - owned_chunks = (0, 3) if cp_rank == 0 else (1, 2) + cp_rank = dist.get_rank( + group=dest_grid.get_pg(bridge._context_parallel_dims(dest_grid)) + ) + owned_chunks = (cp_rank,) if sequence_sharding else ((0, 3) if cp_rank == 0 else (1, 2)) local_gradient = torch.zeros(payload_shape, device="cuda", dtype=torch.float32) for chunk_idx in owned_chunks: chunk_start = chunk_idx * rows_per_chunk