diff --git a/megatron/core/context_parallel_layout.py b/megatron/core/context_parallel_layout.py new file mode 100644 index 00000000000..44014581fd5 --- /dev/null +++ b/megatron/core/context_parallel_layout.py @@ -0,0 +1,307 @@ +# 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/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 64bce710993..4647fb4fd14 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1827,6 +1827,19 @@ def forward( if packed_seq_params is not None else {} ) + if ( + packed_seq_kwargs.get("qkv_format") == "thd" + and packed_seq_kwargs.get("pad_between_seqs") is False + ): + # Megatron represents end padding as dummy THD sequences. TE DPA + # sizes THD outputs from cu_seqlens_q/kv, so pass padded + # boundaries as the effective attention boundaries while keeping + # the original PackedSeqParams metadata intact for downstream + # loss/routing paths. + if packed_seq_kwargs.get("cu_seqlens_q_padded") is not None: + packed_seq_kwargs["cu_seqlens_q"] = packed_seq_kwargs["cu_seqlens_q_padded"] + if packed_seq_kwargs.get("cu_seqlens_kv_padded") is not None: + packed_seq_kwargs["cu_seqlens_kv"] = packed_seq_kwargs["cu_seqlens_kv_padded"] qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format) attention_bias_kwargs = {} diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 5422ead7fd1..212af51ad1b 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -16,6 +16,10 @@ from torch import Tensor from megatron.core import tensor_parallel +from megatron.core.context_parallel_layout import ( + contiguous_to_zigzag_chunks, + zigzag_to_contiguous_chunks, +) from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.fp8_utils import get_fp8_align_size @@ -45,6 +49,7 @@ try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True @@ -58,12 +63,6 @@ logger = logging.getLogger(__name__) -# Triton's autotune key for causal_conv1d includes cdiv(total_tokens, 1024). -# Dynamic CP causes total_tokens to vary per microbatch, triggering repeated -# autotuning. Aligning to this boundary collapses most variations into a -# small number of buckets. -_CONV_PAD_ALIGNMENT = 4096 - @dataclass class GatedDeltaNetSubmodules: @@ -155,21 +154,23 @@ def __init__( self.qk_dim_local_tp = self.qk_dim // self.tp_size self.v_dim_local_tp = self.v_dim // self.tp_size - # GDN uses head-parallel CP: each CP rank handles a slice of heads. - # The static cp_size (== max dynamic cp_size) must evenly divide the - # per-TP head counts so that every possible runtime cp_size also divides. - num_key_heads_per_tp = self.num_key_heads // self.tp_size - num_value_heads_per_tp = self.num_value_heads // self.tp_size - assert num_key_heads_per_tp % self.cp_size == 0, ( - f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " - f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " - f"all runtime dynamic cp_size values divide the static one and so will also divide." - ) - assert num_value_heads_per_tp % self.cp_size == 0, ( - f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " - f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " - f"all runtime dynamic cp_size values divide the static one and so will also divide." - ) + # Headwise CP uses head-parallel layout: each CP rank handles a slice of + # heads. The static cp_size (== max dynamic cp_size) must evenly divide + # the per-TP head counts so that every possible runtime cp_size also + # divides. Chunkwise CP keeps heads local and does not need this split. + if self.config.linear_cp_mode == "headwise": + num_key_heads_per_tp = self.num_key_heads // self.tp_size + num_value_heads_per_tp = self.num_value_heads // self.tp_size + assert num_key_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + assert num_value_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) # Input projection (hidden_states -> q, k, v, gate, beta, alpha) # TODO: for now, output gate is forced for GDN. @@ -276,6 +277,16 @@ def __init__( if self.config.recompute_granularity == "selective" and self.config.recompute_modules: self.recompute_gdn = "gdn" in self.config.recompute_modules + # Cache for CP context objects consumed by FLA kernels. Rebuilding these per-forward + # is unsafe under CUDA graph capture because build_cp_context allocates + # fresh tensors whose memory pointers are baked into the captured graph; + # on the next call those tensors are reallocated, leaving the replayed + # graph pointing at stale memory. For non-packed (SBHD) input the + # cu_seqlens is fully determined by the (static) global sequence length + # and batch size, so we cache the (cu_seqlens, cp_context) pair keyed on + # both values. + self._chunkwise_cp_context_cache = {} + self.reset_parameters() def reset_parameters(self): @@ -308,6 +319,7 @@ def forward( packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, + pg_collection: Optional[ProcessGroupCollection] = None, inference_params: Optional[BaseInferenceContext] = None, **kwargs, ): @@ -331,11 +343,34 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) - cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) - cp_size = cp_group.size() + # Route the CP group to either the headwise (Ulysses-style) path or the + # chunkwise CP path according to config.linear_cp_mode. The two paths + # are mutually exclusive — whichever one is active owns the full CP + # group, and the other is given a size-1 group (None). The unused-path + # helpers already treat a None group as size 1, avoiding a costly and + # CUDA-graph-unsafe `torch.distributed.new_group` on every forward. + base_cp_group = pg_collection.cp if pg_collection is not None else self.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 + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + 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 - seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size * cp_size + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise if inference_context is not None: assert ( @@ -352,17 +387,19 @@ def forward( ), "Packed sequence does not support deterministic mode." # Resolve cu_seqlens with alignment padding handling. + # cu_seqlens in packed_seq_params is the global (pre-CP-split) cu_seqlens, so we + # validate against the global sequence length. cu_seqlens_q = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q, - seq_len, + seq_len_global, "cu_seqlens_q", cp_size=self.cp_size, ) cu_seqlens_kv = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_kv_padded, packed_seq_params.cu_seqlens_kv, - seq_len, + seq_len_global, "cu_seqlens_kv", cp_size=self.cp_size, ) @@ -379,29 +416,84 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + # Non-packed input: the only sources of cu_seqlens are the static + # global sequence length and batch size. Cache both the cu_seqlens + # tensor and the resulting chunkwise CP context so we don't + # reallocate them on every forward — those reallocations break + # CUDA graph capture. + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + if self.recompute_gdn and self.training: def _checkpointed_compute(hidden_states): return self._forward_compute( hidden_states, batch, - seq_len, - cp_size, - cp_group, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, cu_seqlens_q, packed_seq_params, + chunkwise_cp_context, ) out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) else: out, out_bias = self._forward_compute( - hidden_states, batch, seq_len, cp_size, cp_group, cu_seqlens_q, packed_seq_params + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, ) return out, out_bias def _forward_compute( - self, hidden_states, batch, seq_len, cp_size, cp_group, cu_seqlens_q, packed_seq_params + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, ): """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> gated norm -> out_proj). @@ -417,14 +509,37 @@ def _forward_compute( qkvzba, _ = self.in_proj(hidden_states) nvtx_range_pop(suffix="in_proj") + # Chunkwise CP expects the contiguous-time chunk layout (rank r holds chunks + # [2r, 2r+1]) inside conv1d / chunk_gated_delta_rule. Megatron attention CP + # feeds us the zigzag attention-load-balanced layout (rank r holds + # [r, 2*cp-r-1]), so reshuffle chunks over the CP group with a single + # all-to-all — no full-sequence gather required. + # 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 = self._a2a_cp_to_hp( - qkvzba, cp_size, cp_group, cu_seqlens_q, seq_len, packed_seq_params + qkvzba, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, ) if self.gdn_pre_gated_delta_rule_fusion: - # Existing all-to-all CP reaches this point in hidden-parallel layout. - # Chunkwise CP has a different sequence and halo contract; revisit - # fused pre-GDR metadata before enabling this path there. + assert cp_size_chunkwise == 1, ( + "gdn_pre_gated_delta_rule_fusion is not supported with chunkwise CP. " + "Disable gdn_pre_gated_delta_rule_fusion or use linear_cp_mode='headwise'." + ) nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") seq_idx = ( packed_seq_params.seq_idx @@ -437,8 +552,29 @@ def _forward_compute( nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") else: nvtx_range_push(suffix="pre_gated_delta_rule") + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + # TODO: If additional gated delta rule backends are added, handle this + # SBHD + chunkwise CP + batch>1 case per backend instead of + # unconditionally rejecting it. + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "because the FLA gated delta rule backend requires a single batch dimension " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) query, key, value, gate, beta, g = self.pre_gated_delta_rule( - qkvzba, batch, seq_len, cp_size, cp_group, cu_seqlens_q + qkvzba, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, ) nvtx_range_pop(suffix="pre_gated_delta_rule") @@ -453,6 +589,7 @@ def _forward_compute( output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, ) nvtx_range_pop(suffix="gated_delta_rule") @@ -463,11 +600,28 @@ def _forward_compute( # Transpose: b s x --> s b x # From bshd back to sbhd format - norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.reshape(batch, seq_len_post_headwise, -1) norm_out = norm_out.transpose(0, 1).contiguous() + # Inverse of the zigzag -> contiguous reshuffle performed before conv1d. + # Restores the Megatron attention-load-balanced layout that downstream + # layers and loss computation expect. + # 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") + norm_out = self._a2a_hp_to_cp( - norm_out, cp_size, cp_group, packed_seq_params, thd_cp_a2a_inv + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv ) # Output projection @@ -477,7 +631,17 @@ def _forward_compute( return out, out_bias - def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seqlens_q=None): + def pre_gated_delta_rule( + self, + qkvzba, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + ): """Prepare QKV, gate, beta, and decay tensors before the gated delta rule.""" # Transpose: s b x --> b s x @@ -488,10 +652,10 @@ def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seq qkv, gate, beta, alpha = torch.split( qkvzba, [ - (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size, - self.v_dim_local_tp // cp_size, - self.num_value_heads // self.tp_size // cp_size, - self.num_value_heads // self.tp_size // cp_size, + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, ], dim=-1, ) @@ -499,22 +663,30 @@ def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seq beta = beta.reshape(batch, seq_len, -1) alpha = alpha.reshape(batch, seq_len, -1) + kernel_batch = batch + kernel_seq_len = seq_len + # Convolution on qkv nvtx_range_push(suffix="conv1d") - seq_len = qkv.shape[1] + assert ( + qkv.shape[1] == kernel_seq_len + ), f"Shape mismatch: {qkv.shape[1]=} != {kernel_seq_len=}" qkv_channels_split_sections = [ self.qk_dim_local_tp, self.qk_dim_local_tp, self.v_dim_local_tp, ] - conv1d_weight = get_parameter_local_cp( - self.conv1d.weight, dim=0, cp_group=cp_group, split_sections=qkv_channels_split_sections + conv1d_weight = get_parameter_local_cp_headwise( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, ) conv1d_bias = ( - get_parameter_local_cp( + get_parameter_local_cp_headwise( self.conv1d.bias, dim=0, - cp_group=cp_group, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) if self.conv_bias @@ -529,18 +701,31 @@ def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seq stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // cp_size, + groups=self.conv_dim_local_tp // cp_size_headwise, ) - qkv = self.act_fn(conv_out[..., :seq_len]) + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d else: assert self.activation in ["silu", "swish"] _orig_seq = qkv.shape[1] - _pad_n = -_orig_seq % _CONV_PAD_ALIGNMENT - _conv_input = qkv + _pad_n = 0 + _conv_input = qkv.contiguous() _conv_cu_seqlens = cu_seqlens_q + _conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + _pad_n = -_orig_seq % self.config.gdn_conv_pad_alignment if _pad_n > 0: - _conv_input = torch.nn.functional.pad(qkv, (0, 0, 0, _pad_n)) + _conv_input = torch.nn.functional.pad(_conv_input, (0, 0, 0, _pad_n)) # cu_seqlens_q is None in non-packed-sequence mode; only the # last-segment offset needs to grow to cover the padding tail. if cu_seqlens_q is not None: @@ -554,6 +739,7 @@ def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seq initial_state=None, output_final_state=False, cu_seqlens=_conv_cu_seqlens, + cp_context=_conv_cp_context, ) if _pad_n > 0: qkv = qkv[:, :_orig_seq, :] @@ -562,14 +748,18 @@ def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cp_size, cp_group, cu_seq # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len, cp_size + qkv, gate, beta, alpha, kernel_batch, kernel_seq_len, cp_size_headwise=cp_size_headwise ) nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") # Calculate g and beta nvtx_range_push(suffix="g_and_beta") - A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group) - dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group) + A_log_local_cp = get_parameter_local_cp_headwise( + self.A_log, dim=0, cp_group=cp_group_headwise + ) + dt_bias_local_cp = get_parameter_local_cp_headwise( + self.dt_bias, dim=0, cp_group=cp_group_headwise + ) g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") @@ -689,14 +879,18 @@ def _apply_gated_norm(self, x, gate): return y @jit_fuser - def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len, cp_size): + def _prepare_qkv_for_gated_delta_rule( + self, qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise + ): """ Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. """ # Split qkv into query_key and value query_key, value = torch.split( - qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 + qkv, + [2 * self.qk_dim_local_tp // cp_size_headwise, self.v_dim_local_tp // cp_size_headwise], + dim=-1, ) # Reshape query_key and value @@ -708,7 +902,8 @@ def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_l query_key = l2norm(query_key.contiguous()) # Split query and key - query, key = query_key.chunk(2, dim=2) + split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size_headwise + query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) if self.num_value_heads // self.num_key_heads > 1: @@ -984,7 +1179,7 @@ def sh_ten_build_fn( #################### # Context parallel utilities #################### -def get_parameter_local_cp( +def get_parameter_local_cp_headwise( param: torch.Tensor, dim: int, cp_group: torch.distributed.ProcessGroup, @@ -1005,19 +1200,20 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ - cp_size = cp_group.size() - cp_rank = cp_group.rank() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to split if CP size is 1. if cp_size == 1: return param + cp_rank = cp_group.rank() + # Split first if needed. if split_sections is not None: inputs = torch.split(param, split_sections, dim=dim) outputs = [] for p in inputs: - p = get_parameter_local_cp(p, dim, cp_group) + p = get_parameter_local_cp_headwise(p, dim, cp_group) outputs.append(p) return torch.cat(outputs, dim=dim) @@ -1039,6 +1235,8 @@ def tensor_a2a_cp2hp( ): """All-to-all context parallel to hidden parallel. + This communication primitive is used by GDN headwise CP mode. + Args: tensor (torch.Tensor): The tensor to all-to-all. Currently only support (seq_len, batch, head_dim) shaped tensor. @@ -1054,7 +1252,7 @@ def tensor_a2a_cp2hp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: @@ -1102,6 +1300,8 @@ def tensor_a2a_hp2cp( ): """All-to-all hidden parallel to context parallel. + This communication primitive is used by GDN headwise CP mode. + Args: tensor (torch.Tensor): The tensor to all-to-all. Currently only support (seq_len, batch, head_dim) shaped tensor. @@ -1117,7 +1317,7 @@ def tensor_a2a_hp2cp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: @@ -1170,6 +1370,7 @@ def torch_chunk_gated_delta_rule( output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, + cp_context=None, ): # pylint: disable=line-too-long ''' @@ -1182,6 +1383,9 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + assert ( + cp_context is None + ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." initial_dtype = query.dtype if use_qk_l2norm_in_kernel: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 8f01db519de..61433e7c418 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -379,6 +379,11 @@ class TransformerConfig(ModelParallelConfig): gdn_pre_gated_delta_rule_fusion: bool = False """Whether to use the streamed Triton fusion for GatedDeltaNet pre-GDR preprocessing.""" + gdn_conv_pad_alignment: Optional[int] = None + """When set, pad packed GDN causal-conv inputs to this token alignment. + This is only valid without chunkwise CP: padding a chunk-local causal-conv input changes + the sequence seen by later chunks and therefore changes the GDN recurrence numerics.""" + #################### # initialization #################### @@ -978,6 +983,8 @@ class TransformerConfig(ModelParallelConfig): str: all layers share same communication type. List[str]: each layer has its separate communication type. cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". + This option controls standard attention layers. Linear-attention layers use + `linear_cp_mode` instead. "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be overlapped with attention compute. "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not @@ -989,6 +996,19 @@ class TransformerConfig(ModelParallelConfig): and P2P communications in high-level CP groups (e.g., via IBLink). """ + linear_cp_mode: Optional[str] = "chunkwise" + """Context-parallel execution mode for linear-attention layers + (e.g. Gated Delta Net). Independent of `cp_comm_type`, which only controls standard attention. + Can be "chunkwise" or "headwise": + "chunkwise": Keep sequence chunks sharded across CP ranks and use CP-aware linear kernels + (e.g. chunk_gated_delta_rule + causal_conv1d with a CP context). This follows the chunkwise + DeltaNet idea of storing state at chunk boundaries and doing chunk-local matrix work, avoiding + a full per-token recurrent state materialization while keeping tensor-core-friendly matmuls. + See https://sustcsonglin.github.io/blog/2024/deltanet-2/#a-chunkwise-algorithm-for-deltanet. + "headwise": Scatter heads across the CP group with all-to-all (Ulysses-style); each rank runs + the linear-attention kernel on the full sequence for a shard of heads. Correct but memory-heavy. + """ + ################## # Cuda Graphs ################## @@ -1489,16 +1509,36 @@ def __post_init__(self): f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " f"linear_num_key_heads ({self.linear_num_key_heads})." ) + if self.gdn_conv_pad_alignment is not None: + assert self.gdn_conv_pad_alignment > 0, ( + f"gdn_conv_pad_alignment must be positive when set, " + f"got {self.gdn_conv_pad_alignment}." + ) - # Check tensor parallelism compatibility - tp_cp_size = self.tensor_model_parallel_size * self.context_parallel_size - assert self.linear_num_key_heads % tp_cp_size == 0, ( + if self.context_parallel_size > 1: + assert self.linear_cp_mode in ("headwise", "chunkwise"), ( + f"linear_cp_mode must be one of 'headwise' or 'chunkwise', " + f"got {self.linear_cp_mode!r}." + ) + if self.gdn_conv_pad_alignment is not None: + assert self.linear_cp_mode != "chunkwise", ( + "gdn_conv_pad_alignment is incompatible with " + "linear_cp_mode='chunkwise' when context_parallel_size > 1. " + "Padding chunk-local GDN causal-conv inputs can change later " + "chunk numerics." + ) + # Check tensor parallelism compatibility. Headwise CP splits linear-attention heads + # across CP ranks; chunkwise CP keeps all TP-local heads on each CP rank. + linear_head_parallel_size = self.tensor_model_parallel_size + if self.context_parallel_size > 1 and self.linear_cp_mode == "headwise": + linear_head_parallel_size *= self.context_parallel_size + assert self.linear_num_key_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_key_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) - assert self.linear_num_value_heads % tp_cp_size == 0, ( + assert self.linear_num_value_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_value_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) elif self.experimental_attention_variant == "dsa": pass diff --git a/pyproject.toml b/pyproject.toml index 95a965e959d..444da6253fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,7 @@ dev = [ "opentelemetry-api~=1.33.1", "mamba-ssm~=2.2", "causal-conv1d~=1.5", - "flash-linear-attention~=0.4.0", + "flash-linear-attention>=0.4.2,<0.5", "megatron-energon[av_decode]~=6.0", "av", "flashinfer-python>=0.5.0,<0.7.0", diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index acabd1ceccd..08c98dde4b1 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -159,7 +159,9 @@ "linear_key_head_dim": 128, "linear_num_key_heads": 16, "linear_num_value_heads": 32, + "gdn_conv_pad_alignment": None, "gdn_pre_gated_delta_rule_fusion": False, + "linear_cp_mode": "chunkwise", "linear_value_head_dim": 128, "log_max_attention_logit": False, "log_moe_overload_factor": False, diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index da209735e39..8c09a9c927b 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -119,16 +119,58 @@ def test_gdn_pre_gated_delta_rule_fusion_requires_gdn_variant(): ) +def test_gdn_conv_pad_alignment_rejects_chunkwise_cp(): + with pytest.raises(AssertionError, match="gdn_conv_pad_alignment is incompatible"): + _make_gdn_config( + context_parallel_size=2, linear_cp_mode="chunkwise", gdn_conv_pad_alignment=4096 + ) + + +def test_gdn_chunkwise_cp_head_divisibility_ignores_cp_size(): + config = _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="chunkwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + assert config.linear_cp_mode == "chunkwise" + + +def test_gdn_headwise_cp_head_divisibility_includes_cp_size(): + with pytest.raises(AssertionError, match="linear_head_parallel_size"): + _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="headwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + + @pytest.mark.parametrize( - ("tp_size", "sp", "cp_size"), - [(1, False, 1), (2, False, 1), (2, True, 1), (1, False, 2), (2, False, 2), (2, True, 2)], + ("tp_size", "sp", "cp_size", "linear_cp_mode"), + [ + # cp_size=1: the CP path is inactive, so linear_cp_mode choice is irrelevant. + # Cover the "chunkwise" default and skip the "headwise" variants for brevity. + (1, False, 1, None), + (2, False, 1, None), + (2, True, 1, None), + # cp_size=2: exercise both CP paths. + (1, False, 2, "headwise"), + (2, False, 2, "headwise"), + (2, True, 2, "headwise"), + (1, False, 2, "chunkwise"), + (2, False, 2, "chunkwise"), + (2, True, 2, "chunkwise"), + ], ) @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.internal class TestGatedDeltaNet: @pytest.fixture(scope='function', autouse=True) - def setup_method(self, tp_size, sp, cp_size): + def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): # Initialize parallel and random seed Utils.initialize_model_parallel( tensor_model_parallel_size=tp_size, @@ -139,6 +181,18 @@ def setup_method(self, tp_size, sp, cp_size): self.tp_size = tp_size self.cp_size = cp_size self.sp_size = tp_size if sp else 1 + self.linear_cp_mode = linear_cp_mode + if self.linear_cp_mode == "headwise": + self.cp_size_chunkwise = 1 + self.cp_size_headwise = self.cp_size + elif self.linear_cp_mode == "chunkwise": + self.cp_size_chunkwise = self.cp_size + self.cp_size_headwise = 1 + elif self.cp_size == 1: + self.cp_size_chunkwise = 1 + self.cp_size_headwise = 1 + else: + raise ValueError(f"Invalid linear CP mode: {self.linear_cp_mode}") # Get TP and CP process groups from device mesh tp_group = parallel_state.get_tensor_model_parallel_group() @@ -166,6 +220,7 @@ def setup_method(self, tp_size, sp, cp_size): context_parallel_size=cp_size, experimental_attention_variant="gated_delta_net", linear_attention_freq=[1], + linear_cp_mode=self.linear_cp_mode, transformer_impl="transformer_engine", ) gdn_submodules = get_experimental_attention_variant_module_spec( @@ -191,7 +246,7 @@ def teardown_method(self): def test_gpu_forward(self): gdn = self.gdn - micro_batch_size = 2 + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 seq_length = 64 hidden_states = torch.ones( (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), @@ -229,7 +284,7 @@ def test_selective_recompute_gdn(self): gdn = self.gdn gdn.train() - micro_batch_size = 2 + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 seq_length = 64 torch.manual_seed(1234) base_input = torch.randn( @@ -270,6 +325,43 @@ def run(recompute): msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", ) + def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): + if not (self.linear_cp_mode == "chunkwise" and self.cp_size > 1): + pytest.skip("Only chunkwise CP with CP>1 uses the FLA CP batch guard.") + + gdn = self.gdn + + micro_batch_size = 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with pytest.raises(ValueError, match="requires micro_batch_size == 1"): + gdn(hidden_states, None) + + def test_gpu_forward_rejects_sbhd_conv_padding(self): + gdn = self.gdn + gdn.config.gdn_conv_pad_alignment = 4096 + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + expected_error = ( + "incompatible with GDN chunkwise CP" + if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 + else "only supported with packed sequence" + ) + with pytest.raises(ValueError, match=expected_error): + gdn(hidden_states, None) + def test_jit_compiled_helpers(self): import torch._dynamo @@ -277,9 +369,9 @@ def test_jit_compiled_helpers(self): batch = 2 seq_len = 16 - num_v_heads_local = gdn.num_value_heads // gdn.tp_size // gdn.cp_size + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // self.cp_size_headwise - qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size + qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // self.cp_size_headwise qkv = torch.randn( batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 ) @@ -311,7 +403,7 @@ def test_jit_compiled_helpers(self): with torch._dynamo.config.patch(disable=True): query, key, value, gate_out, beta_out, alpha_out = ( gdn._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len, gdn.cp_size + qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise=self.cp_size_headwise ) ) @@ -339,6 +431,8 @@ def test_jit_compiled_helpers(self): def test_gpu_forward_thd_correctness(self): if self.sp_size > 1: pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") atol, rtol = 3e-4, 3e-4 @@ -382,6 +476,8 @@ def test_gpu_forward_thd_correctness(self): def test_gpu_forward_thd_padding_correctness(self): if self.sp_size > 1: pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") atol, rtol = 3e-4, 3e-4 sequence_length = 32 @@ -424,14 +520,30 @@ def test_gpu_forward_thd_padding_correctness(self): ) assert output_thd_no_padding.shape == output_thd_padded.shape - # C) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + # C) explicit causal-conv padding is only applied to packed inputs and + # should not affect the original unpadded token outputs. + self.gdn.config.gdn_conv_pad_alignment = 48 + output_thd_conv_pad, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + self.gdn.config.gdn_conv_pad_alignment = None + assert output_thd_conv_pad.shape == output_thd_no_padding.shape + torch.testing.assert_close( + output_thd_conv_pad, + output_thd_no_padding, + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD conv-padded output mismatch ({rank=}): {msg}", + ) + + # D) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. 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] ) with pytest.raises(ValueError, match="does not match"): self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) - # D) actual mismatch branch without *_padded: should raise. + # E) actual mismatch branch without *_padded: should raise. actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) with pytest.raises(ValueError, match="does not match"): self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) @@ -516,8 +628,8 @@ def _packed_pre_gated_delta_rule_reference(self, gdn, qkvzba, cu_seqlens): qkvzba[start:end], batch=1, seq_len=end - start, - cp_size=gdn.cp_size, - cp_group=gdn.pg_collection.cp, + cp_size_headwise=gdn.cp_size, + cp_group_headwise=gdn.pg_collection.cp, ) for output_list, output in zip(segment_outputs, outputs): output_list.append(output) @@ -971,17 +1083,22 @@ def test_cp1_still_validates_total(self, mock_gdn): @pytest.mark.parametrize("sequence_packing", [False, True]) @pytest.mark.parametrize( - ("tp", "sp", "cp"), + ("tp", "sp", "cp", "linear_cp_mode"), [ - (4, False, 1), # TP w/o SP - (4, True, 1), # TP w/ SP - (1, False, 2), # CP - (2, False, 2), # TP w/o SP + CP - (2, True, 2), # TP w/ SP + CP + (4, False, 1, None), # TP w/o SP + (4, True, 1, None), # TP w/ SP + (1, False, 2, "headwise"), # Headwise CP + (2, False, 2, "headwise"), # TP w/o SP + Headwise CP + (2, True, 2, "headwise"), # TP w/ SP + Headwise CP + (1, False, 2, "chunkwise"), # Chunkwise CP + (2, False, 2, "chunkwise"), # TP w/o SP + chunkwise CP + (2, True, 2, "chunkwise"), # TP w/ SP + chunkwise CP ], ) @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") -def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packing, tp, sp, cp): +def test_parallel_gated_delta_net_correctness( + tmp_path_dist_ckpt, sequence_packing, tp, sp, cp, linear_cp_mode +): transformer_config = TransformerConfig( hidden_size=128, linear_conv_kernel_dim=2, @@ -998,6 +1115,7 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi bf16=True, experimental_attention_variant="gated_delta_net", linear_attention_freq=[1], + linear_cp_mode=linear_cp_mode, transformer_impl="transformer_engine", ) @@ -1005,10 +1123,16 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi config=transformer_config, vp_stage=None, pp_rank=0 ) - if cp: - atol, rtol = 5e-3, 5e-3 + cosine_similarity_threshold = None + if cp > 1: + atol, rtol = 2e-3, 1e-2 + cosine_similarity_threshold = 0.9999 else: - atol, rtol = 5e-4, 5e-4 + atol, rtol = 2e-4, 2e-3 + cosine_similarity_threshold = 0.99999 + + is_chunkwise_cp = linear_cp_mode == "chunkwise" and cp > 1 + micro_batch_size = 1 if is_chunkwise_cp and not sequence_packing else 4 _test_parallel_attention_correctness( transformer_config=transformer_config, @@ -1016,12 +1140,13 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi tmp_path_dist_ckpt=tmp_path_dist_ckpt, atol=atol, rtol=rtol, + cosine_similarity_threshold=cosine_similarity_threshold, tp=tp, sp=sp, cp=cp, seed=123, sequence_length=256, - micro_batch_size=4, + micro_batch_size=micro_batch_size, sequence_packing=sequence_packing, ) diff --git a/tests/unit_tests/test_context_parallel_layout.py b/tests/unit_tests/test_context_parallel_layout.py new file mode 100644 index 00000000000..b762e4594a1 --- /dev/null +++ b/tests/unit_tests/test_context_parallel_layout.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.context_parallel_layout import get_thd_context_parallel_rank_indices + + +def _token_ranges(*spans): + return [token for start, end in spans for token in range(start, end)] + + +def test_thd_context_parallel_rank_indices_match_per_sequence_chunk_order(): + cu_seqlens = torch.tensor([0, 16, 40]) + + 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) + ) + + +@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 + + rank_indices = [ + get_thd_context_parallel_rank_indices(cu_seqlens, cp_size, rank, layout) + for rank in range(cp_size) + ] + + assert [indices.numel() for indices in rank_indices] == [32, 32, 32, 32] + assert torch.cat(rank_indices).sort().values.tolist() == list(range(128)) + + +@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]) + + 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), + ) + + +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") + + +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_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") diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index 4f2c96e9ef9..6a65a91a2fc 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -706,6 +706,8 @@ def _test_parallel_attention_correctness( tmp_path_dist_ckpt, atol, rtol, + cosine_similarity_threshold=None, + relative_l2_threshold=0.1, tp=1, sp=False, cp=1, @@ -862,27 +864,45 @@ def get_tensor_on_this_rank(tensor): ~torch.isinf(bias_hidden_states_parallel) ), "bias_hidden_states_parallel contains inf" - torch.testing.assert_close( - output_hidden_states_baseline, - output_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in output_hidden_states: {msg}", - ) - torch.testing.assert_close( - input_grad_baseline, - input_grad_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in input_grad: {msg}", + def assert_close_or_cosine_similarity(baseline, parallel, tensor_name): + try: + torch.testing.assert_close( + baseline, + parallel, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Mismatch in {tensor_name}: {msg}", + ) + return + except AssertionError as close_error: + if cosine_similarity_threshold is None: + raise close_error + baseline_flat = baseline.flatten().float() + parallel_flat = parallel.flatten().float() + cosine_sim = torch.nn.functional.cosine_similarity( + baseline_flat.unsqueeze(0), parallel_flat.unsqueeze(0) + ).item() + diff_norm = torch.linalg.vector_norm((parallel_flat - baseline_flat).float()) + baseline_norm = torch.linalg.vector_norm(baseline_flat.float()).clamp_min(1e-12) + relative_l2 = (diff_norm / baseline_norm).item() + assert cosine_sim >= cosine_similarity_threshold, ( + f"Mismatch in {tensor_name}: cosine similarity " + f"{cosine_sim} < {cosine_similarity_threshold}, " + f"while assert_close failed: {close_error}" + ) + assert relative_l2 <= relative_l2_threshold, ( + f"Mismatch in {tensor_name}: relative L2 " + f"{relative_l2} > {relative_l2_threshold}, " + f"while assert_close failed: {close_error}" + ) + + assert_close_or_cosine_similarity( + output_hidden_states_baseline, output_hidden_states_parallel, "output_hidden_states" ) + assert_close_or_cosine_similarity(input_grad_baseline, input_grad_parallel, "input_grad") if has_bias: - torch.testing.assert_close( - bias_hidden_states_baseline, - bias_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in bias_hidden_states: {msg}", + assert_close_or_cosine_similarity( + bias_hidden_states_baseline, bias_hidden_states_parallel, "bias_hidden_states" ) Utils.destroy_model_parallel() diff --git a/uv.lock b/uv.lock index 062d09d4217..ba195557c34 100644 --- a/uv.lock +++ b/uv.lock @@ -3309,7 +3309,7 @@ requires-dist = [ { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'lts'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "fastapi", marker = "extra == 'lts'", specifier = "~=0.50" }, - { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "~=0.4.0" }, + { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = ">=0.4.2,<0.5" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, { name = "flashinfer-python", marker = "extra == 'lts'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" },