diff --git a/experimental/lite/megatron/lite/primitive/parallel/cp.py b/experimental/lite/megatron/lite/primitive/parallel/cp.py index ab4b1049248..0eff4fc7ebd 100644 --- a/experimental/lite/megatron/lite/primitive/parallel/cp.py +++ b/experimental/lite/megatron/lite/primitive/parallel/cp.py @@ -310,8 +310,8 @@ def get_thd_context_parallel_rank_indices( The returned indices are ordered exactly as the rank-local THD tensor is stored. ``"zigzag"`` follows Megatron's per-sequence load-balanced chunk order; ``"contiguous"`` - partitions the flattened packed THD buffer into rank-contiguous spans. Mirrors upstream - Megatron ``context_parallel_layout.get_thd_context_parallel_rank_indices``. + partitions the flattened packed THD buffer into rank-contiguous spans. This matches the + rank-index semantics used by Megatron's THD CP route builder. """ if layout not in ("zigzag", "contiguous"): raise ValueError(f"Unsupported context-parallel layout {layout!r}.") @@ -384,9 +384,8 @@ def _zigzag_contiguous_thd_swap( The packed THD tensor stays packed: we first group local tokens by their target CP rank, exchange those groups once, then scatter received tokens back into the - target rank-local order. Mirrors upstream Megatron - ``context_parallel_layout._zigzag_contiguous_thd_swap`` (packing-aware routing - from the *global* ``cu_seqlens``). + target rank-local order. This matches Megatron's route-based THD layout conversion + semantics in ``megatron.core.context_parallel_layout.conversion``. """ cp_size = dist.get_world_size(cp_group) if cp_group is not None else 1 if cp_size <= 1: diff --git a/megatron/core/context_parallel_layout.py b/megatron/core/context_parallel_layout.py deleted file mode 100644 index 44014581fd5..00000000000 --- a/megatron/core/context_parallel_layout.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - -"""Context parallel tensor layout helpers.""" - -from typing import List, Optional, Tuple - -import torch - -from megatron.core.tensor_parallel import all_to_all - - -def get_thd_context_parallel_rank_indices( - cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, layout: str -) -> torch.Tensor: - """Return global THD token indices owned by one CP rank in a layout. - - Args: - cu_seqlens: Global packed-sequence cumulative lengths before CP partitioning. - cp_size: Context-parallel group size. - cp_rank: Context-parallel rank. - layout: Either ``"zigzag"`` or ``"contiguous"``. - - The returned indices are ordered exactly as the rank-local THD tensor is stored. - ``"zigzag"`` follows Megatron's per-sequence load-balanced chunk order; ``"contiguous"`` - partitions the flattened packed THD buffer into rank-contiguous spans. - """ - if layout not in ("zigzag", "contiguous"): - raise ValueError(f"Unsupported context-parallel layout {layout!r}.") - if cp_size < 1: - raise ValueError(f"cp_size must be >= 1, got {cp_size}.") - if not 0 <= cp_rank < cp_size: - raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}.") - if cu_seqlens.dim() != 1: - raise ValueError(f"cu_seqlens must be 1-D, got shape {tuple(cu_seqlens.shape)}.") - - cu = cu_seqlens.to(dtype=torch.long) - if cu.numel() == 0 or cu[0].item() != 0: - raise ValueError(f"cu_seqlens must start at 0, got {cu_seqlens}.") - - if torch.any(torch.diff(cu) < 0): - raise ValueError(f"cu_seqlens must be nondecreasing, got {cu_seqlens}.") - - nonduplicate_boundaries = torch.ones(cu.numel(), device=cu.device, dtype=torch.bool) - nonduplicate_boundaries[1:] = cu[1:] != cu[:-1] - cu = cu[nonduplicate_boundaries] - - total_tokens = int(cu[-1].item()) - positions = torch.arange(total_tokens, device=cu.device, dtype=torch.long) - if total_tokens == 0: - return positions - - seq_lens = torch.diff(cu) - chunk_divisor = 2 * cp_size - if torch.any(seq_lens % chunk_divisor != 0): - raise ValueError( - "All packed sequence lengths must be divisible by " - f"2 * cp_size ({chunk_divisor}) for zigzag/contiguous CP layout conversion, " - f"got {seq_lens}." - ) - - if layout == "contiguous": - part_len = total_tokens // cp_size - rank_start = cp_rank * part_len - return positions[rank_start : rank_start + part_len] - - seq_idx = torch.bucketize(positions, cu[1:], right=True) - global_starts = cu[:-1] - pos_in_seq = positions - global_starts[seq_idx] - chunk_lens = (seq_lens // chunk_divisor)[seq_idx] - chunk = pos_in_seq // chunk_lens - offset = pos_in_seq - chunk * chunk_lens - - owner = torch.where(chunk < cp_size, chunk, 2 * cp_size - chunk - 1) - local_slot = torch.where(chunk < cp_size, torch.zeros_like(chunk), torch.ones_like(chunk)) - - local_starts = (global_starts // cp_size)[seq_idx] - local_pos = local_starts + local_slot * chunk_lens + offset - - rank_mask = owner == cp_rank - rank_positions = positions[rank_mask] - rank_local_pos = local_pos[rank_mask] - return rank_positions[torch.argsort(rank_local_pos)] - - -def zigzag_to_contiguous_chunks( - x: torch.Tensor, - cp_group: torch.distributed.ProcessGroup, - seq_dim: int = 0, - cu_seqlens: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Permute CP chunks from Megatron zigzag layout to contiguous-time layout. - - SBHD tensors have two equal chunks per rank along ``seq_dim`` and use a - chunk-level all-to-all. THD tensors pass global ``cu_seqlens`` and use one - packed-token all-to-all over the whole local THD tensor. - """ - if cu_seqlens is not None: - return _zigzag_contiguous_thd_swap( - x, cp_group, seq_dim, cu_seqlens, source_layout="zigzag", target_layout="contiguous" - ) - return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=True) - - -def contiguous_to_zigzag_chunks( - x: torch.Tensor, - cp_group: torch.distributed.ProcessGroup, - seq_dim: int = 0, - cu_seqlens: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Inverse of :func:`zigzag_to_contiguous_chunks`.""" - if cu_seqlens is not None: - return _zigzag_contiguous_thd_swap( - x, cp_group, seq_dim, cu_seqlens, source_layout="contiguous", target_layout="zigzag" - ) - return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=False) - - -def _zigzag_contiguous_thd_swap( - x: torch.Tensor, - cp_group: Optional[torch.distributed.ProcessGroup], - seq_dim: int, - cu_seqlens: torch.Tensor, - source_layout: str, - target_layout: str, -) -> torch.Tensor: - """Single-all-to-all THD permutation between zigzag and contiguous layouts. - - The packed THD tensor stays packed: we first group local tokens by their - target CP rank, exchange those groups once, then scatter received tokens - back into the target rank-local order. - """ - cp_size = cp_group.size() if cp_group is not None else 1 - if cp_size == 1: - return x - cp_rank = cp_group.rank() - - if seq_dim != 0: - x = x.movedim(seq_dim, 0) - x = x.contiguous() - - cu = cu_seqlens.to(device=x.device, dtype=torch.long) - # TODO: Let a future CP layout scheduler precompute this routing once per - # microbatch from immutable cu_seqlens and pass it through both THD swaps. - # Do not cache it across microbatches because packed sequence boundaries change. - source_by_rank = [ - get_thd_context_parallel_rank_indices(cu, cp_size, rank, source_layout) - for rank in range(cp_size) - ] - target_by_rank = [ - get_thd_context_parallel_rank_indices(cu, cp_size, rank, target_layout) - for rank in range(cp_size) - ] - - local_source_indices = source_by_rank[cp_rank] - local_target_indices = target_by_rank[cp_rank] - if x.size(0) != local_source_indices.numel(): - raise ValueError( - f"Local THD tensor length ({x.size(0)}) does not match {source_layout} " - f"rank-{cp_rank} partition length ({local_source_indices.numel()})." - ) - - total_tokens = int(cu[-1].item()) - target_owner = torch.empty(total_tokens, device=x.device, dtype=torch.long) - target_local_pos = torch.empty(total_tokens, device=x.device, dtype=torch.long) - for rank, indices in enumerate(target_by_rank): - target_owner[indices] = rank - target_local_pos[indices] = torch.arange(indices.numel(), device=x.device) - - local_target_owner = target_owner[local_source_indices] - local_target_pos = target_local_pos[local_source_indices] - - send_parts: List[torch.Tensor] = [] - input_split_sizes: List[int] = [] - for dst_rank in range(cp_size): - dst_mask = local_target_owner == dst_rank - dst_rows = dst_mask.nonzero(as_tuple=False).flatten() - if dst_rows.numel() > 0: - dst_rows = dst_rows[torch.argsort(local_target_pos[dst_rows])] - send_part = x.index_select(0, dst_rows) - else: - send_part = x.narrow(0, 0, 0) - send_parts.append(send_part) - input_split_sizes.append(send_part.size(0)) - send_buf = torch.cat(send_parts, dim=0).contiguous() - - output_split_sizes: List[int] = [] - recv_target_positions: List[torch.Tensor] = [] - for src_rank in range(cp_size): - src_indices = source_by_rank[src_rank] - src_to_this_rank = target_owner[src_indices] == cp_rank - recv_global_indices = src_indices[src_to_this_rank] - if recv_global_indices.numel() > 0: - recv_positions = target_local_pos[recv_global_indices] - recv_positions = recv_positions[torch.argsort(recv_positions)] - else: - recv_positions = local_target_indices.narrow(0, 0, 0) - recv_target_positions.append(recv_positions) - output_split_sizes.append(recv_positions.numel()) - - recv_buf = all_to_all(cp_group, send_buf, output_split_sizes, input_split_sizes) - - out_shape = (local_target_indices.numel(),) + tuple(x.shape[1:]) - out = x.new_empty(out_shape) - offset = 0 - for recv_positions in recv_target_positions: - recv_len = recv_positions.numel() - if recv_len > 0: - out[recv_positions] = recv_buf[offset : offset + recv_len] - offset += recv_len - - if seq_dim != 0: - out = out.movedim(0, seq_dim) - return out.contiguous() - - -def _zigzag_contiguous_chunk_swap( - x: torch.Tensor, - cp_group: Optional[torch.distributed.ProcessGroup], - seq_dim: int, - to_contiguous: bool, -) -> torch.Tensor: - """Single-all-to-all chunk permutation between zigzag and contiguous layouts. - - Each rank holds exactly two chunks along ``seq_dim``. The mapping from - local (rank, slot) to (rank, slot) in the target layout is deterministic - and depends only on ``cp_size`` and ``cp_rank``, so we pack send data in - destination-rank order and use one ``all_to_all_single`` with unequal - splits to route each chunk to its target rank. - """ - cp_size = cp_group.size() if cp_group is not None else 1 - if cp_size == 1: - return x - cp_rank = cp_group.rank() - - # Work with seq_dim at position 0. - if seq_dim != 0: - x = x.movedim(seq_dim, 0) - x = x.contiguous() - - seq_len_local = x.size(0) - assert seq_len_local % 2 == 0, ( - f"zigzag/contiguous chunk swap requires an even local sequence length, " - f"got {seq_len_local}." - ) - chunk_len = seq_len_local // 2 - - def _rank_to_chunks(rank: int, in_zigzag: bool) -> Tuple[int, int]: - """Global chunk indices at (slot 0, slot 1) for this rank.""" - if in_zigzag: - return (rank, 2 * cp_size - rank - 1) - return (2 * rank, 2 * rank + 1) - - def _chunk_to_dest(chunk_idx: int, target_zigzag: bool) -> Tuple[int, int]: - """Destination (rank, slot) for a given global chunk index in the target layout.""" - if target_zigzag: - if chunk_idx < cp_size: - return chunk_idx, 0 - return 2 * cp_size - chunk_idx - 1, 1 - return chunk_idx // 2, chunk_idx % 2 - - source_in_zigzag = to_contiguous - target_in_zigzag = not to_contiguous - - local_chunk_indices = _rank_to_chunks(cp_rank, source_in_zigzag) - local_dests = [_chunk_to_dest(c, target_in_zigzag) for c in local_chunk_indices] - - # Pack the send buffer so chunks are ordered by (dst_rank, dst_slot). - local_slot_order = sorted(range(2), key=lambda s: local_dests[s]) - local_chunks = [x[:chunk_len], x[chunk_len:]] - send_buf = torch.cat([local_chunks[s] for s in local_slot_order], dim=0).contiguous() - - input_split_chunks = [0] * cp_size - for dst_rank, _ in local_dests: - input_split_chunks[dst_rank] += 1 - - # Mirror every source rank's packing logic so we know which received chunk - # belongs in which local target slot. - output_split_chunks = [0] * cp_size - recv_dst_slots_per_source: List[List[int]] = [[] for _ in range(cp_size)] - for src in range(cp_size): - src_chunks = _rank_to_chunks(src, source_in_zigzag) - src_dests = [_chunk_to_dest(c, target_in_zigzag) for c in src_chunks] - src_slot_order = sorted(range(2), key=lambda s: src_dests[s]) - for s in src_slot_order: - dst_rank, dst_slot = src_dests[s] - if dst_rank == cp_rank: - output_split_chunks[src] += 1 - recv_dst_slots_per_source[src].append(dst_slot) - - input_split_sizes = [n * chunk_len for n in input_split_chunks] - output_split_sizes = [n * chunk_len for n in output_split_chunks] - - recv_buf = all_to_all(cp_group, send_buf, output_split_sizes, input_split_sizes) - - # Reassemble local chunks in target-layout slot order. - target_slots: List[Optional[torch.Tensor]] = [None, None] - offset = 0 - for src in range(cp_size): - for dst_slot in recv_dst_slots_per_source[src]: - target_slots[dst_slot] = recv_buf[offset : offset + chunk_len] - offset += chunk_len - assert all(t is not None for t in target_slots), "Incomplete chunk reassembly in CP swap" - - out = torch.cat(target_slots, dim=0) - if seq_dim != 0: - out = out.movedim(0, seq_dim) - return out.contiguous() diff --git a/megatron/core/context_parallel_layout/__init__.py b/megatron/core/context_parallel_layout/__init__.py new file mode 100644 index 00000000000..06bb67e652c --- /dev/null +++ b/megatron/core/context_parallel_layout/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Public context parallel sequence partition-mode APIs. + +The implementation is split by responsibility; internal conversion and route-building +helpers remain in their respective submodules rather than being re-exported here. + +Ownership summary: + +- model builders choose the pipeline-stage input CP layout; +- blocks convert rank-local sequence tensors between layer preferences; +- model postprocess restores the public output boundary to the input layout; +- MTP validates its inner-layer layout preference but does not own outer conversion. +""" + +from megatron.core.context_parallel_layout.conversion import ( + CpPartitionModeConverter, + convert_module_input_tensors_cp_partition_mode, +) +from megatron.core.context_parallel_layout.routes import prebuild_thd_cp_partition_routes +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import finalize_packed_seq_params + +__all__ = [ + "CpPartitionMode", + "CpPartitionModeConverter", + "ThdCpRoute", + "convert_module_input_tensors_cp_partition_mode", + "finalize_packed_seq_params", + "prebuild_thd_cp_partition_routes", +] diff --git a/megatron/core/context_parallel_layout/conversion.py b/megatron/core/context_parallel_layout/conversion.py new file mode 100644 index 00000000000..2763f2a8f06 --- /dev/null +++ b/megatron/core/context_parallel_layout/conversion.py @@ -0,0 +1,668 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Tensor operations for converting between CP partition modes.""" + +import warnings +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union, cast + +import torch + +from megatron.core.context_parallel_layout.routes import ( + build_thd_cp_partition_route, + get_thd_cp_partition_route, +) +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import ( + get_packed_seq_params_cp_partition_cu_seqlens, +) +from megatron.core.tensor_parallel.mappings import all_to_all +from megatron.core.utils import nvtx_range + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + + +class CpPartitionModeConverter: + """Convert tensors across one CP layout edge.""" + + def __init__( + self, + *, + packed_seq_params: Optional["PackedSeqParams"], + source_partition_mode: CpPartitionMode, + target_partition_mode: CpPartitionMode, + config: Any, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + ) -> None: + self.cp_group = cp_group + self.packed_seq_params = packed_seq_params + self.source_partition_mode = source_partition_mode + self.target_partition_mode = target_partition_mode + self.config = config + self.tp_group = tp_group + self.tp_cp_group = tp_cp_group + if ( + self.conversion_needed + and getattr(self.packed_seq_params, "qkv_format", None) == "thd" + and self.config.cuda_graph_impl == "full_iteration" + ): + raise ValueError( + "Full-iteration CUDA graph is not supported for THD CP layout conversion: " + f"source={self.source_partition_mode!r}, target={self.target_partition_mode!r}." + ) + + @property + def conversion_needed(self) -> bool: + """Return whether this edge needs a real layout conversion.""" + return ( + self.source_partition_mode != self.target_partition_mode + and self.cp_group is not None + and self.cp_group.size() > 1 + ) + + def assert_no_dense_attention_inputs( + self, + *, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + hidden_states: Optional[torch.Tensor] = None, + ) -> None: + """Reject dense attention tensors when this edge would reorder tokens.""" + if not self.conversion_needed: + return + if attention_mask is not None: + self._raise_unsupported_dense_attention( + "an explicit attention_mask", hidden_states=hidden_states + ) + if attention_bias is not None: + self._raise_unsupported_dense_attention("attention_bias", hidden_states=hidden_states) + + def convert( + self, + value: Any, + *, + seq_dim: Union[int, Callable[[torch.Tensor], int]] = 0, + sequence_parallel: bool = False, + ) -> Any: + """Convert a tensor or nested tensor container across this layout edge.""" + if not self.conversion_needed or value is None: + return value + # Nested values may contain optional tensors; traverse containers while + # preserving their original shape. + if isinstance(value, tuple): + return tuple( + self.convert(part, seq_dim=seq_dim, sequence_parallel=sequence_parallel) + for part in value + ) + if isinstance(value, list): + return [ + self.convert(part, seq_dim=seq_dim, sequence_parallel=sequence_parallel) + for part in value + ] + if not torch.is_tensor(value): + return value + + resolved_seq_dim = seq_dim(value) if callable(seq_dim) else seq_dim + converted = convert_cp_partition_mode( + x=value, + source_partition_mode=self.source_partition_mode, + target_partition_mode=self.target_partition_mode, + seq_dim=resolved_seq_dim, + cu_seqlens=get_packed_seq_params_cp_partition_cu_seqlens(self.packed_seq_params), + sequence_parallel=sequence_parallel, + cp_group=self.cp_group, + tp_group=self.tp_group, + tp_cp_group=self.tp_cp_group, + thd_cp_partition_route=get_thd_cp_partition_route( + self.packed_seq_params, self.source_partition_mode, self.target_partition_mode + ), + ) + if self.packed_seq_params is not None: + self.packed_seq_params.cp_partition_mode = self.target_partition_mode + return converted + + def _raise_unsupported_dense_attention( + self, tensor_name: str, *, hidden_states: Optional[torch.Tensor] + ) -> None: + hidden_shape = tuple(hidden_states.shape) if hidden_states is not None else None + raise NotImplementedError( + "Changing CP partition mode with " + f"{tensor_name} is not supported yet: " + f"source={self.source_partition_mode!r}, " + f"target={self.target_partition_mode!r}, " + f"qkv_format={getattr(self.packed_seq_params, 'qkv_format', None)!r}, " + f"hidden_shape={hidden_shape}." + ) + + +def convert_module_input_tensors_cp_partition_mode( + *, + hidden_states: torch.Tensor, + packed_seq_params: Optional["PackedSeqParams"], + target_partition_mode: CpPartitionMode, + sequence_parallel: bool, + config: Any, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + key_value_states: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[CpPartitionModeConverter]]: + """Convert a module's rank-local sequence tensors to a target CP layout. + + This helper performs the common "entry conversion" pattern used by modules + that need to consume a different CP layout than their caller supplied. It + returns a converter for the opposite edge so the module output can be + converted back to the original input layout. + """ + if cp_group is None or cp_group.size() <= 1: + return hidden_states, None + + source_partition_mode = getattr(config, "cp_partition_mode", None) + if source_partition_mode is None: + raise ValueError( + "config.cp_partition_mode is required before module input CP layout conversion when " + "context parallelism is active." + ) + if source_partition_mode == target_partition_mode: + return hidden_states, None + + input_to_target_converter = CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode=source_partition_mode, + target_partition_mode=target_partition_mode, + config=config, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + input_to_target_converter.assert_no_dense_attention_inputs( + attention_mask=attention_mask, attention_bias=attention_bias, hidden_states=hidden_states + ) + if key_value_states is not None: + raise NotImplementedError( + "Changing CP partition mode with cross-attention key/value states is not supported " + f"yet: source={source_partition_mode!r}, target={target_partition_mode!r}." + ) + hidden_states = input_to_target_converter.convert( + value=hidden_states, seq_dim=0, sequence_parallel=sequence_parallel + ) + + target_to_input_converter = CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode=target_partition_mode, + target_partition_mode=source_partition_mode, + config=config, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + return (hidden_states, target_to_input_converter) + + +def convert_cp_partition_mode( + x: torch.Tensor, + *, + source_partition_mode: Optional[str], + target_partition_mode: Optional[str], + seq_dim: int = 0, + cu_seqlens: Optional[torch.Tensor] = None, + sequence_parallel: bool = False, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + thd_cp_partition_route: Optional[ThdCpRoute] = None, +) -> torch.Tensor: + """Convert a sequence tensor between CP zigzag and contiguous layouts. + + SBHD tensors use one unified all-to-all-v redistribution path over CP or + TPxCP. THD tensors use their packed-token CP route and, when sequence + parallelism shards the packed sequence, retain the naive TP gather/scatter + fallback. + """ + + if source_partition_mode == target_partition_mode: + return x + + cp_size = cp_group.size() if cp_group is not None else 1 + if cp_size == 1: + return x + assert cp_group is not None + + if source_partition_mode not in ("zigzag", "contiguous") or target_partition_mode not in ( + "zigzag", + "contiguous", + ): + cp_rank = cp_group.rank() if cp_group is not None else 0 + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r}; " + f"shape={tuple(x.shape)}, seq_dim={seq_dim}, cp_size={cp_size}, cp_rank={cp_rank}." + ) + source_layout = cast(CpPartitionMode, source_partition_mode) + target_layout = cast(CpPartitionMode, target_partition_mode) + + if cu_seqlens is None: + moved = x.movedim(seq_dim, 0) if seq_dim != 0 else x + converted = _redistribute_sbhd_layout( + input_=moved, + cp_group=cp_group, + source_layout=source_layout, + target_layout=target_layout, + sequence_parallel=sequence_parallel, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + return converted.movedim(0, seq_dim).contiguous() if seq_dim != 0 else converted + + if sequence_parallel and tp_group is not None and tp_group.size() > 1: + from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, + ) + + # TODO(yuzhongw): replace the naive THD TP gather -> CP all-to-all -> TP scatter + # fallback with a direct packed THD TPxCP redistribution path. + warnings.warn( + "THD CP layout conversion with sequence parallelism uses the naive " + "TP gather -> CP all-to-all -> TP scatter fallback.", + RuntimeWarning, + stacklevel=2, + ) + moved = x.movedim(seq_dim, 0) if seq_dim != 0 else x + # This gather is only used to run a duplicated CP layout permutation before + # scattering back to SP shards. Its backward must split, not reduce-scatter; + # otherwise every TP rank contributes the same full-sequence gradient. + gathered = gather_from_sequence_parallel_region( + input_=moved, tensor_parallel_output_grad=False, group=tp_group + ) + converted = _redistribute_thd_layout( + x=gathered, + cp_group=cp_group, + seq_dim=0, + cu_seqlens=cu_seqlens, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + thd_cp_partition_route=thd_cp_partition_route, + ) + scattered = scatter_to_sequence_parallel_region(input_=converted, group=tp_group) + return scattered.movedim(0, seq_dim).contiguous() if seq_dim != 0 else scattered + + return _redistribute_thd_layout( + x=x, + cp_group=cp_group, + seq_dim=seq_dim, + cu_seqlens=cu_seqlens, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + thd_cp_partition_route=thd_cp_partition_route, + ) + + +def _pack_thd_cp_route_send_buffer( + x: torch.Tensor, send_index: Optional[torch.Tensor] +) -> torch.Tensor: + if send_index is None: + return x + return x.index_select(0, send_index) + + +def _scatter_thd_cp_route_recv_buffer( + recv_buf: torch.Tensor, recv_index: Optional[torch.Tensor], out_shape: Tuple[int, ...] +) -> torch.Tensor: + if recv_index is None: + return recv_buf + out = recv_buf.new_empty(out_shape) + if recv_index.numel() > 0: + out.index_copy_(0, recv_index, recv_buf) + return out + + +def _redistribute_thd_layout( + x: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], + seq_dim: int, + cu_seqlens: torch.Tensor, + source_partition_mode: str, + target_partition_mode: str, + thd_cp_partition_route: Optional[ThdCpRoute] = None, +) -> torch.Tensor: + """Single-all-to-all THD permutation between zigzag and contiguous layouts. + + The packed THD tensor stays packed: we first group local tokens by their + target CP rank, exchange those groups once, then scatter received tokens + back into the target rank-local order. + """ + cp_size = cp_group.size() if cp_group is not None else 1 + if cp_size == 1: + return x + assert cp_group is not None + cp_rank = cp_group.rank() + conversion_name = f"{source_partition_mode}_to_{target_partition_mode}" + with nvtx_range(f"cp_layout/thd/swap/{conversion_name}"): + if seq_dim != 0: + x = x.movedim(seq_dim, 0) + x = x.contiguous() + + route = thd_cp_partition_route + if route is None: + route = build_thd_cp_partition_route( + cu_seqlens=cu_seqlens, cp_size=cp_size, cp_rank=cp_rank, device=x.device + ) + + if source_partition_mode == "zigzag" and target_partition_mode == "contiguous": + send_index = route.zigzag_index + recv_index = route.contiguous_index + input_split_sizes = route.zigzag_split_sizes + output_split_sizes = route.contiguous_split_sizes + elif source_partition_mode == "contiguous" and target_partition_mode == "zigzag": + send_index = route.contiguous_index + recv_index = route.zigzag_index + input_split_sizes = route.contiguous_split_sizes + output_split_sizes = route.zigzag_split_sizes + else: + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r} for THD route." + ) + + local_source_length = sum(input_split_sizes) + local_target_length = sum(output_split_sizes) + + if x.size(0) != local_source_length: + raise ValueError( + f"Local THD tensor length ({x.size(0)}) does not match {source_partition_mode} " + f"rank-{cp_rank} partition length ({local_source_length})." + ) + if local_target_length != x.size(0): + raise ValueError( + "THD CP layout conversion must preserve the local token count, " + f"got source={local_source_length}, target={local_target_length}, " + f"cp_size={cp_size}, cp_rank={cp_rank}, " + f"source_layout={source_partition_mode!r}, " + f"target_layout={target_partition_mode!r}." + ) + + with nvtx_range(f"cp_layout/thd/pack/{conversion_name}"): + send_buf = _pack_thd_cp_route_send_buffer(x=x, send_index=send_index) + if not send_buf.is_contiguous(): + send_buf = send_buf.contiguous() + + with nvtx_range(f"cp_layout/thd/all_to_all/{conversion_name}"): + recv_buf = all_to_all( + group=cp_group, + input_=send_buf, + output_split_sizes_=output_split_sizes, + input_split_sizes=input_split_sizes, + ) + + with nvtx_range(f"cp_layout/thd/scatter/{conversion_name}"): + out_shape = (local_target_length,) + tuple(x.shape[1:]) + out = _scatter_thd_cp_route_recv_buffer( + recv_buf=recv_buf, recv_index=recv_index, out_shape=out_shape + ) + + if seq_dim != 0: + out = out.movedim(0, seq_dim) + return out.contiguous() + + +@dataclass(frozen=True) +class _SbhdLayoutRedistributionPlan: + """Rank-local SBHD all-to-all plan expressed in sequence-segment counts.""" + + send_slots: tuple[int, ...] + input_segment_counts: tuple[int, ...] + output_segment_counts: tuple[int, ...] + receive_permutation: tuple[int, ...] + + +def _sbhd_segments_per_rank(tp_size: int) -> int: + """Return two SBHD segments for CP-only conversion and one for even-TP SP conversion.""" + if tp_size == 1: + return 2 + if tp_size % 2 != 0: + raise ValueError( + "Sequence-parallel SBHD CP layout conversion requires an even tensor-parallel size, " + f"got {tp_size}" + ) + return 1 + + +def _local_sbhd_segment_ids( + layout: CpPartitionMode, cp_size: int, cp_rank: int, tp_size: int = 1, tp_rank: int = 0 +) -> tuple[int, ...]: + """Return the atomic SBHD sequence segments owned by one TP×CP rank.""" + segments_per_rank = _sbhd_segments_per_rank(tp_size=tp_size) + if layout == "contiguous": + first_segment = segments_per_rank * (cp_rank * tp_size + tp_rank) + return tuple(range(first_segment, first_segment + segments_per_rank)) + if layout == "zigzag": + segments_per_cp_half = tp_size * segments_per_rank // 2 + front_start = cp_rank * segments_per_cp_half + back_start = (2 * cp_size - cp_rank - 1) * segments_per_cp_half + cp_segments = tuple(range(front_start, front_start + segments_per_cp_half)) + tuple( + range(back_start, back_start + segments_per_cp_half) + ) + sp_start = segments_per_rank * tp_rank + return cp_segments[sp_start : sp_start + segments_per_rank] + raise ValueError(f"Unsupported CP layout: {layout}") + + +@lru_cache(maxsize=None) +def _sbhd_segment_owner( + segment_id: int, layout: CpPartitionMode, cp_size: int, tp_size: int +) -> tuple[int, int]: + for cp_rank in range(cp_size): + for tp_rank in range(tp_size): + if segment_id in _local_sbhd_segment_ids( + layout=layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ): + return cp_rank, tp_rank + raise ValueError( + f"SBHD segment {segment_id} is not present in the {layout} layout for " + f"{cp_size=} and {tp_size=}" + ) + + +@lru_cache(maxsize=None) +def _build_sbhd_group_rank_by_logical_rank( + cp_global_ranks: tuple[int, ...], + tp_global_ranks: tuple[int, ...], + tp_cp_global_ranks: tuple[int, ...], + current_global_rank: int, +) -> tuple[int, ...]: + """Map logical ``cp_rank * tp_size + tp_rank`` coordinates to group ranks for SBHD.""" + group_rank_by_global_rank = { + global_rank: group_rank for group_rank, global_rank in enumerate(tp_cp_global_ranks) + } + group_rank_by_logical_rank = [] + for cp_global_rank in cp_global_ranks: + for tp_global_rank in tp_global_ranks: + target_global_rank = cp_global_rank + tp_global_rank - current_global_rank + if target_global_rank not in group_rank_by_global_rank: + raise RuntimeError( + "TP and CP process groups do not form the expected Cartesian product" + ) + group_rank_by_logical_rank.append(group_rank_by_global_rank[target_global_rank]) + return tuple(group_rank_by_logical_rank) + + +def _get_sbhd_group_rank_by_logical_rank( + cp_group: torch.distributed.ProcessGroup, + tp_group: torch.distributed.ProcessGroup, + tp_cp_group: torch.distributed.ProcessGroup, +) -> tuple[int, ...]: + return _build_sbhd_group_rank_by_logical_rank( + cp_global_ranks=tuple(torch.distributed.get_process_group_ranks(cp_group)), + tp_global_ranks=tuple(torch.distributed.get_process_group_ranks(tp_group)), + tp_cp_global_ranks=tuple(torch.distributed.get_process_group_ranks(tp_cp_group)), + current_global_rank=torch.distributed.get_rank(), + ) + + +@lru_cache(maxsize=None) +def _build_sbhd_layout_redistribution_plan( + source_layout: CpPartitionMode, + target_layout: CpPartitionMode, + cp_size: int, + cp_rank: int, + tp_size: int = 1, + tp_rank: int = 0, + group_rank_by_logical_rank: tuple[int, ...] | None = None, +) -> _SbhdLayoutRedistributionPlan: + """Build the SBHD all-to-all-v plan for one rank of a CP layout conversion.""" + if cp_size < 1: + raise ValueError(f"cp_size must be positive, got {cp_size}") + if tp_size < 1: + raise ValueError(f"tp_size must be positive, got {tp_size}") + if not 0 <= cp_rank < cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + if not 0 <= tp_rank < tp_size: + raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}") + + group_size = cp_size * tp_size + if group_rank_by_logical_rank is None: + group_rank_by_logical_rank = tuple(range(group_size)) + if sorted(group_rank_by_logical_rank) != list(range(group_size)): + raise ValueError("group_rank_by_logical_rank must be a permutation of the group ranks") + + source_ids = _local_sbhd_segment_ids( + layout=source_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + target_ids = _local_sbhd_segment_ids( + layout=target_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + + def destination_group_rank(segment_id: int) -> int: + destination_cp_rank, destination_tp_rank = _sbhd_segment_owner( + segment_id=segment_id, layout=target_layout, cp_size=cp_size, tp_size=tp_size + ) + destination_logical_rank = destination_cp_rank * tp_size + destination_tp_rank + return group_rank_by_logical_rank[destination_logical_rank] + + send_entries = sorted( + (destination_group_rank(segment_id=segment_id), slot) + for slot, segment_id in enumerate(source_ids) + ) + send_slots = tuple(slot for _, slot in send_entries) + input_segment_counts = tuple( + sum(destination == rank for destination, _ in send_entries) for rank in range(group_size) + ) + + received_ids = [] + output_segment_counts = [] + source_logical_ranks = sorted( + range(group_size), key=lambda logical_rank: group_rank_by_logical_rank[logical_rank] + ) + for source_logical_rank in source_logical_ranks: + source_cp_rank, source_tp_rank = divmod(source_logical_rank, tp_size) + rank_source_ids = _local_sbhd_segment_ids( + layout=source_layout, + cp_size=cp_size, + cp_rank=source_cp_rank, + tp_size=tp_size, + tp_rank=source_tp_rank, + ) + ids_from_source = [ + segment_id + for segment_id in rank_source_ids + if _sbhd_segment_owner( + segment_id=segment_id, layout=target_layout, cp_size=cp_size, tp_size=tp_size + ) + == (cp_rank, tp_rank) + ] + received_ids.extend(ids_from_source) + output_segment_counts.append(len(ids_from_source)) + + if sorted(received_ids) != sorted(target_ids): + raise RuntimeError( + f"Invalid {source_layout}-to-{target_layout} SBHD redistribution plan for " + f"CP rank {cp_rank}, TP rank {tp_rank}: received {received_ids}, " + f"expected {target_ids}" + ) + receive_permutation = tuple(received_ids.index(segment_id) for segment_id in target_ids) + + return _SbhdLayoutRedistributionPlan( + send_slots=send_slots, + input_segment_counts=input_segment_counts, + output_segment_counts=tuple(output_segment_counts), + receive_permutation=receive_permutation, + ) + + +def _redistribute_sbhd_layout( + input_: torch.Tensor, + cp_group: torch.distributed.ProcessGroup, + source_layout: CpPartitionMode, + target_layout: CpPartitionMode, + sequence_parallel: bool, + tp_group: Optional[torch.distributed.ProcessGroup], + tp_cp_group: Optional[torch.distributed.ProcessGroup], +) -> torch.Tensor: + """Redistribute local SBHD sequence segments with a differentiable all-to-all-v.""" + cp_size = cp_group.size() + if cp_size == 1 or source_layout == target_layout: + return input_ + + cp_rank = cp_group.rank() + tp_size, tp_rank = 1, 0 + communication_group = cp_group + group_rank_by_logical_rank = None + if sequence_parallel and tp_group is not None and tp_group.size() > 1: + if tp_cp_group is None: + raise ValueError( + "tp_cp_group is required for direct sequence-parallel SBHD layout conversion" + ) + tp_size, tp_rank = tp_group.size(), tp_group.rank() + communication_group = tp_cp_group + group_rank_by_logical_rank = _get_sbhd_group_rank_by_logical_rank( + cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group + ) + + plan = _build_sbhd_layout_redistribution_plan( + source_layout=source_layout, + target_layout=target_layout, + cp_size=cp_size, + cp_rank=cp_rank, + tp_size=tp_size, + tp_rank=tp_rank, + group_rank_by_logical_rank=group_rank_by_logical_rank, + ) + + input_contiguous = input_.contiguous() + local_seq_len = input_contiguous.shape[0] + local_segment_count = _sbhd_segments_per_rank(tp_size=tp_size) + if local_seq_len % local_segment_count != 0: + raise ValueError( + "SBHD CP layout conversion requires the sequence length local to each TP×CP rank to " + f"be divisible by {local_segment_count}, got {local_seq_len}" + ) + segment_len = local_seq_len // local_segment_count + segment_shape = (local_segment_count, segment_len, *input_contiguous.shape[1:]) + segments = input_contiguous.reshape(segment_shape) + + if plan.send_slots == tuple(range(local_segment_count)): + send_buffer = input_contiguous + else: + send_buffer = segments.flip(0).reshape(input_contiguous.shape) + input_split_sizes = [count * segment_len for count in plan.input_segment_counts] + output_split_sizes = [count * segment_len for count in plan.output_segment_counts] + received = all_to_all( + group=communication_group, + input_=send_buffer, + output_split_sizes_=output_split_sizes, + input_split_sizes=input_split_sizes, + ) + + received_segments = received.reshape(segment_shape) + if plan.receive_permutation == tuple(range(local_segment_count)): + output = received + else: + output = received_segments.flip(0).reshape(input_contiguous.shape) + return output.contiguous() diff --git a/megatron/core/context_parallel_layout/routes.py b/megatron/core/context_parallel_layout/routes.py new file mode 100644 index 00000000000..66981dfca52 --- /dev/null +++ b/megatron/core/context_parallel_layout/routes.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""THD context-parallel route helpers.""" + +import warnings +from typing import TYPE_CHECKING, List, Optional, Tuple + +import torch + +from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute +from megatron.core.context_parallel_layout.utils import ( + get_packed_seq_params_cp_partition_cu_seqlens, +) +from megatron.core.utils import nvtx_range + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + +_ThdLayoutSegment = Tuple[int, int, int] + + +def _compact_thd_cu_seqlens_to_list(cu_seqlens: torch.Tensor) -> List[int]: + if cu_seqlens.dim() != 1: + raise ValueError(f"cu_seqlens must be 1-D, got shape {tuple(cu_seqlens.shape)}.") + + cu = cu_seqlens.detach().to(device="cpu", dtype=torch.long).tolist() + if not cu or cu[0] != 0: + raise ValueError(f"cu_seqlens must start at 0, got {cu_seqlens}.") + + compact_cu: List[int] = [cu[0]] + prev = cu[0] + for value in cu[1:]: + if value < prev: + raise ValueError(f"cu_seqlens must be nondecreasing, got {cu_seqlens}.") + if value != prev: + compact_cu.append(value) + prev = value + return compact_cu + + +def _validate_thd_route_partitioning(cu: List[int], cp_size: int) -> None: + total_tokens = cu[-1] + if total_tokens % cp_size != 0: + raise ValueError( + f"Contiguous CP partitioning requires total_tokens={total_tokens} " + f"to be divisible by cp_size={cp_size}." + ) + + chunk_divisor = 2 * cp_size + bad_seq_lens = [ + seq_end - seq_start + for seq_start, seq_end in zip(cu[:-1], cu[1:]) + if (seq_end - seq_start) % chunk_divisor != 0 + ] + if bad_seq_lens: + raise ValueError( + "All packed sequence lengths must be divisible by " + f"2 * cp_size ({chunk_divisor}) for zigzag CP layout conversion, " + f"got {bad_seq_lens}." + ) + + +def _build_thd_layout_segments( + cu: List[int], cp_size: int, cp_rank: int, cp_partition_mode: CpPartitionMode +) -> Tuple[List[_ThdLayoutSegment], int]: + total_tokens = cu[-1] + if cp_partition_mode == "contiguous": + part_len = total_tokens // cp_size + if part_len == 0: + return [], 0 + return [(cp_rank * part_len, part_len, 0)], part_len + + if cp_partition_mode != "zigzag": + raise ValueError( + f"Unsupported context-parallel partition mode {cp_partition_mode!r} " + f"for THD layout segments with cp_size={cp_size}, rank={cp_rank}." + ) + + segments: List[_ThdLayoutSegment] = [] + local_start = 0 + for seq_start, seq_end in zip(cu[:-1], cu[1:]): + seq_len = seq_end - seq_start + chunk_len = seq_len // (2 * cp_size) + first_chunk = cp_rank + second_chunk = 2 * cp_size - cp_rank - 1 + segments.append((seq_start + first_chunk * chunk_len, chunk_len, local_start)) + segments.append((seq_start + second_chunk * chunk_len, chunk_len, local_start + chunk_len)) + local_start += 2 * chunk_len + + return segments, local_start + + +def _intersect_thd_layout_segments( + source_segments: List[_ThdLayoutSegment], target_segments: List[_ThdLayoutSegment] +) -> List[Tuple[int, int, int]]: + intersections: List[Tuple[int, int, int]] = [] + source_index = 0 + target_index = 0 + while source_index < len(source_segments) and target_index < len(target_segments): + source_global_start, source_len, source_local_start = source_segments[source_index] + target_global_start, target_len, target_local_start = target_segments[target_index] + source_global_end = source_global_start + source_len + target_global_end = target_global_start + target_len + + overlap_start = max(source_global_start, target_global_start) + overlap_end = min(source_global_end, target_global_end) + if overlap_start < overlap_end: + intersections.append( + ( + source_local_start + overlap_start - source_global_start, + target_local_start + overlap_start - target_global_start, + overlap_end - overlap_start, + ) + ) + + if source_global_end <= target_global_end: + source_index += 1 + else: + target_index += 1 + + return intersections + + +def _build_thd_layout_side_route( + local_segments: List[_ThdLayoutSegment], + target_segments_by_rank: List[List[_ThdLayoutSegment]], + *, + device: torch.device, +) -> Tuple[Optional[torch.Tensor], List[int]]: + row_order: List[int] = [] + split_sizes: List[int] = [] + for peer_rank in range(len(target_segments_by_rank)): + intersections = _intersect_thd_layout_segments( + local_segments, target_segments_by_rank[peer_rank] + ) + intersections.sort(key=lambda item: item[1]) + split_size = 0 + for source_row, _, length in intersections: + row_order.extend(range(source_row, source_row + length)) + split_size += length + split_sizes.append(split_size) + + if all(row == index for index, row in enumerate(row_order)): + return None, split_sizes + return torch.tensor(row_order, device=device, dtype=torch.long), split_sizes + + +def build_thd_cp_partition_route( + cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, *, device: Optional[torch.device] = None +) -> ThdCpRoute: + """Precompute the rank-local THD CP layout route for a microbatch. + + The route stores both zigzag and contiguous layout views and can be reused + for either conversion direction over tensors with the same THD sequence + axis in the same microbatch. + """ + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}.") + if not 0 <= cp_rank < cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}.") + if device is None: + device = cu_seqlens.device + + with nvtx_range("cp_layout/thd/route"): + cu = _compact_thd_cu_seqlens_to_list(cu_seqlens) + _validate_thd_route_partitioning(cu, cp_size) + + zigzag_segments_by_rank: List[List[_ThdLayoutSegment]] = [] + zigzag_lengths: List[int] = [] + contiguous_segments_by_rank: List[List[_ThdLayoutSegment]] = [] + for rank in range(cp_size): + zigzag_segments, zigzag_length = _build_thd_layout_segments(cu, cp_size, rank, "zigzag") + contiguous_segments, contiguous_length = _build_thd_layout_segments( + cu, cp_size, rank, "contiguous" + ) + if zigzag_length != contiguous_length: + raise ValueError( + "THD CP layout conversion must preserve local token count, " + f"got zigzag={zigzag_length}, contiguous={contiguous_length} " + f"for cp_size={cp_size}, rank={rank}." + ) + zigzag_segments_by_rank.append(zigzag_segments) + zigzag_lengths.append(zigzag_length) + contiguous_segments_by_rank.append(contiguous_segments) + + zigzag_index, zigzag_split_sizes = _build_thd_layout_side_route( + zigzag_segments_by_rank[cp_rank], contiguous_segments_by_rank, device=device + ) + contiguous_index, contiguous_split_sizes = _build_thd_layout_side_route( + contiguous_segments_by_rank[cp_rank], zigzag_segments_by_rank, device=device + ) + + local_length = zigzag_lengths[cp_rank] + if sum(zigzag_split_sizes) != local_length: + raise ValueError( + "Zigzag THD CP route split sizes do not match the local token count: " + f"splits={zigzag_split_sizes}, local_length={local_length}." + ) + if sum(contiguous_split_sizes) != local_length: + raise ValueError( + "Contiguous THD CP route split sizes do not match the local token count: " + f"splits={contiguous_split_sizes}, local_length={local_length}." + ) + + return ThdCpRoute( + zigzag_index=zigzag_index, + zigzag_split_sizes=zigzag_split_sizes, + contiguous_index=contiguous_index, + contiguous_split_sizes=contiguous_split_sizes, + ) + + +def get_thd_cp_partition_route( + packed_seq_params: Optional["PackedSeqParams"], + source_partition_mode: CpPartitionMode, + target_partition_mode: CpPartitionMode, +) -> Optional[ThdCpRoute]: + """Return the precomputed THD CP partition route for one direction. + + The fallback below is intentionally only a compatibility path: it performs + a blocking device-to-host copy while compacting ``cu_seqlens`` and mutates + ``packed_seq_params`` by storing the resulting route. Production callers + should prebuild routes when constructing the batch. + """ + if source_partition_mode == target_partition_mode: + return None + if source_partition_mode not in ("zigzag", "contiguous") or target_partition_mode not in ( + "zigzag", + "contiguous", + ): + raise ValueError( + f"Unsupported CP partition mode conversion " + f"{source_partition_mode!r} -> {target_partition_mode!r} for THD route." + ) + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return None + + route = getattr(packed_seq_params, "cp_partition_route", None) + if route is not None: + return route + + warnings.warn( + "THD PackedSeqParams is missing precomputed context-parallel layout routes. " + "This lookup will attempt to build them from packed_seq_params.cp_group as " + "a compatibility fallback. The fallback synchronizes cu_seqlens to CPU " + "and mutates packed_seq_params.cp_partition_route, so it should not be " + "used on the steady-state forward path. Callers should prebuild THD CP " + "routes when constructing the batch; a future release will require the " + "routes to be present before layout conversion.", + FutureWarning, + stacklevel=2, + ) + prebuild_thd_cp_partition_routes(packed_seq_params) + return getattr(packed_seq_params, "cp_partition_route", None) + + +def prebuild_thd_cp_partition_routes( + packed_seq_params: Optional["PackedSeqParams"], + cp_group: Optional[torch.distributed.ProcessGroup] = None, + *, + device: Optional[torch.device] = None, +) -> None: + """Prebuild the THD CP layout route for a packed microbatch.""" + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return + if cp_group is None: + cp_group = getattr(packed_seq_params, "cp_group", None) + if cp_group is None or cp_group.size() <= 1: + return + cp_size = cp_group.size() + cp_rank = cp_group.rank() + cu_seqlens = get_packed_seq_params_cp_partition_cu_seqlens(packed_seq_params) + if cu_seqlens is None: + return + if device is None: + device = cu_seqlens.device + + packed_seq_params.cp_partition_route = build_thd_cp_partition_route( + cu_seqlens, cp_size, cp_rank, device=device + ) diff --git a/megatron/core/context_parallel_layout/types.py b/megatron/core/context_parallel_layout/types.py new file mode 100644 index 00000000000..97a54bef25a --- /dev/null +++ b/megatron/core/context_parallel_layout/types.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Leaf type definitions for context-parallel layout helpers.""" + +from dataclasses import dataclass +from typing import List, Literal, Optional + +import torch + +CpPartitionMode = Literal["zigzag", "contiguous"] + + +@dataclass +class ThdCpRoute: + """Rank-local route plan for THD zigzag/contiguous CP layout conversion. + + The route stores each layout's local communication view exactly once. A + directional conversion interprets the source layout fields as send metadata + and the target layout fields as receive metadata. + """ + + zigzag_index: Optional[torch.Tensor] + zigzag_split_sizes: List[int] + contiguous_index: Optional[torch.Tensor] + contiguous_split_sizes: List[int] diff --git a/megatron/core/context_parallel_layout/utils.py b/megatron/core/context_parallel_layout/utils.py new file mode 100644 index 00000000000..66d9b9183ae --- /dev/null +++ b/megatron/core/context_parallel_layout/utils.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Packed-sequence metadata helpers for CP partition-mode tracking.""" + +from typing import TYPE_CHECKING, Optional + +import torch + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + + +def get_packed_seq_params_cp_partition_cu_seqlens( + packed_seq_params: Optional["PackedSeqParams"], +) -> Optional[torch.Tensor]: + """Return THD cumulative sequence lengths used for CP layout conversion. + + ``packed_seq_params=None`` represents the ordinary SBHD path. Only THD + metadata carries global packed-token boundaries. + """ + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + return None + return ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + + +def finalize_packed_seq_params( + packed_seq_params: Optional["PackedSeqParams"], +) -> Optional["PackedSeqParams"]: + """Resolve CP metadata and prebuild the THD layout route for a microbatch.""" + if packed_seq_params is None: + return None + + # Keep these imports local: routes depends on this module for metadata access. + from megatron.core.context_parallel_layout.routes import prebuild_thd_cp_partition_routes + from megatron.core.packed_seq_params import resolve_cp_group + from megatron.core.parallel_state import get_context_parallel_group + + cp_group = resolve_cp_group(get_context_parallel_group(), packed_seq_params) + packed_seq_params.cp_group = cp_group + prebuild_thd_cp_partition_routes(packed_seq_params, cp_group) + return packed_seq_params diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 120513c1819..acb7d796e8a 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import lru_cache from math import ceil, log2 @@ -82,7 +82,10 @@ def get_cp_slice_for_thd( if cp_partition_mode != "zigzag": raise ValueError(f"Unsupported CP partition mode: {cp_partition_mode}") - index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + cu_seqlens_for_index = ( + cu_seqlens if cu_seqlens.dtype == torch.int32 else cu_seqlens.to(dtype=torch.int32) + ) + index = get_thd_partitioned_indices(cu_seqlens_for_index, total_tokens, cp_size, cp_rank) for key in keys: if key in batch and batch[key] is not None: batch[key] = batch[key].index_select(0, index) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8d797e816db..348847e7399 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import copy @@ -1753,11 +1753,13 @@ def __init__( # These fields are MCore-only and should not be forwarded to TE attention. # total_tokens and seq_idx are only for Mamba; tokens_per_sample is only for - # MoE sequence-level aux loss reshaping; cp_partition_mode is MCore CP metadata. + # MoE sequence-level aux loss reshaping; cp_partition_mode and cp_partition_route + # are MCore CP metadata. self.kept_packed_seq_params.discard("total_tokens") self.kept_packed_seq_params.discard("seq_idx") self.kept_packed_seq_params.discard("tokens_per_sample") self.kept_packed_seq_params.discard("cp_partition_mode") + self.kept_packed_seq_params.discard("cp_partition_route") if config.qk_clip or config.log_max_attention_logit: # qk-clip is only supported in TE 2.9.0 and later diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index f88b8bdb5a1..e0f7c9ff8c1 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -373,6 +373,9 @@ def _preprocess( rotary_pos_sin = None # this is used to store combined cos/sin embeddings, exclusively for flash infer rope rotary_pos_cos_sin = None + # Model-level rotary_pos_emb is only for regular attention. Regular + # attention uses the default zigzag CP RoPE layout; MLA/CSA/DSv4-style + # variants must ignore this external RoPE and build/apply RoPE internally. if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention: use_flash_infer_fused_rope = ( @@ -730,6 +733,7 @@ def _postprocess( output_weight = None if self.share_embeddings_and_output_weights: output_weight = self.shared_embedding_or_output_weight() + if mtp_in_postprocess and not (in_inference_mode or is_spec_decode): hidden_states = self.mtp( input_ids=input_ids, diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 47da3af197f..8e71fcb0106 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2024, Tri Dao, Albert Gu. # Some of this code was adopted from https://github.com/state-spaces/mamba/ @@ -429,6 +429,7 @@ def _call_inner_transformer_layer_without_local_bda( inference_context=inference_context, padding_mask=padding_mask, input_ids=input_ids, + packed_seq_params=packed_seq_params, ) if layer.mlp_norm_manager is not None: output_with_bias = layer._group_offload_output_with_bias( diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 0d3ebcfe9f4..d10639e5a12 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -478,6 +478,9 @@ def forward( decoder_input = None rotary_pos_emb = None + # Model-level rotary_pos_emb is only for regular attention. Regular + # attention uses the default zigzag CP RoPE layout; MLA/CSA/DSv4-style + # variants must ignore this external RoPE and build/apply RoPE internally. if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention: rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( inference_context, self.decoder, decoder_input, self.config, packed_seq_params @@ -485,6 +488,7 @@ def forward( rotary_pos_emb = self.rotary_pos_emb( rotary_seq_len, packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd', + cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None, ) elif self.position_embedding_type == 'yarn': rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( @@ -494,6 +498,7 @@ def forward( rotary_pos_emb, _ = self.rotary_pos_emb( rotary_seq_len, packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd', + cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None, ) # Wrap decoder_input to allow the decoder (HybridStack) to delete the diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index a0957a39eab..47af6be146f 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -1,18 +1,25 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from dataclasses import dataclass -from typing import Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union import torch import torch.distributed as dist import torch.nn.functional as F from torch import Tensor +if TYPE_CHECKING: + from megatron.core.context_parallel_layout import ThdCpRoute + @dataclass class PackedSeqParams: ''' parameters to TEDotProductAttention and fused rope kernels for the `thd` (packed) sequence format + + ``cp_partition_route`` is a per-microbatch THD CP layout conversion route. + Metadata annotation helpers update the current partition mode in-place while + preserving the route identity. ''' qkv_format: str = None @@ -29,6 +36,7 @@ class PackedSeqParams: pad_between_seqs: Optional[bool] = None cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" tokens_per_sample: int = None + cp_partition_route: Optional["ThdCpRoute"] = None def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index 79c58a41710..515d9bd132e 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -12,10 +12,7 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -from megatron.core.context_parallel_layout import ( - contiguous_to_zigzag_chunks, - zigzag_to_contiguous_chunks, -) +from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group @@ -113,7 +110,8 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) - base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp cp_group = resolve_cp_group(base_cp_group, packed_seq_params) if self.config.linear_cp_mode == "chunkwise": cp_group_chunkwise = cp_group @@ -132,6 +130,18 @@ def forward( cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 cp_size_runtime = cp_group.size() + back_to_input_converter = None + if self.config.linear_cp_mode == "chunkwise": + hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + target_partition_mode="contiguous", + sequence_parallel=self.config.sequence_parallel, + config=self.config, + ) seq_len_local, batch, _ = hidden_states.shape seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise @@ -145,6 +155,22 @@ def forward( # TODO: support inference raise NotImplementedError("GDN does not support inference for now.") + if cp_size_headwise > 1 and ( + ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and packed_seq_params.cp_partition_mode != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "GatedDeltaNet with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling GatedDeltaNet." + ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': assert batch == 1, "Packed sequence expects batch dimension to be 1" assert ( @@ -238,6 +264,11 @@ def _checkpointed_compute(hidden_states): chunkwise_cp_context, ) + if back_to_input_converter is not None: + out = back_to_input_converter.convert( + out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) + return out, out_bias def _forward_compute( @@ -259,18 +290,6 @@ def _forward_compute( qkvzba, _ = self.in_proj(hidden_states) nvtx_range_pop(suffix="in_proj") - # TODO: Move CP layout ownership to a model/region-level scheduler so hybrid models can - # enter contiguous layout before GDN regions instead of paying module-local conversions. - if cp_size_chunkwise > 1: - nvtx_range_push(suffix="zigzag_to_contiguous") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - qkvzba = zigzag_to_contiguous_chunks( - qkvzba, cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q - ) - else: - qkvzba = zigzag_to_contiguous_chunks(qkvzba, cp_group_chunkwise, seq_dim=0) - nvtx_range_pop(suffix="zigzag_to_contiguous") - qkvzba, thd_cp_a2a_inv = a2a_cp_to_hp( qkvzba, self.in_proj_split_sections, @@ -414,20 +433,6 @@ def _gated_norm_and_layout_restore( norm_out = norm_out.reshape(batch, seq_len, -1) norm_out = norm_out.transpose(0, 1).contiguous() - # TODO: The planned CP layout refactor should keep consecutive GDN layers contiguous and - # restore zigzag only at SDPA/canonical-layout boundaries. - if cp_size_chunkwise > 1: - nvtx_range_push(suffix="contiguous_to_zigzag") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - norm_out = contiguous_to_zigzag_chunks( - norm_out, cp_group=cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q - ) - else: - norm_out = contiguous_to_zigzag_chunks( - norm_out, cp_group=cp_group_chunkwise, seq_dim=0 - ) - nvtx_range_pop(suffix="contiguous_to_zigzag") - return a2a_hp_to_cp( norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 1f29c93eef3..f08dc7d85bb 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -11,6 +11,7 @@ from torch import Tensor from megatron.core import tensor_parallel +from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ( ReplicaId, @@ -329,7 +330,9 @@ def __init__( self.kv_projection_size = self.config.kv_channels * self.config.num_query_groups if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp', 'tp_cp'] + ) else: assert hasattr( pg_collection, 'tp' @@ -530,7 +533,7 @@ def _checkpointed_attention_forward( checkpoint_inputs.append(kwarg_value) def custom_forward(*inputs): - (query, key, value, attention_mask, _, attn_mask_type, *tensor_kwarg_values) = inputs + query, key, value, attention_mask, _, attn_mask_type, *tensor_kwarg_values = inputs attn_mask_type = AttnMaskType(attn_mask_type.item()) extra_kwargs = dict(core_attention_extra_kwargs) for name, kwarg_value in zip(tensor_kwarg_names, tensor_kwarg_values): @@ -1361,6 +1364,19 @@ def forward( if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" self.pg_collection.cp = packed_seq_params.cp_group + hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + key_value_states=key_value_states, + packed_seq_params=packed_seq_params, + cp_group=self.pg_collection.cp, + tp_group=self.pg_collection.tp, + tp_cp_group=getattr(self.pg_collection, "tp_cp", None), + target_partition_mode="zigzag", + sequence_parallel=self.config.sequence_parallel, + config=self.config, + attention_mask=attention_mask, + attention_bias=attention_bias, + ) # Check if we need to skip RoPE # no_rope is 0-indexed array and self.layer_number is 1-indexed @@ -1507,6 +1523,10 @@ def forward( out = output.transpose(0, 1).contiguous() context_layer = out.view(out.size(0), out.size(1), -1) output, bias = apply_module(self.linear_proj)(context_layer) + if back_to_input_converter is not None: + output = back_to_input_converter.convert( + output, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) self.pg_collection.cp = _orig_cp_group return output, bias @@ -1700,6 +1720,11 @@ def forward( output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) nvtx_range_pop(suffix="linear_proj") + if back_to_input_converter is not None: + output = back_to_input_converter.convert( + output, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) + self.pg_collection.cp = _orig_cp_group return output, bias @@ -2003,9 +2028,9 @@ def get_query_key_value_tensors( ] if SplitAlongDim is not None: - (query, gate, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + query, gate, key, value = SplitAlongDim(mixed_qkv, 3, split_arg_list) else: - (query, gate, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) + query, gate, key, value = torch.split(mixed_qkv, split_arg_list, dim=3) else: # If no output gate: [sq, b, ng, (np/ng + 2) * hn] # --> [sq, b, ng, np/ng * hn], None, [sq, b, ng, hn], [sq, b, ng, hn] @@ -2020,9 +2045,9 @@ def get_query_key_value_tensors( return mixed_qkv, split_arg_list if SplitAlongDim is not None: - (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + query, key, value = SplitAlongDim(mixed_qkv, 3, split_arg_list) else: - (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) + query, key, value = torch.split(mixed_qkv, split_arg_list, dim=3) # Query [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head) @@ -2354,7 +2379,7 @@ def get_query_key_value_tensors( mixed_kv = mixed_kv.view(*new_tensor_shape) # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn] - (key, value) = tensor_parallel.split_tensor_along_last_dim(mixed_kv, 2) + key, value = tensor_parallel.split_tensor_along_last_dim(mixed_kv, 2) # Attention head [sq, b, h] --> [sq, b, hp] query, _ = apply_module(self.linear_q)(hidden_states) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index ecfd2d52ee5..6c0f38eaef5 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -64,11 +64,8 @@ def _restore_packed_thd_batch_dim( core_attn_out: torch.Tensor, hidden_states: torch.Tensor, packed_seq_params ) -> torch.Tensor: """Restore the singleton packed-THD batch dim only when core attention omitted it.""" - if ( - packed_seq_params is not None - and packed_seq_params.qkv_format == 'thd' - and core_attn_out.ndim == hidden_states.ndim - 1 - ): + thd_packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if thd_packed_seq and core_attn_out.ndim == hidden_states.ndim - 1: core_attn_out = core_attn_out.unsqueeze(1) return core_attn_out @@ -464,13 +461,13 @@ def get_query_key_value_tensors( mscale = 1.0 rotary_pos_cos = None rotary_pos_sin = None - packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + thd_packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' if self.config.rope_type == "rope": - rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=thd_packed_seq) else: if self.config.apply_rope_fusion: rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( - rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=thd_packed_seq ) rotary_pos_emb = None assert inference_context is None, "Inference with MLA RoPE fusion is not supported" @@ -479,9 +476,11 @@ def get_query_key_value_tensors( and fused_apply_mla_rope_for_kv is not None ), "Fused MLA RoPE apply is not imported successfully" else: - rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + rotary_pos_emb, mscale = self.rotary_pos_emb( + rotary_seq_len, packed_seq=thd_packed_seq + ) - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if thd_packed_seq: if packed_seq_params.cu_seqlens_q_padded is not None: cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded else: @@ -545,7 +544,7 @@ def get_query_key_value_tensors( # k_pos_emb: [s, b, qk_pos_emb_head_dim] k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb, group=self.tp_group) - if packed_seq_params is not None: + if thd_packed_seq: assert q_compressed.ndim == 3 and q_compressed.size(1) == 1 assert kv_compressed.ndim == 3 and kv_compressed.size(1) == 1 assert k_pos_emb.ndim == 3 and k_pos_emb.size(1) == 1 @@ -648,7 +647,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po sequence_start = inference_context.sequence_len_offset sequence_end = sequence_start + q_len rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] - elif packed_seq_params is None or self.config.context_parallel_size == 1: + elif not thd_packed_seq or self.config.context_parallel_size == 1: # Shorten rotary_pos_emb to the sequence length when inference_params # is not provided. This makes sure we can run forward directly with # any sequence length. During training, the sequence length is always @@ -914,13 +913,24 @@ def forward( inference_context is None and inference_params is None ), "Inference is not supported for AbsorbedMLA" - # Set the right cp group for dynamic-cp. Mirrors Attention.forward: - # downstream RoPE uses self.pg_collection.cp, which must point at this - # microbatch's dynamic CP group. Restored before every return. + # Set the right cp group for dynamic-cp. Downstream RoPE and CSA core + # attention use self.pg_collection.cp, which must point at this + # microbatch's dynamic CP group. Restored before returning. _orig_cp_group = self.pg_collection.cp if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" self.pg_collection.cp = packed_seq_params.cp_group + thd_packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == "thd" + if ( + thd_packed_seq + and self.pg_collection.cp is not None + and get_pg_size(self.pg_collection.cp) > 1 + and packed_seq_params.cp_partition_mode != "zigzag" + ): + raise ValueError( + "AbsorbedMLASelfAttention requires cp_partition_mode='zigzag'. " + "CP partition conversion must be handled before entering AbsorbedMLA." + ) # ===================== # Query, Key, and Value diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 9e7a41e8246..ae24626c380 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -2089,6 +2089,21 @@ def forward( """ nvtx_range_push("compressed_sparse_attn") + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + + cp_size = self.pg_collection.cp.size() if self.pg_collection.cp is not None else 1 + qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None + if cp_size > 1 and qkv_format != 'thd': + raise ValueError("CompressedSparseAttention with CP requires qkv_format='thd'.") + if cp_size > 1 and packed_seq_params.cp_partition_mode != "contiguous": + raise ValueError( + "CompressedSparseAttention requires cp_partition_mode='contiguous'. " + "CP partition conversion must be handled before entering CSA." + ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': if self.pg_collection.cp is not None and self.pg_collection.cp.size() > 1: output = self._forward_thd_cp( @@ -2096,6 +2111,7 @@ def forward( ) else: output = self._forward_thd(query, key, x, qr, packed_seq_params) + self.pg_collection.cp = _orig_cp_group nvtx_range_pop("compressed_sparse_attn") return output @@ -2140,6 +2156,7 @@ def forward( if indexer_loss is not None: output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + self.pg_collection.cp = _orig_cp_group nvtx_range_pop("compressed_sparse_attn") return output @@ -2199,7 +2216,7 @@ def _forward_unfused_csa_thd( # Physical padded offsets define the packed address space; # unpadded lengths identify the real rows within each segment. - (varlen_starts, varlen_ends, query_valid_rows, compressed_offsets) = ( + varlen_starts, varlen_ends, query_valid_rows, compressed_offsets = ( _build_compressed_thd_indexer_metadata( cu_seqlens_q, cu_seqlens_compressed_idx, diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 439a0f0649e..4304662db21 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -349,6 +349,18 @@ def forward( if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" self.pg_collection.cp = packed_seq_params.cp_group + if ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and self.pg_collection.cp is not None + and get_pg_size(self.pg_collection.cp) > 1 + and packed_seq_params.cp_partition_mode != "zigzag" + ): + raise ValueError( + "MultiLatentAttention requires cp_partition_mode='zigzag', but " + f"packed_seq_params has {packed_seq_params.cp_partition_mode!r}. CP partition " + "conversion must be handled before entering MLA." + ) # ===================== # Query, Key, and Value @@ -760,7 +772,7 @@ def get_query_key_value_tensors( # k_pos_emb: [s, b, qk_pos_emb_head_dim] k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb, group=self.tp_group) - if packed_seq_params is not None: + if thd_packed_seq: # If sequence packing, TE expect [t, h, d] shaped qkv input. # In Megatron-Core, the qkv shape is [t, 1, h, d]. # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. @@ -915,7 +927,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po sequence_start = inference_context.sequence_len_offset sequence_end = sequence_start + q_len rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] - elif packed_seq_params is None or self.config.context_parallel_size == 1: + elif not thd_packed_seq or self.config.context_parallel_size == 1: # Shorten rotary_pos_emb to the sequence length when inference_params # is not provided. This makes sure we can run forward directly with # any sequence length. During training, the sequence length is always diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 220e9662d8e..d6afac69c48 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1569,23 +1569,52 @@ def __post_init__(self): if self.sequence_packing_scheduler is None: raise ValueError( "cp_partition_mode='contiguous' with context parallelism requires THD " - "inputs from a sequence_packing_scheduler; BSHD inputs are not supported." + "inputs from a sequence_packing_scheduler; BSHD inputs are not supported. " + "The legacy non-scheduler CP batch slicing path only supports zigzag layout." ) - - if self.context_parallel_size > 1: - if ( - self.experimental_attention_variant == "dsv4_hybrid" - and self.cp_partition_mode != "contiguous" - ): - raise ValueError("DSv4 Hybrid with CP requires cp_partition_mode='contiguous'.") if ( - self.experimental_attention_variant != "dsv4_hybrid" - and self.cp_partition_mode != "zigzag" + (self.mtp_num_layers or 0) > 0 + and self.tensor_model_parallel_size > 1 + and self.sequence_parallel ): raise ValueError( - "cp_partition_mode='contiguous' currently is only supported with dsv4_hybrid." + "MTP with tensor_model_parallel_size > 1, sequence_parallel=True, and " + "cp_partition_mode='contiguous' has a known token-side padding-mask layout " + "bug. This combination is temporarily unsupported and will be fixed in a " + "follow-up change." ) + if self.context_parallel_size > 1: + if self.cp_partition_mode == "contiguous": + if ( + self.multi_latent_attention + and self.experimental_attention_variant != "dsv4_hybrid" + ): + raise ValueError( + "cp_partition_mode='contiguous' is not supported with " + "multi_latent_attention outside dsv4_hybrid." + ) + if self.experimental_attention_variant not in ("dsv4_hybrid", "gated_delta_net"): + raise ValueError( + "cp_partition_mode='contiguous' with context parallelism currently " + "requires experimental_attention_variant to be either 'dsv4_hybrid' " + "or 'gated_delta_net'." + ) + if ( + self.experimental_attention_variant == "gated_delta_net" + and self.linear_cp_mode == "headwise" + ): + raise ValueError( + "cp_partition_mode='contiguous' is incompatible with " + "gated_delta_net linear_cp_mode='headwise'." + ) + elif self.cp_partition_mode == "zigzag": + if self.experimental_attention_variant == "dsv4_hybrid": + raise ValueError( + "DSv4 Hybrid with context parallelism requires " + "cp_partition_mode='contiguous'." + ) + # Normalize the deprecated DSv4 kernel switch only after all deprecated attention # selectors have been folded into experimental_attention_variant, and immediately # before the centralized attention-variant validation consumes dsa_kernel_backend. @@ -3138,6 +3167,30 @@ def _scope_to_str(s): "path is unvalidated." ) + cuda_graph_captures_attention = self.cuda_graph_impl == "full_iteration" or ( + self.cuda_graph_impl in ("local", "transformer_engine") + and (not self.cuda_graph_modules or CudaGraphModule.attn in self.cuda_graph_modules) + ) + + cp_layout_conversion_required = self.experimental_attention_variant == "gated_delta_net" + # TODO: Extend this predicate as GDN2/KDA are introduced, and for DSv4 when + # dsa_cp_balance_indexer is introduced; those paths will also require module-local THD CP + # layout conversion. + if ( + (self.context_parallel_size > 1 or self.dynamic_context_parallel) + and self.sequence_packing_scheduler is not None + and cuda_graph_captures_attention + and cp_layout_conversion_required + ): + raise ValueError( + "THD context parallel layout conversion is required for this model " + "configuration, but it is not supported by CUDA graph capture that includes " + "attention " + f"(experimental_attention_variant={self.experimental_attention_variant!r}, " + f"cuda_graph_impl={self.cuda_graph_impl!r}, " + f"cuda_graph_modules={self.cuda_graph_modules!r})." + ) + if self.cuda_graph_impl != "none": if self.cpu_offloading and self.cuda_graph_impl != "full_iteration": diff --git a/megatron/core/utils.py b/megatron/core/utils.py index ac97916d21b..dbc87c1f4be 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2737,6 +2737,16 @@ def nvtx_range_pop(msg=None, suffix=None) -> None: torch.cuda.nvtx.range_pop() +@contextmanager +def nvtx_range(msg=None, suffix=None): + """Create an NVTX range controlled by ``configure_nvtx_profiling``.""" + nvtx_range_push(msg, suffix) + try: + yield + finally: + nvtx_range_pop(msg, suffix) + + @lru_cache(maxsize=None) def _nvtx_decorator_get_func_path(func): """Get the path of a function. diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 94229390984..e7a8ea68ce3 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """General utilities.""" import json @@ -805,7 +805,7 @@ def get_nvtx_range(): time: If True, also track with Megatron timers (default: False) log_level: Timer log level (0=always, 1=default, 2=verbose). Default: 1 """ - from megatron.core.utils import nvtx_range_pop, nvtx_range_push + from megatron.core.utils import nvtx_range as core_nvtx_range @contextmanager def nvtx_range(msg, time=False, log_level=1): @@ -813,10 +813,9 @@ def nvtx_range(msg, time=False, log_level=1): timers = get_timers() timers(msg, log_level=log_level).start() try: - nvtx_range_push(msg) - yield + with core_nvtx_range(msg): + yield finally: - nvtx_range_pop(msg) if time: timers(msg, log_level=log_level).stop() diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 9e99200b27a..c883828e809 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain and SFT GPT.""" @@ -25,6 +25,7 @@ from gpt_builders import gpt_builder from megatron.core import mpu +from megatron.core.context_parallel_layout import finalize_packed_seq_params from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.data_schedule import get_batch_on_this_rank_for_sequence_packing from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset @@ -132,7 +133,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): if args.sequence_packing_scheduler is not None: # `get_batch_on_this_rank_for_sequence_packing` owns scheduler THD metadata # and returns a 7-tuple including `padding_mask`. - return get_batch_on_this_rank_for_sequence_packing( + batch = get_batch_on_this_rank_for_sequence_packing( data_iterator, vpp_size=config.virtual_pipeline_model_parallel_size, mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage), @@ -140,6 +141,8 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): dynamic_cp=args.dynamic_context_parallel, config=config, ) + finalize_packed_seq_params(batch[5]) + return batch # TODO: this is pretty hacky, find a better way is_packed_sequence = args.sft or (args.use_varlen_dataset and not args.varlen_sbhd_validation) @@ -174,19 +177,21 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): # For middle pipeline stages with packed sequences, only cu_seqlens and # max_seqlen are needed (for attention masking); skip the full batch. if not is_first_or_last_pipeline_stage(vp_stage) and is_packed_sequence: + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=int(max_seqlen[0].item()), + max_seqlen_kv=int(max_seqlen[0].item()), + qkv_format='thd', + ) + finalize_packed_seq_params(packed_seq_params) return ( None, None, None, None, None, - PackedSeqParams( - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=int(max_seqlen[0].item()), - max_seqlen_kv=int(max_seqlen[0].item()), - qkv_format='thd', - ), + packed_seq_params, None, ) @@ -239,6 +244,8 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): if 'position_ids' in batch: batch['position_ids'] = position_ids + finalize_packed_seq_params(packed_seq_params) + # Unpack explicitly to avoid relying on dict insertion order. return ( batch.get('tokens'), diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 5fdc7d8645a..daacd3f2bf3 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain and SFT Hybrid.""" # Capture the true program start time BEFORE any heavy imports. @@ -24,6 +24,7 @@ from hybrid_builders import hybrid_builder from megatron.core import mpu +from megatron.core.context_parallel_layout import finalize_packed_seq_params from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.data_schedule import get_batch_on_this_rank_for_sequence_packing from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset @@ -125,6 +126,7 @@ def get_batch(data_iterator, vp_stage=None): dynamic_cp=is_dynamic_cp, config=config, ) + finalize_packed_seq_params(packed_seq_params) return ( attention_mask, None, @@ -328,6 +330,7 @@ def forward_step(data_iterator, model: HybridModel): total_tokens=int(cu_seqlens_for_params[-1].item()), tokens_per_sample=args.seq_length, ) + finalize_packed_seq_params(packed_seq_params) timers('batch-generator').stop() diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py index 93d91ff05ec..943c0bc2417 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -66,6 +66,18 @@ def _make_gdn_config(**overrides): return TransformerConfig(**config_kwargs) +def _set_gdn_test_cp_partition_mode(packed_seq_params, cp_size, linear_cp_mode): + if cp_size <= 1: + return packed_seq_params + if linear_cp_mode == "headwise": + packed_seq_params.cp_partition_mode = "zigzag" + elif linear_cp_mode == "chunkwise": + packed_seq_params.cp_partition_mode = "contiguous" + else: + raise ValueError(f"Invalid linear CP mode: {linear_cp_mode}") + return packed_seq_params + + def test_gdn_pre_gated_delta_rule_fusion_defaults_to_disabled(): config = _make_gdn_config() assert not config.gdn_pre_gated_delta_rule_fusion @@ -203,7 +215,8 @@ def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): # Get TP and CP process groups from device mesh tp_group = parallel_state.get_tensor_model_parallel_group() cp_group = parallel_state.get_context_parallel_group() - pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + tp_cp_group = parallel_state.get_tensor_and_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group, tp_cp=tp_cp_group) # Initialize model, with the same config as Qwen Next except `num_layers` self.transformer_config = TransformerConfig( @@ -667,6 +680,7 @@ def test_gpu_forward_thd_correctness(self): hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) attention_mask_thd = None packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + _set_gdn_test_cp_partition_mode(packed_seq_params, self.cp_size, self.linear_cp_mode) # THD format output_thd, _ = self.gdn( @@ -718,6 +732,7 @@ def test_gpu_forward_thd_padding_correctness(self): padded_params = make_test_packed_seq_params_with_padding( cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] ) + _set_gdn_test_cp_partition_mode(padded_params, self.cp_size, self.linear_cp_mode) output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) output_thd2bshd = output_thd_padded.view(*output_bshd.shape) torch.testing.assert_close( @@ -730,6 +745,7 @@ def test_gpu_forward_thd_padding_correctness(self): # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + _set_gdn_test_cp_partition_mode(no_padding_params, self.cp_size, self.linear_cp_mode) output_thd_no_padding, _ = self.gdn( hidden_states_thd, None, packed_seq_params=no_padding_params ) @@ -755,11 +771,13 @@ def test_gpu_forward_thd_padding_correctness(self): padded_mismatch_params = make_test_packed_seq_params_with_padding( cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] ) + _set_gdn_test_cp_partition_mode(padded_mismatch_params, self.cp_size, self.linear_cp_mode) with pytest.raises(ValueError, match="does not match"): self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) # E) actual mismatch branch without *_padded: should raise. actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + _set_gdn_test_cp_partition_mode(actual_mismatch_params, self.cp_size, self.linear_cp_mode) with pytest.raises(ValueError, match="does not match"): self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py index 6b12ec6a1c8..0f6b4f539d9 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py @@ -167,6 +167,15 @@ def _assert_pre_gated_delta_rule_outputs_close( msg=lambda msg, output_name=name: f"{output_name} mismatch: {msg}", ) + def _make_pre_gated_delta_rule_grad_outputs(self, outputs): + grad_outputs = [] + for output_idx, output in enumerate(outputs): + grad = torch.linspace( + -0.1, 0.1, output.numel(), device=output.device, dtype=torch.float32 + ).reshape(output.shape) + grad_outputs.append(grad + (output_idx - 2.5) * 0.01) + return grad_outputs + def test_fused_and_unfused_forward_match(self): hidden_states = torch.randn( (32, 2, self.unfused_gdn.config.hidden_size), @@ -360,6 +369,7 @@ def test_fused_and_unfused_pre_gated_delta_rule_backward_match(self): batch = 2 seq_len = 32 + torch.manual_seed(1234) qkvzba = torch.randn( (seq_len, batch, reference_gdn.in_proj_dim), device=torch.cuda.current_device(), @@ -375,7 +385,7 @@ def test_fused_and_unfused_pre_gated_delta_rule_backward_match(self): qkvzba_unfused, batch, seq_len, reference_gdn.cp_size, reference_gdn.pg_collection.cp ) fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule(qkvzba_fused) - grad_outputs = [torch.randn_like(output.float()) for output in unfused_outputs] + grad_outputs = self._make_pre_gated_delta_rule_grad_outputs(unfused_outputs) unfused_loss = sum( (output.float() * grad).sum() for output, grad in zip(unfused_outputs, grad_outputs) @@ -497,6 +507,7 @@ def test_fused_and_unfused_packed_pre_gated_delta_rule_backward_match(self): [0, 1, 4, 6, 11], device=torch.cuda.current_device(), dtype=torch.int32 ) seq_len = cu_seqlens[-1].item() + torch.manual_seed(1234) qkvzba = torch.randn( (seq_len, batch, reference_gdn.in_proj_dim), device=torch.cuda.current_device(), @@ -514,7 +525,7 @@ def test_fused_and_unfused_packed_pre_gated_delta_rule_backward_match(self): fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( qkvzba_fused, cu_seqlens_q=cu_seqlens ) - grad_outputs = [torch.randn_like(output.float()) for output in unfused_outputs] + grad_outputs = self._make_pre_gated_delta_rule_grad_outputs(unfused_outputs) unfused_loss = sum( (output.float() * grad).sum() for output, grad in zip(unfused_outputs, grad_outputs) @@ -581,7 +592,7 @@ def test_fused_and_unfused_packed_pre_gated_delta_rule_backward_repeat4_match(se fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( qkvzba_fused, cu_seqlens_q=cu_seqlens ) - grad_outputs = [torch.randn_like(output.float()) for output in unfused_outputs] + grad_outputs = self._make_pre_gated_delta_rule_grad_outputs(unfused_outputs) unfused_loss = sum( (output.float() * grad).sum() for output, grad in zip(unfused_outputs, grad_outputs) @@ -820,6 +831,7 @@ def _make_packed_seq_params(cu_seqlens): cu_seqlens[i + 1] - cu_seqlens[i] for i in range(len(cu_seqlens) - 1) ), total_tokens=cu_seqlens[-1] // cp_size, + cp_partition_mode="contiguous", ) @staticmethod diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py index 781326e1fdd..c8c010d75e7 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py @@ -112,6 +112,8 @@ def test_parallel_gated_delta_net_correctness( sequence_length=256, micro_batch_size=micro_batch_size, sequence_packing=sequence_packing, + cp_partition_mode="contiguous" if is_chunkwise_cp else "zigzag", + compare_param_grads=is_chunkwise_cp and tp == 1 and not sequence_packing, ) diff --git a/tests/unit_tests/test_context_parallel_layout.py b/tests/unit_tests/test_context_parallel_layout.py index b762e4594a1..c2e1edc27ca 100644 --- a/tests/unit_tests/test_context_parallel_layout.py +++ b/tests/unit_tests/test_context_parallel_layout.py @@ -1,68 +1,647 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace import pytest import torch -from megatron.core.context_parallel_layout import get_thd_context_parallel_rank_indices +import megatron.core.context_parallel_layout.conversion as context_parallel_layout_conversion +from megatron.core import parallel_state +from megatron.core.context_parallel_layout import ( + CpPartitionModeConverter, + ThdCpRoute, + convert_module_input_tensors_cp_partition_mode, + prebuild_thd_cp_partition_routes, +) +from megatron.core.context_parallel_layout.routes import ( + build_thd_cp_partition_route, + get_thd_cp_partition_route, +) +from tests.unit_tests.test_utilities import Utils -def _token_ranges(*spans): - return [token for start, end in spans for token in range(start, end)] +class _FakeGroup: + def __init__(self, size, rank): + self._size = size + self._rank = rank -def test_thd_context_parallel_rank_indices_match_per_sequence_chunk_order(): - cu_seqlens = torch.tensor([0, 16, 40]) + def size(self): + return self._size - assert get_thd_context_parallel_rank_indices( - cu_seqlens, 2, 0, "zigzag" - ).tolist() == _token_ranges((0, 4), (12, 16), (16, 22), (34, 40)) - assert get_thd_context_parallel_rank_indices( - cu_seqlens, 2, 1, "zigzag" - ).tolist() == _token_ranges((4, 12), (22, 34)) - assert get_thd_context_parallel_rank_indices(cu_seqlens, 2, 0, "contiguous").tolist() == list( - range(0, 20) - ) - assert get_thd_context_parallel_rank_indices(cu_seqlens, 2, 1, "contiguous").tolist() == list( - range(20, 40) + def rank(self): + return self._rank + + +def _make_sequence_tensor(total_seq_len, seq_dim, device): + if seq_dim == 0: + shape = (total_seq_len, 3, 5) + elif seq_dim == 1: + shape = (3, total_seq_len, 5) + else: + raise ValueError(f"Unsupported test seq_dim {seq_dim}.") + return torch.arange(torch.prod(torch.tensor(shape)), device=device, dtype=torch.float32).view( + *shape ) -@pytest.mark.parametrize("layout", ["zigzag", "contiguous"]) -def test_thd_context_parallel_rank_indices_cover_all_tokens_once(layout): - cu_seqlens = torch.tensor([0, 32, 96, 128]) - cp_size = 4 +def _get_sequence_parallel_shard(tensor, seq_dim, tp_group): + tp_size = tp_group.size() + tp_rank = tp_group.rank() + assert tensor.size(seq_dim) % tp_size == 0 + return tensor.chunk(tp_size, dim=seq_dim)[tp_rank].contiguous() + + +def _get_sbhd_tensor_on_this_cp_rank(tensor, seq_dim, cp_group, cp_partition_mode): + cp_size = cp_group.size() + cp_rank = cp_group.rank() + if cp_partition_mode == "zigzag": + cp_idx = torch.tensor([cp_rank, 2 * cp_size - cp_rank - 1], device=tensor.device) + elif cp_partition_mode == "contiguous": + cp_idx = torch.tensor([2 * cp_rank, 2 * cp_rank + 1], device=tensor.device) + else: + raise ValueError(f"Unsupported test CP partition mode {cp_partition_mode!r}.") + tensor = tensor.view(*tensor.shape[:seq_dim], 2 * cp_size, -1, *tensor.shape[(seq_dim + 1) :]) + tensor = tensor.index_select(seq_dim, cp_idx) + return tensor.view(*tensor.shape[:seq_dim], -1, *tensor.shape[(seq_dim + 2) :]) + - rank_indices = [ - get_thd_context_parallel_rank_indices(cu_seqlens, cp_size, rank, layout) +def _get_test_thd_token_indices(cu_seqlens, cp_size, cp_rank, cp_partition_mode): + cu = cu_seqlens.to(dtype=torch.long).tolist() + compact_cu = [cu[0]] + for value in cu[1:]: + if value != compact_cu[-1]: + compact_cu.append(value) + + total_tokens = compact_cu[-1] + if cp_partition_mode == "contiguous": + part_len = total_tokens // cp_size + start = cp_rank * part_len + return torch.arange(start, start + part_len, dtype=torch.long) + if cp_partition_mode != "zigzag": + raise ValueError(f"Unsupported test CP partition mode {cp_partition_mode!r}.") + + token_indices = [] + for seq_start, seq_end in zip(compact_cu[:-1], compact_cu[1:]): + chunk_len = (seq_end - seq_start) // (2 * cp_size) + first_start = seq_start + cp_rank * chunk_len + second_chunk = 2 * cp_size - cp_rank - 1 + second_start = seq_start + second_chunk * chunk_len + token_indices.extend(range(first_start, first_start + chunk_len)) + token_indices.extend(range(second_start, second_start + chunk_len)) + return torch.tensor(token_indices, dtype=torch.long) + + +@pytest.mark.parametrize( + ("source_layout", "target_layout"), [("zigzag", "contiguous"), ("contiguous", "zigzag")] +) +@pytest.mark.parametrize( + ("cp_size", "tp_size", "group_rank_by_logical_rank"), + [(3, 1, (0, 1, 2)), (2, 2, (0, 2, 1, 3)), (2, 4, tuple(range(8)))], +) +def test_sbhd_layout_redistribution_plan_reassembles_target_segments( + source_layout, target_layout, cp_size, tp_size, group_rank_by_logical_rank +): + group_size = cp_size * tp_size + plans = [None] * group_size + sends = [[None] * group_size for _ in range(group_size)] + + for logical_rank, group_rank in enumerate(group_rank_by_logical_rank): + cp_rank, tp_rank = divmod(logical_rank, tp_size) + source_ids = context_parallel_layout_conversion._local_sbhd_segment_ids( + layout=source_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + plan = context_parallel_layout_conversion._build_sbhd_layout_redistribution_plan( + source_layout=source_layout, + target_layout=target_layout, + cp_size=cp_size, + cp_rank=cp_rank, + tp_size=tp_size, + tp_rank=tp_rank, + group_rank_by_logical_rank=group_rank_by_logical_rank, + ) + plans[group_rank] = plan + packed_ids = tuple(source_ids[slot] for slot in plan.send_slots) + offset = 0 + for destination, count in enumerate(plan.input_segment_counts): + sends[group_rank][destination] = packed_ids[offset : offset + count] + offset += count + + for logical_rank, group_rank in enumerate(group_rank_by_logical_rank): + cp_rank, tp_rank = divmod(logical_rank, tp_size) + plan = plans[group_rank] + received_ids = tuple( + segment_id + for source_group_rank in range(group_size) + for segment_id in sends[source_group_rank][group_rank] + ) + output_ids = tuple(received_ids[index] for index in plan.receive_permutation) + assert output_ids == context_parallel_layout_conversion._local_sbhd_segment_ids( + layout=target_layout, cp_size=cp_size, cp_rank=cp_rank, tp_size=tp_size, tp_rank=tp_rank + ) + + +def test_sbhd_layout_redistribution_rejects_odd_tensor_parallel_size(): + with pytest.raises(ValueError, match="even tensor-parallel size"): + context_parallel_layout_conversion._sbhd_segments_per_rank(tp_size=3) + + +@pytest.mark.parametrize( + ("source_layout", "target_layout"), [("zigzag", "contiguous"), ("contiguous", "zigzag")] +) +@pytest.mark.parametrize( + ("cu_seqlens", "cp_size"), + [ + (torch.tensor([0, 16, 40]), 2), + (torch.tensor([0, 32, 96, 128]), 4), + (torch.tensor([0, 32, 96, 128, 128, 128]), 4), + ], +) +def test_thd_cp_partition_route_reassembles_target_layout( + source_layout, target_layout, cu_seqlens, cp_size +): + source_indices = [ + _get_test_thd_token_indices(cu_seqlens, cp_size, rank, source_layout) for rank in range(cp_size) ] + target_indices = [ + _get_test_thd_token_indices(cu_seqlens, cp_size, rank, target_layout) + for rank in range(cp_size) + ] + routes = [build_thd_cp_partition_route(cu_seqlens, cp_size, rank) for rank in range(cp_size)] + if source_layout == "zigzag" and target_layout == "contiguous": + selected_routes = [ + ( + route.zigzag_index, + route.contiguous_index, + route.zigzag_split_sizes, + route.contiguous_split_sizes, + ) + for route in routes + ] + else: + selected_routes = [ + ( + route.contiguous_index, + route.zigzag_index, + route.contiguous_split_sizes, + route.zigzag_split_sizes, + ) + for route in routes + ] + for rank, (_, _, send_split_sizes, recv_split_sizes) in enumerate(selected_routes): + assert sum(send_split_sizes) == source_indices[rank].numel() + assert sum(recv_split_sizes) == target_indices[rank].numel() + assert source_indices[rank].numel() == target_indices[rank].numel() + + send_buffers = [] + for rank, (send_index, _, _, _) in enumerate(selected_routes): + send_buffers.append( + source_indices[rank] + if send_index is None + else source_indices[rank].index_select(0, send_index) + ) + + for dst_rank in range(cp_size): + recv_chunks = [] + for src_rank in range(cp_size): + _, _, send_split_sizes, _ = selected_routes[src_rank] + send_offset = sum(send_split_sizes[:dst_rank]) + send_len = send_split_sizes[dst_rank] + recv_chunks.append(send_buffers[src_rank].narrow(0, send_offset, send_len)) + recv_buf = torch.cat(recv_chunks, dim=0) + _, recv_index, _, recv_split_sizes = selected_routes[dst_rank] + local_target_length = sum(recv_split_sizes) + if recv_index is None: + out = recv_buf + else: + out = torch.empty(local_target_length, dtype=recv_buf.dtype) + out.index_copy_(0, recv_index, recv_buf) + assert torch.equal(out, target_indices[dst_rank]) + - assert [indices.numel() for indices in rank_indices] == [32, 32, 32, 32] - assert torch.cat(rank_indices).sort().values.tolist() == list(range(128)) +def test_thd_cp_partition_route_stores_bidirectional_layout_views(): + route = build_thd_cp_partition_route(torch.tensor([0, 8, 12, 16]), cp_size=2, cp_rank=0) + assert isinstance(route, ThdCpRoute) + assert route.zigzag_index is None + assert route.zigzag_split_sizes == [4, 4] + assert route.contiguous_index.tolist() == [0, 1, 6, 7, 2, 3, 4, 5] + assert route.contiguous_split_sizes == [4, 4] -@pytest.mark.parametrize("layout", ["zigzag", "contiguous"]) -def test_thd_context_parallel_rank_indices_ignore_duplicate_boundaries(layout): - compact_cu_seqlens = torch.tensor([0, 16, 40]) - padded_cu_seqlens = torch.tensor([0, 16, 40, 40, 40]) + c2z_send_index = route.contiguous_index + c2z_recv_index = route.zigzag_index + c2z_send_splits = route.contiguous_split_sizes + c2z_recv_splits = route.zigzag_split_sizes + z2c_send_index = route.zigzag_index + z2c_recv_index = route.contiguous_index + z2c_send_splits = route.zigzag_split_sizes + z2c_recv_splits = route.contiguous_split_sizes - for rank in range(2): - assert torch.equal( - get_thd_context_parallel_rank_indices(padded_cu_seqlens, 2, rank, layout), - get_thd_context_parallel_rank_indices(compact_cu_seqlens, 2, rank, layout), + assert c2z_send_index is route.contiguous_index + assert c2z_recv_index is route.zigzag_index + assert c2z_send_splits is route.contiguous_split_sizes + assert c2z_recv_splits is route.zigzag_split_sizes + assert z2c_send_index is route.zigzag_index + assert z2c_recv_index is route.contiguous_index + assert z2c_send_splits is route.zigzag_split_sizes + assert z2c_recv_splits is route.contiguous_split_sizes + + +def test_build_thd_cp_partition_route_rejects_decreasing_boundaries(): + with pytest.raises(ValueError, match="nondecreasing"): + build_thd_cp_partition_route( + torch.tensor([0, 16, 8], dtype=torch.int32), cp_size=2, cp_rank=0 + ) + + +@pytest.mark.internal +@pytest.mark.parametrize( + ("source_layout", "target_layout"), [("zigzag", "contiguous"), ("contiguous", "zigzag")] +) +@pytest.mark.parametrize("seq_dim", [0, 1]) +def test_sbhd_convert_cp_partition_mode_matches_direct_target_shard( + source_layout, target_layout, seq_dim +): + if not torch.cuda.is_available() or Utils.world_size < 2: + pytest.skip("SBHD CP partition-mode conversion needs at least two CUDA ranks.") + + cp_size = 2 + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=cp_size) + try: + cp_group = parallel_state.get_context_parallel_group() + full_tensor = _make_sequence_tensor( + total_seq_len=32, + seq_dim=seq_dim, + device=torch.device(f"cuda:{torch.cuda.current_device()}"), + ) + source_shard = _get_sbhd_tensor_on_this_cp_rank( + full_tensor, seq_dim, cp_group, cp_partition_mode=source_layout + ) + + converted = context_parallel_layout_conversion.convert_cp_partition_mode( + x=source_shard, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + seq_dim=seq_dim, + cp_group=cp_group, + ) + expected = _get_sbhd_tensor_on_this_cp_rank( + full_tensor, seq_dim, cp_group, cp_partition_mode=target_layout + ) + + torch.testing.assert_close(converted, expected, atol=0.0, rtol=0.0) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.parametrize( + ("source_layout", "target_layout", "seq_dim", "sequence_parallel"), + [ + pytest.param("zigzag", "contiguous", 0, False, id="zigzag-contiguous-seq0"), + pytest.param("zigzag", "contiguous", 1, False, id="zigzag-contiguous-seq1"), + pytest.param("contiguous", "zigzag", 0, False, id="contiguous-zigzag-seq0"), + pytest.param("contiguous", "zigzag", 1, False, id="contiguous-zigzag-seq1"), + pytest.param("zigzag", "contiguous", 0, True, id="sp-zigzag-contiguous-seq0"), + pytest.param("zigzag", "contiguous", 1, True, id="sp-zigzag-contiguous-seq1"), + pytest.param("contiguous", "zigzag", 0, True, id="sp-contiguous-zigzag-seq0"), + pytest.param("contiguous", "zigzag", 1, True, id="sp-contiguous-zigzag-seq1"), + ], +) +def test_sbhd_convert_cp_partition_mode_backward_matches_direct_source_shard( + source_layout, target_layout, seq_dim, sequence_parallel +): + min_world_size = 4 if sequence_parallel else 2 + if not torch.cuda.is_available() or Utils.world_size < min_world_size: + pytest.skip( + f"SBHD CP partition-mode conversion backward needs at least {min_world_size} " + "CUDA ranks." + ) + + cp_size = 2 + tp_size = 2 if sequence_parallel else 1 + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, context_parallel_size=cp_size + ) + try: + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() if sequence_parallel else None + tp_cp_group = ( + parallel_state.get_tensor_and_context_parallel_group() if sequence_parallel else None + ) + full_tensor = _make_sequence_tensor( + total_seq_len=32, + seq_dim=seq_dim, + device=torch.device(f"cuda:{torch.cuda.current_device()}"), + ) + full_upstream_grad = full_tensor.mul(0.125).add(1.0) + source_shard = _get_sbhd_tensor_on_this_cp_rank( + full_tensor, seq_dim, cp_group, cp_partition_mode=source_layout + ) + if sequence_parallel: + source_shard = _get_sequence_parallel_shard(source_shard, seq_dim, tp_group) + source_shard = source_shard.detach().requires_grad_(True) + + convert_kwargs = ( + {"sequence_parallel": True, "tp_group": tp_group, "tp_cp_group": tp_cp_group} + if sequence_parallel + else {} + ) + converted = context_parallel_layout_conversion.convert_cp_partition_mode( + x=source_shard, + source_partition_mode=source_layout, + target_partition_mode=target_layout, + seq_dim=seq_dim, + cp_group=cp_group, + **convert_kwargs, + ) + expected_target = _get_sbhd_tensor_on_this_cp_rank( + full_tensor, seq_dim, cp_group, cp_partition_mode=target_layout + ) + if sequence_parallel: + expected_target = _get_sequence_parallel_shard(expected_target, seq_dim, tp_group) + torch.testing.assert_close(converted, expected_target, atol=0.0, rtol=0.0) + + target_upstream_grad = _get_sbhd_tensor_on_this_cp_rank( + full_upstream_grad, seq_dim, cp_group, cp_partition_mode=target_layout + ) + if sequence_parallel: + target_upstream_grad = _get_sequence_parallel_shard( + target_upstream_grad, seq_dim, tp_group + ) + converted.mul(target_upstream_grad).sum().backward() + expected_source_grad = _get_sbhd_tensor_on_this_cp_rank( + full_upstream_grad, seq_dim, cp_group, cp_partition_mode=source_layout ) + if sequence_parallel: + expected_source_grad = _get_sequence_parallel_shard( + expected_source_grad, seq_dim, tp_group + ) + torch.testing.assert_close(source_shard.grad, expected_source_grad, atol=0.0, rtol=0.0) + finally: + Utils.destroy_model_parallel() + + +def test_prebuild_thd_cp_partition_routes_populates_direct_fields(): + packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 16, 40]), + cu_seqlens_q_padded=None, + cp_partition_route=None, + ) + cp_group = _FakeGroup(size=2, rank=0) + prebuild_thd_cp_partition_routes(packed_seq_params, cp_group) + + route = get_thd_cp_partition_route(packed_seq_params, "zigzag", "contiguous") + same_route = get_thd_cp_partition_route(packed_seq_params, "zigzag", "contiguous") + reverse_route = get_thd_cp_partition_route(packed_seq_params, "contiguous", "zigzag") + + assert same_route is route + assert reverse_route is route + assert packed_seq_params.cp_partition_route is route + + +def test_prebuild_thd_cp_partition_routes_raises_route_errors(): + packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 10, 18]), + cu_seqlens_q_padded=None, + cp_partition_route=None, + ) + cp_group = _FakeGroup(size=2, rank=0) -def test_thd_context_parallel_rank_indices_reject_uneven_chunks(): with pytest.raises(ValueError, match="divisible"): - get_thd_context_parallel_rank_indices(torch.tensor([0, 10]), 2, 0, "zigzag") + prebuild_thd_cp_partition_routes(packed_seq_params, cp_group) -def test_thd_context_parallel_rank_indices_reject_decreasing_boundaries(): - with pytest.raises(ValueError, match="nondecreasing"): - get_thd_context_parallel_rank_indices(torch.tensor([0, 16, 8]), 2, 0, "zigzag") +def test_cp_partition_mode_converter_recurses_over_tensor_containers(monkeypatch): + calls = [] + + def fake_convert(*, x, cp_group, **kwargs): + calls.append((x, cp_group, kwargs)) + return x + 10 + + monkeypatch.setattr( + context_parallel_layout_conversion, "convert_cp_partition_mode", fake_convert + ) + cp_group = SimpleNamespace(size=lambda: 2) + tp_cp_group = object() + config = SimpleNamespace(cuda_graph_impl=None) + cu_seqlens = torch.tensor([0, 8]) + untouched = object() + value = (torch.tensor([1]), [None, untouched, torch.tensor([2])]) + route = object() + packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=None, + cp_partition_mode="zigzag", + cp_partition_route=route, + ) + + converter = CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode="zigzag", + target_partition_mode="contiguous", + config=config, + tp_cp_group=tp_cp_group, + ) + converted = converter.convert(value=value, seq_dim=lambda tensor: tensor.dim() - 1) + + assert torch.equal(converted[0], torch.tensor([11])) + assert converted[1][0] is None + assert converted[1][1] is untouched + assert torch.equal(converted[1][2], torch.tensor([12])) + assert [call[1] for call in calls] == [cp_group, cp_group] + assert [call[2]["seq_dim"] for call in calls] == [0, 0] + assert all(call[2]["cu_seqlens"] is cu_seqlens for call in calls) + assert all(call[2]["tp_cp_group"] is tp_cp_group for call in calls) + assert packed_seq_params.cp_partition_mode == "contiguous" + assert packed_seq_params.cp_partition_route is route + + +def test_cp_partition_mode_converter_rejects_thd_full_iteration_cuda_graph_conversion(): + cp_group = SimpleNamespace(size=lambda: 2) + packed_seq_params = SimpleNamespace(qkv_format="thd") + config = SimpleNamespace(cuda_graph_impl="full_iteration") + + CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode="zigzag", + target_partition_mode="zigzag", + config=config, + ) + + with pytest.raises(ValueError, match="Full-iteration CUDA graph"): + CpPartitionModeConverter( + cp_group=cp_group, + packed_seq_params=packed_seq_params, + source_partition_mode="zigzag", + target_partition_mode="contiguous", + config=config, + ) + + +def test_module_input_conversion_treats_missing_packed_seq_params_as_sbhd(monkeypatch): + calls = [] + + def fake_convert(*, x, cp_group, **kwargs): + calls.append((x, cp_group, kwargs)) + return x + 1 + + monkeypatch.setattr( + context_parallel_layout_conversion, "convert_cp_partition_mode", fake_convert + ) + cp_group = SimpleNamespace(size=lambda: 2) + tp_cp_group = object() + hidden_states = torch.ones(8, 1, 4) + converted, converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=None, + cp_group=cp_group, + tp_group=None, + tp_cp_group=tp_cp_group, + target_partition_mode="contiguous", + sequence_parallel=False, + config=SimpleNamespace(cp_partition_mode="zigzag", cuda_graph_impl=None), + ) + + assert converter is not None + assert torch.equal(converted, hidden_states + 1) + assert calls[0][2]["source_partition_mode"] == "zigzag" + assert calls[0][2]["target_partition_mode"] == "contiguous" + assert calls[0][2]["cu_seqlens"] is None + assert calls[0][2]["tp_cp_group"] is tp_cp_group + assert converter.tp_cp_group is tp_cp_group + + +def test_public_conversion_apis_default_to_no_cp_group(): + hidden_states = torch.ones(8, 1, 4) + config = SimpleNamespace(cp_partition_mode="zigzag", cuda_graph_impl=None) + converter = CpPartitionModeConverter( + packed_seq_params=None, + source_partition_mode="zigzag", + target_partition_mode="contiguous", + config=config, + ) + + assert converter.convert(value=hidden_states) is hidden_states + converted, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=None, + target_partition_mode="contiguous", + sequence_parallel=False, + config=config, + ) + assert converted is hidden_states + assert back_to_input_converter is None + assert ( + context_parallel_layout_conversion.convert_cp_partition_mode( + x=hidden_states, source_partition_mode="zigzag", target_partition_mode="contiguous" + ) + is hidden_states + ) + + +@pytest.mark.parametrize( + ("sequence_parallel", "tp_size"), + [(False, None), (False, 2), (True, None), (True, 1), (True, 2)], +) +def test_sbhd_conversion_uses_one_redistribution_path(monkeypatch, sequence_parallel, tp_size): + calls = [] + cp_group = _FakeGroup(size=2, rank=0) + tp_group = _FakeGroup(size=tp_size, rank=0) if tp_size is not None else None + tp_cp_group = _FakeGroup(size=2 * tp_size, rank=0) if tp_size is not None else None + x = torch.arange(24).view(2, 6, 2) + + def fake_redistribute(**kwargs): + calls.append(kwargs) + return kwargs["input_"] + 1 + + monkeypatch.setattr( + context_parallel_layout_conversion, "_redistribute_sbhd_layout", fake_redistribute + ) + + converted = context_parallel_layout_conversion.convert_cp_partition_mode( + x=x, + source_partition_mode="zigzag", + target_partition_mode="contiguous", + seq_dim=1, + sequence_parallel=sequence_parallel, + cp_group=cp_group, + tp_group=tp_group, + tp_cp_group=tp_cp_group, + ) + + torch.testing.assert_close(converted, x + 1) + assert len(calls) == 1 + call = calls[0] + assert torch.equal(call.pop("input_"), x.movedim(1, 0)) + assert call == { + "cp_group": cp_group, + "source_layout": "zigzag", + "target_layout": "contiguous", + "sequence_parallel": sequence_parallel, + "tp_group": tp_group, + "tp_cp_group": tp_cp_group, + } + + +def test_sequence_parallel_thd_conversion_warns_about_naive_fallback(monkeypatch): + from megatron.core.tensor_parallel import mappings + + calls = [] + cp_group = _FakeGroup(size=2, rank=0) + tp_group = _FakeGroup(size=2, rank=0) + cu_seqlens = torch.tensor([0, 12]) + x = torch.arange(24).view(2, 6, 2) + + def fake_gather(*, input_, tensor_parallel_output_grad, group): + calls.append(("gather", tensor_parallel_output_grad, group)) + return input_ + + def fake_redistribute(**kwargs): + calls.append(("thd", kwargs)) + return kwargs["x"] + + def fake_scatter(*, input_, group): + calls.append(("scatter", group)) + return input_ + + monkeypatch.setattr(mappings, "gather_from_sequence_parallel_region", fake_gather) + monkeypatch.setattr( + context_parallel_layout_conversion, "_redistribute_thd_layout", fake_redistribute + ) + monkeypatch.setattr(mappings, "scatter_to_sequence_parallel_region", fake_scatter) + + with pytest.warns(RuntimeWarning, match="naive TP gather"): + converted = context_parallel_layout_conversion.convert_cp_partition_mode( + x=x, + source_partition_mode="zigzag", + target_partition_mode="contiguous", + seq_dim=1, + cu_seqlens=cu_seqlens, + sequence_parallel=True, + cp_group=cp_group, + tp_group=tp_group, + ) -def test_thd_context_parallel_rank_indices_reject_unknown_layout(): - with pytest.raises(ValueError, match="Unsupported"): - get_thd_context_parallel_rank_indices(torch.tensor([0, 16]), 2, 0, "interleaved") + assert torch.equal(converted, x) + assert calls[0] == ("gather", False, tp_group) + call_name, thd_call = calls[1] + assert call_name == "thd" + assert torch.equal(thd_call.pop("x"), x.movedim(1, 0)) + assert thd_call.pop("cu_seqlens") is cu_seqlens + assert thd_call == { + "cp_group": cp_group, + "seq_dim": 0, + "source_partition_mode": "zigzag", + "target_partition_mode": "contiguous", + "thd_cp_partition_route": None, + } + assert calls[2] == ("scatter", tp_group) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_integration.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_integration.py index 85f4c8fa564..2cecc98009a 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_integration.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_integration.py @@ -312,7 +312,7 @@ def forward(self, query, key, *, value, attention_mask, **kwargs): def test_absorbed_mla_forward_uses_and_restores_dynamic_cp_group(): original_cp_group = object() - dynamic_cp_group = object() + dynamic_cp_group = SimpleNamespace(size=lambda: 2) pg_collection = SimpleNamespace(cp=original_cp_group) observed_groups = [] diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index 6a65a91a2fc..e38b06a9c9a 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy from unittest import mock @@ -10,10 +10,8 @@ from torch.nn import functional as F import megatron.core.parallel_state as parallel_state +from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd from megatron.core.hyper_comm_grid import HyperCommGrid -from megatron.core.models.common.embeddings.rope_utils import ( - get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank, -) from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_spec, get_gpt_layer_with_transformer_engine_spec, @@ -715,6 +713,8 @@ def _test_parallel_attention_correctness( sequence_length=256, micro_batch_size=4, sequence_packing=False, + cp_partition_mode="zigzag", + compare_param_grads=False, ): # Model initialization function def initialize_gpt_model( @@ -766,8 +766,24 @@ def initialize_gpt_model( mock_args.no_load_rng = True save_checkpoint(10, gpt_model, None, None, 0) + def get_param_grad(param): + grad = param.grad + if grad is None: + grad = getattr(param, "main_grad", None) + if grad is None: + return torch.zeros_like(param, dtype=torch.float32) + return grad + + def zero_param_grads(module): + for param in module.parameters(): + param.grad = None + main_grad = getattr(param, "main_grad", None) + if main_grad is not None: + main_grad.zero_() + # Calculate baseline output attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(attention) output_hidden_states_baseline, bias_hidden_states_baseline = attention( input_hidden_states, attention_mask=None ) @@ -775,6 +791,13 @@ def initialize_gpt_model( # Save baseline output input_grad_baseline = input_hidden_states.grad.detach() + param_grads_baseline = None + if compare_param_grads: + param_grads_baseline = { + name: get_param_grad(param).detach().float().clone() + for name, param in attention.named_parameters() + if param.requires_grad + } output_hidden_states_baseline = output_hidden_states_baseline.detach() bias_hidden_states_baseline = bias_hidden_states_baseline if bias_hidden_states_baseline is not None: @@ -793,9 +816,11 @@ def initialize_gpt_model( transformer_config.context_parallel_size = cp transformer_config.tensor_model_parallel_size = tp transformer_config.sequence_parallel = sp + transformer_config.cp_partition_mode = cp_partition_mode init_basic_mock_args(mock_args, tp, 1, bf16=True) mock_args.context_parallel_size = cp mock_args.sequence_parallel = sp + mock_args.cp_partition_mode = cp_partition_mode gpt_model = unwrap_model(get_model(initialize_gpt_model, config=transformer_config)) with mock.patch('megatron.training.checkpointing.check_checkpoint_args'): with mock.patch('megatron.training.checkpointing.update_num_microbatches'): @@ -806,10 +831,31 @@ def initialize_gpt_model( tp_rank = parallel_state.get_tensor_model_parallel_rank() def get_tensor_on_this_rank(tensor): - if cp > 1: - tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) if sequence_packing: tensor = tensor.transpose(0, 1).contiguous().view(-1, 1, *tensor.shape[2:]) + if cp > 1: + cu_seqlens_tensor = torch.tensor( + [i * sequence_length for i in range(micro_batch_size + 1)], + device=tensor.device, + ) + batch = {"hidden_states": tensor, "cu_seqlens_padded": cu_seqlens_tensor} + get_cp_slice_for_thd( + batch, + cp_group, + keys=("hidden_states",), + cp_partition_mode=cp_partition_mode, + ) + tensor = batch["hidden_states"] + elif cp > 1: + cp_rank = torch.distributed.get_rank(cp_group) + if cp_partition_mode == "zigzag": + cp_idx = torch.tensor([cp_rank, 2 * cp - cp_rank - 1], device=tensor.device) + elif cp_partition_mode == "contiguous": + cp_idx = torch.tensor([2 * cp_rank, 2 * cp_rank + 1], device=tensor.device) + else: + raise ValueError(f"Unsupported test CP partition mode {cp_partition_mode!r}.") + tensor = tensor.view(2 * cp, -1, *tensor.shape[1:]) + tensor = tensor.index_select(0, cp_idx).view(-1, *tensor.shape[2:]) if tp > 1 and sp: sp_seg = tensor.shape[0] // tp tensor = tensor[tp_rank * sp_seg : (tp_rank + 1) * sp_seg] @@ -819,15 +865,27 @@ def get_tensor_on_this_rank(tensor): if sequence_packing: cu_seqlens = [i * sequence_length for i in range(micro_batch_size + 1)] packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + packed_seq_params.cp_partition_mode = cp_partition_mode else: packed_seq_params = None input_hidden_states = get_tensor_on_this_rank(input_hidden_states) input_hidden_states = input_hidden_states.detach().requires_grad_(True) parallel_attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(parallel_attention) output_hidden_states_parallel, bias_hidden_states_parallel = parallel_attention( input_hidden_states, attention_mask=None, packed_seq_params=packed_seq_params ) output_hidden_states_parallel.sum().backward() + param_grads_parallel = None + if compare_param_grads: + assert tp == 1, "Parameter gradient parity only supports unsharded parameters." + param_grads_parallel = {} + for name, param in parallel_attention.named_parameters(): + if param.requires_grad: + grad = get_param_grad(param) + if cp > 1: + torch.distributed.all_reduce(grad, group=cp_group) + param_grads_parallel[name] = grad.detach().float().clone() input_grad_parallel = input_hidden_states.grad.detach() # Check if the output is close @@ -904,6 +962,12 @@ def assert_close_or_cosine_similarity(baseline, parallel, tensor_name): assert_close_or_cosine_similarity( bias_hidden_states_baseline, bias_hidden_states_parallel, "bias_hidden_states" ) + if compare_param_grads: + assert param_grads_baseline.keys() == param_grads_parallel.keys() + for name, grad_baseline in param_grads_baseline.items(): + assert_close_or_cosine_similarity( + grad_baseline, param_grads_parallel[name], f"param_grad[{name}]" + ) Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/test_thd_correctness.py b/tests/unit_tests/transformer/test_thd_correctness.py index ac759c2bdce..2f9815d366e 100644 --- a/tests/unit_tests/transformer/test_thd_correctness.py +++ b/tests/unit_tests/transformer/test_thd_correctness.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """ Compare THD format against SBHD format. @@ -30,14 +30,31 @@ import torch.nn as nn from megatron.core import parallel_state +from megatron.core.context_parallel_layout import prebuild_thd_cp_partition_routes +from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_transformer_block_with_experimental_attention_variant_spec, +) +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + get_gpt_mtp_block_spec, +) +from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler, MTPLossLoggingHelper from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer from tests.unit_tests.test_utilities import Utils +try: + import fla # noqa: F401 + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + # ============================================================================= # Constants # ============================================================================= @@ -217,6 +234,7 @@ def to_cu_seqlens(lens): max_seqlen_q=max(padded), max_seqlen_kv=max(padded), qkv_format='thd', + cp_partition_mode="zigzag", ) @@ -1061,3 +1079,282 @@ def test_dynamic_cp_format(tc: DynamicCPTestCase): # === Cleanup === Utils.destroy_model_parallel() + + +# ============================================================================= +# Mixed GDN/GQA Model Correctness +# ============================================================================= + + +def _make_mixed_model_config( + *, + linear_cp_mode: str, + cp_partition_mode: str, + dynamic_context_parallel: bool, + context_parallel_size: int, +) -> TransformerConfig: + layer_pattern = [1, 1, 0, 1, 0] + return TransformerConfig( + num_layers=len(layer_pattern), + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=8, + num_query_groups=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=8, + activation_func=torch.nn.functional.silu, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=layer_pattern, + linear_cp_mode=linear_cp_mode, + cp_partition_mode=cp_partition_mode, + context_parallel_size=context_parallel_size, + dynamic_context_parallel=dynamic_context_parallel, + cp_comm_type="p2p", + sequence_packing_scheduler=( + "default_dynamic_cp" if dynamic_context_parallel else "dp_balanced" + ), + pad_packed_seq_alignment="max", + max_seqlen_per_dp_cp_rank=128, + hidden_dropout=0.0, + attention_dropout=0.0, + calculate_per_token_loss=True, + bf16=True, + params_dtype=torch.bfloat16, + mtp_num_layers=1, + ) + + +def _build_mixed_model(model_type: str, config: TransformerConfig): + model_kwargs = { + "config": config, + "vocab_size": 512, + "max_sequence_length": config.max_seqlen_per_dp_cp_rank, + "position_embedding_type": "rope", + } + if model_type == "gpt": + transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec( + config=config + ) + mtp_block_spec = get_gpt_mtp_block_spec( + config=config, spec=transformer_layer_spec, use_transformer_engine=True + ) + model = GPTModel( + transformer_layer_spec=transformer_layer_spec, + mtp_block_spec=mtp_block_spec, + **model_kwargs, + ) + elif model_type == "hybrid": + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + from megatron.core.models.hybrid.hybrid_model import HybridModel + + model = HybridModel( + hybrid_stack_spec=hybrid_stack_spec, hybrid_layer_pattern="GG*G*/*", **model_kwargs + ) + else: + raise ValueError(f"Unsupported model type: {model_type}") + + return model.cuda() + + +def _prepare_mixed_model_batch(seq_indices, cp_group, config, vocab_size): + all_sequence_lengths = (37, 53, 61, 71, 43, 59, 67, 79) + sequence_lengths = [all_sequence_lengths[i] for i in seq_indices] + padded_lengths = [_round_up(sequence_length, 16) for sequence_length in sequence_lengths] + total_tokens = sum(padded_lengths) + device = torch.device("cuda", torch.cuda.current_device()) + + tokens = torch.zeros(total_tokens, device=device, dtype=torch.long) + labels = torch.zeros_like(tokens) + loss_mask = torch.zeros(total_tokens, device=device, dtype=torch.float32) + padding_mask = torch.ones(total_tokens, device=device, dtype=torch.bool) + position_ids = torch.zeros_like(tokens) + cu_seqlens = [0] + cu_seqlens_padded = [0] + + offset = 0 + for sequence_index, sequence_length, padded_length in zip( + seq_indices, sequence_lengths, padded_lengths + ): + valid_end = offset + sequence_length + sequence_tokens = ( + torch.arange(sequence_length, device=device, dtype=torch.long) + 17 * sequence_index + 3 + ) % vocab_size + tokens[offset:valid_end] = sequence_tokens + labels[offset:valid_end] = (sequence_tokens + 11) % vocab_size + loss_mask[offset:valid_end] = 1.0 + padding_mask[offset:valid_end] = False + position_ids[offset:valid_end] = torch.arange( + sequence_length, device=device, dtype=torch.long + ) + cu_seqlens.append(cu_seqlens[-1] + sequence_length) + offset += padded_length + cu_seqlens_padded.append(offset) + + cu_seqlens = torch.tensor(cu_seqlens, device=device, dtype=torch.int32) + cu_seqlens_padded = torch.tensor(cu_seqlens_padded, device=device, dtype=torch.int32) + sequence_keys = ("tokens", "labels", "loss_mask", "padding_mask", "position_ids") + batch = { + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "padding_mask": padding_mask, + "position_ids": position_ids, + "cu_seqlens_padded": cu_seqlens_padded, + } + get_cp_slice_for_thd( + batch, cp_group, keys=sequence_keys, cp_partition_mode=config.cp_partition_mode + ) + batch = {name: batch[name].view(1, -1) for name in sequence_keys} + + packed_seq_params = PackedSeqParams( + qkv_format="thd", + 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=max(padded_lengths), + max_seqlen_kv=max(padded_lengths), + local_cp_size=cp_group.size() if config.dynamic_context_parallel else None, + cp_group=cp_group, + cp_partition_mode=config.cp_partition_mode, + pad_between_seqs=True, + ) + prebuild_thd_cp_partition_routes(packed_seq_params, cp_group) + return batch, packed_seq_params + + +def _run_mixed_model(model, batch, packed_seq_params, dp_cp_group): + MTPLossLoggingHelper.tracker = {} + MTPLossLoggingHelper.configure_acceptance_collection(enabled=False) + + loss = model( + input_ids=batch["tokens"], + position_ids=batch["position_ids"], + attention_mask=None, + labels=batch["labels"], + loss_mask=batch["loss_mask"], + packed_seq_params=packed_seq_params, + padding_mask=batch["padding_mask"], + ) + local_numerator = (loss.float() * batch["loss_mask"]).sum() + local_denominator = batch["loss_mask"].sum() + global_stats = torch.stack([local_numerator.detach(), local_denominator.detach()]) + dist.all_reduce(global_stats, group=dp_cp_group) + global_denominator = global_stats[1].clamp(min=1) + + MTPLossAutoScaler.set_loss_scale(global_denominator.reciprocal()) + (local_numerator / global_denominator).backward() + + MTPLossLoggingHelper.reduce_loss_in_tracker() + assert "values" in MTPLossLoggingHelper.tracker + mtp_loss = MTPLossLoggingHelper.tracker["values"].detach().float().clone() + + grads = [ + (name, param.grad) for name, param in model.named_parameters() if param.grad is not None + ] + assert grads, "Mixed GDN/GQA model did not produce parameter gradients." + grad_names, grad_tensors = zip(*grads) + grad_vector = torch.cat([grad.detach().float().reshape(-1) for grad in grad_tensors]) + dist.all_reduce(grad_vector, group=dp_cp_group) + return global_stats[0] / global_denominator, mtp_loss, grad_names, grad_vector + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.parametrize("model_type", ("gpt", "hybrid")) +@pytest.mark.parametrize("dynamic_context_parallel", (False, True), ids=("thd_cp", "dcp")) +def test_mixed_gdn_gqa_model_cp_correctness(model_type, dynamic_context_parallel): + """Compare mixed GDN/GQA/MTP models against a THD CP=1 baseline. + + Each of the {fixed THD CP, DCP} x {GPTModel, HybridModel} cases uses a no-CP + reference that processes one complete packed sequence per rank. + """ + if not torch.cuda.is_available() or Utils.world_size != 8: + pytest.skip("Mixed GDN/GQA model CP correctness requires exactly 8 CUDA ranks.") + + seed = 1234 + reference_config = _make_mixed_model_config( + linear_cp_mode="chunkwise", + cp_partition_mode="zigzag", + dynamic_context_parallel=False, + context_parallel_size=1, + ) + Utils.initialize_model_parallel(context_parallel_size=1) + try: + reference_dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True) + reference_cp_group = parallel_state.get_context_parallel_group() + reference_seq_indices = [dist.get_rank(group=reference_dp_cp_group)] + + torch.manual_seed(seed) + model_parallel_cuda_manual_seed(seed) + reference_model = _build_mixed_model(model_type, reference_config) + reference_batch, reference_packed_seq_params = _prepare_mixed_model_batch( + reference_seq_indices, reference_cp_group, reference_config, reference_model.vocab_size + ) + reference_stats = _run_mixed_model( + reference_model, reference_batch, reference_packed_seq_params, reference_dp_cp_group + ) + reference_state_dict = { + name: value.detach().cpu().clone() if torch.is_tensor(value) else value + for name, value in reference_model.state_dict().items() + } + + assert reference_cp_group.size() == 1 + assert reference_packed_seq_params.qkv_format == "thd" + assert reference_packed_seq_params.cp_partition_mode == "zigzag" + finally: + MTPLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel() + + del reference_model, reference_batch, reference_packed_seq_params + + candidate_config = _make_mixed_model_config( + linear_cp_mode="chunkwise", + cp_partition_mode="contiguous", + dynamic_context_parallel=dynamic_context_parallel, + context_parallel_size=2, + ) + Utils.initialize_model_parallel( + context_parallel_size=2, dynamic_context_parallel=dynamic_context_parallel + ) + try: + candidate_dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True) + if dynamic_context_parallel: + dp_cp_rank = dist.get_rank(group=candidate_dp_cp_group) + if dp_cp_rank < 4: + local_cp_size, candidate_seq_indices = 4, range(4) + elif dp_cp_rank < 6: + local_cp_size, candidate_seq_indices = 2, range(4, 6) + else: + local_cp_size, candidate_seq_indices = 1, [dp_cp_rank] + candidate_cp_group = parallel_state.get_dynamic_data_context_parallel_groups( + group_size=local_cp_size + ) + else: + dp_rank = parallel_state.get_data_parallel_rank() + candidate_seq_indices = range(2 * dp_rank, 2 * dp_rank + 2) + candidate_cp_group = parallel_state.get_context_parallel_group() + + torch.manual_seed(seed) + model_parallel_cuda_manual_seed(seed) + candidate_model = _build_mixed_model(model_type, candidate_config) + candidate_model.load_state_dict(reference_state_dict) + candidate_batch, candidate_packed_seq_params = _prepare_mixed_model_batch( + candidate_seq_indices, candidate_cp_group, candidate_config, candidate_model.vocab_size + ) + candidate_stats = _run_mixed_model( + candidate_model, candidate_batch, candidate_packed_seq_params, candidate_dp_cp_group + ) + reference_loss, reference_mtp_loss, reference_grad_names, reference_grads = reference_stats + candidate_loss, candidate_mtp_loss, candidate_grad_names, candidate_grads = candidate_stats + + assert reference_grad_names == candidate_grad_names + assert candidate_packed_seq_params.cp_partition_mode == "contiguous" + torch.testing.assert_close(candidate_loss, reference_loss, atol=5e-3, rtol=0.0) + torch.testing.assert_close(candidate_mtp_loss, reference_mtp_loss, atol=5e-3, rtol=0.0) + assert_close("aggregated parameter gradients", candidate_grads, reference_grads, False) + finally: + MTPLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel()