From 8e7b82977b74e6ce35e6e0c12ae891b9739503ba Mon Sep 17 00:00:00 2001 From: guapisolo Date: Mon, 6 Apr 2026 23:27:41 +0000 Subject: [PATCH 1/3] merge common utils and directly pass cu_seqlens instead of caluclate --- miles/backends/megatron_utils/actor.py | 4 +- miles/backends/training_utils/cp_utils.py | 39 ++++++++++++++++++- miles_plugins/models/qwen3_5.py | 21 +--------- miles_plugins/models/qwen3_next.py | 21 +--------- .../fast/backends/training_utils/__init__.py | 1 + 5 files changed, 45 insertions(+), 41 deletions(-) create mode 100644 tests/fast/backends/training_utils/__init__.py diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index f3dbef0ac85..5e5373079e6 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -112,12 +112,12 @@ def init( from megatron.core import mpu - if mpu.get_context_parallel_world_size() > 1: + cp_world_size = mpu.get_context_parallel_world_size() + if cp_world_size > 1: from miles.backends.training_utils.cp_utils import setup_hybrid_cp cp_group = mpu.get_context_parallel_group() cp_rank = mpu.get_context_parallel_rank() - cp_world_size = mpu.get_context_parallel_world_size() for model_chunk in self.model: setup_hybrid_cp(model_chunk, cp_group, cp_rank, cp_world_size) diff --git a/miles/backends/training_utils/cp_utils.py b/miles/backends/training_utils/cp_utils.py index 3cfee2d74f6..b0da338eda4 100644 --- a/miles/backends/training_utils/cp_utils.py +++ b/miles/backends/training_utils/cp_utils.py @@ -8,6 +8,11 @@ from .parallel import get_parallel_state +try: + from fla.ops.cp import build_cp_context as _fla_build_cp_context +except ImportError: + _fla_build_cp_context = None + logger = logging.getLogger(__name__) @@ -342,14 +347,46 @@ def slice_log_prob_with_cp( return torch.cat([chunk_1, chunk_2], dim=0) +def build_gdn_cp_context(module: nn.Module, cu_seqlens: torch.Tensor, device: torch.device): + """Build fla CP context for a GatedDeltaNet module from packed sequence boundaries. + + Args: + module: GDN module with ``cp_group`` / ``cp_world_size`` / ``conv_kernel_size``. + cu_seqlens: Global packed sequence boundaries (e.g. ``packed_seq_params.cu_seqlens_q``). + device: Target device. + + Returns ``None`` when CP is not configured on the module (``cp_group`` not set). + Raises ``RuntimeError`` if hybrid CP is configured but ``fla.ops.cp`` is missing. + """ + cp_group = getattr(module, "cp_group", None) + if cp_group is None: + return None + if _fla_build_cp_context is None: + raise RuntimeError( + "Hybrid CP requires fla.ops.cp (flash-linear-attention >= 0.4.2) " "but it could not be imported." + ) + if cu_seqlens is None or cu_seqlens.numel() < 2: + raise ValueError(f"Hybrid CP requires valid cu_seqlens (at least 2 elements) but got {cu_seqlens}") + return _fla_build_cp_context( + cu_seqlens=cu_seqlens.to(device=device, dtype=torch.int32), + group=cp_group, + conv1d_kernel_size=module.conv_kernel_size, + ) + + def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: """Configure GatedDeltaNet modules for native fla CP instead of all-gather duplication. Walks the model tree looking for HuggingfaceAttention submodules that have a ``linear_attn`` child (i.e. DeltaNet layers). For each one it sets the CP - metadata so that ``_build_cp_context`` produces a valid context, and flips + metadata so that ``build_gdn_cp_context`` produces a valid context, and flips ``hybrid_cp`` so the parent skips the all-gather path. """ + if _fla_build_cp_context is None: + raise RuntimeError( + "setup_hybrid_cp requires fla.ops.cp (flash-linear-attention >= 0.4.2) " + "but it could not be imported. Cannot enable hybrid CP without the fla CP backend." + ) from miles_plugins.models.hf_attention import HuggingfaceAttention count = 0 diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index 075ceec5078..794cf738081 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -15,11 +15,7 @@ except ImportError: pass -try: - from fla.ops.cp import FLACPContext, build_cp_context -except ImportError: - FLACPContext = None - build_cp_context = None +from miles.backends.training_utils.cp_utils import build_gdn_cp_context from .hf_attention import HuggingfaceAttention, _load_hf_config @@ -87,19 +83,6 @@ def __init__(self, config, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - def _build_cp_context(self, local_seq_len: int, device: torch.device): - """Build fla CP context from the local (sharded) sequence length.""" - cp_group = getattr(self, "cp_group", None) - if cp_group is None or build_cp_context is None: - return None - global_seq_len = local_seq_len * self.cp_world_size - global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device) - return build_cp_context( - cu_seqlens=global_cu_seqlens, - group=cp_group, - conv1d_kernel_size=self.conv_kernel_size, - ) - def forward( self, hidden_states: torch.Tensor, @@ -107,7 +90,7 @@ def forward( ): batch_size, seq_len, _ = hidden_states.shape - cp_context = self._build_cp_context(seq_len, hidden_states.device) + cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) # Projections (flat layout: [Q_all, K_all, V_all]) mixed_qkv = self.in_proj_qkv(hidden_states) diff --git a/miles_plugins/models/qwen3_next.py b/miles_plugins/models/qwen3_next.py index 71cae337653..1dbee8acd01 100644 --- a/miles_plugins/models/qwen3_next.py +++ b/miles_plugins/models/qwen3_next.py @@ -18,11 +18,7 @@ except ImportError: pass -try: - from fla.ops.cp import FLACPContext, build_cp_context -except ImportError: - FLACPContext = None - build_cp_context = None +from miles.backends.training_utils.cp_utils import build_gdn_cp_context from .hf_attention import HuggingfaceAttention @@ -80,19 +76,6 @@ def __init__(self, config, layer_idx: int): self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - def _build_cp_context(self, local_seq_len: int, device: torch.device): - """Build fla CP context from the local (sharded) sequence length.""" - cp_group = getattr(self, "cp_group", None) - if cp_group is None or build_cp_context is None: - return None - global_seq_len = local_seq_len * self.cp_world_size - global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device) - return build_cp_context( - cu_seqlens=global_cu_seqlens, - group=cp_group, - conv1d_kernel_size=self.conv_kernel_size, - ) - def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): """ Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`. @@ -127,7 +110,7 @@ def forward( hidden_states: torch.Tensor, cu_seqlens: torch.Tensor = None, ): - cp_context = self._build_cp_context(hidden_states.shape[1], hidden_states.device) + cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) projected_states_qkvz = self.in_proj_qkvz(hidden_states) projected_states_ba = self.in_proj_ba(hidden_states) diff --git a/tests/fast/backends/training_utils/__init__.py b/tests/fast/backends/training_utils/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/fast/backends/training_utils/__init__.py @@ -0,0 +1 @@ + From 6ef2aec2914184de88c8146629d44bfadbab774e Mon Sep 17 00:00:00 2001 From: guapisolo Date: Tue, 7 Apr 2026 21:17:27 +0000 Subject: [PATCH 2/3] try fix qwen35 --- miles/backends/megatron_utils/actor.py | 4 +- miles/backends/training_utils/cp_utils.py | 19 +--- miles_plugins/models/hf_attention.py | 121 +++++++++++++++++++++- 3 files changed, 126 insertions(+), 18 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 5e5373079e6..993c4952c46 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -114,12 +114,12 @@ def init( cp_world_size = mpu.get_context_parallel_world_size() if cp_world_size > 1: - from miles.backends.training_utils.cp_utils import setup_hybrid_cp + from miles.backends.training_utils.cp_utils import detect_and_setup_hybrid_cp cp_group = mpu.get_context_parallel_group() cp_rank = mpu.get_context_parallel_rank() for model_chunk in self.model: - setup_hybrid_cp(model_chunk, cp_group, cp_rank, cp_world_size) + detect_and_setup_hybrid_cp(model_chunk, cp_group, cp_rank, cp_world_size) verify_megatron_parallel_state(self.model) diff --git a/miles/backends/training_utils/cp_utils.py b/miles/backends/training_utils/cp_utils.py index b0da338eda4..12088552903 100644 --- a/miles/backends/training_utils/cp_utils.py +++ b/miles/backends/training_utils/cp_utils.py @@ -374,19 +374,10 @@ def build_gdn_cp_context(module: nn.Module, cu_seqlens: torch.Tensor, device: to ) -def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: - """Configure GatedDeltaNet modules for native fla CP instead of all-gather duplication. - - Walks the model tree looking for HuggingfaceAttention submodules that have a - ``linear_attn`` child (i.e. DeltaNet layers). For each one it sets the CP - metadata so that ``build_gdn_cp_context`` produces a valid context, and flips - ``hybrid_cp`` so the parent skips the all-gather path. - """ - if _fla_build_cp_context is None: - raise RuntimeError( - "setup_hybrid_cp requires fla.ops.cp (flash-linear-attention >= 0.4.2) " - "but it could not be imported. Cannot enable hybrid CP without the fla CP backend." - ) +def detect_and_setup_hybrid_cp( + model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int +) -> None: + """Scan for GatedDeltaNet modules and configure them for native fla CP.""" from miles_plugins.models.hf_attention import HuggingfaceAttention count = 0 @@ -401,4 +392,4 @@ def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, count += 1 if count > 0: - logger.info(f"Configured hybrid CP on {count} DeltaNet modules (fla native state passing)") + logger.info(f"Configured hybrid CP on {count} GDN modules (fla native state passing)") diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index a5e63ae8b19..8b12cb78ea6 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -38,6 +38,105 @@ def _fix_dtype(d): return ns +class _ZigzagSequentialExchange(torch.autograd.Function): + """P2P exchange to convert between zigzag and sequential CP layouts. + + For CP=2, zigzag rank 0 holds [sub_0, sub_3] and rank 1 holds [sub_1, sub_2]. + Sequential rank 0 needs [sub_0, sub_1] and rank 1 needs [sub_2, sub_3]. + This exchanges the misplaced sub-chunk between ranks via a single sendrecv. + """ + + @staticmethod + def forward(ctx, send_buf, cp_group, cp_rank): + ctx.cp_group = cp_group + ctx.cp_rank = cp_rank + recv_buf = torch.empty_like(send_buf) + peer = 1 - cp_rank + if cp_rank == 0: + dist.send(send_buf.contiguous(), group_dst=peer, group=cp_group) + dist.recv(recv_buf, group_src=peer, group=cp_group) + else: + dist.recv(recv_buf, group_src=peer, group=cp_group) + dist.send(send_buf.contiguous(), group_dst=peer, group=cp_group) + return recv_buf + + @staticmethod + def backward(ctx, grad_recv): + grad_send = torch.empty_like(grad_recv) + peer = 1 - ctx.cp_rank + if ctx.cp_rank == 0: + dist.send(grad_recv.contiguous(), group_dst=peer, group=ctx.cp_group) + dist.recv(grad_send, group_src=peer, group=ctx.cp_group) + else: + dist.recv(grad_send, group_src=peer, group=ctx.cp_group) + dist.send(grad_recv.contiguous(), group_dst=peer, group=ctx.cp_group) + return grad_send, None, None + + +def _zigzag_to_sequential(hidden_states, local_cu_seqlens, cp_group, cp_rank): + """Convert zigzag CP layout to sequential via P2P exchange (CP=2 only). + + Rank 0 zigzag: [sub_0, sub_3] → sequential: [sub_0, sub_1] + Rank 1 zigzag: [sub_1, sub_2] → sequential: [sub_2, sub_3] + """ + # Split each sample into ascending (first half) and descending (second half) + keep_parts, send_parts = [], [] + for i in range(len(local_cu_seqlens) - 1): + start, end = local_cu_seqlens[i], local_cu_seqlens[i + 1] + mid = (start + end) // 2 + if cp_rank == 0: + keep_parts.append(hidden_states[start:mid]) # sub_0 + send_parts.append(hidden_states[mid:end]) # sub_3 → send to rank 1 + else: + send_parts.append(hidden_states[start:mid]) # sub_1 → send to rank 0 + keep_parts.append(hidden_states[mid:end]) # sub_2 + + send_buf = torch.cat(send_parts, dim=0) + recv_buf = _ZigzagSequentialExchange.apply(send_buf, cp_group, cp_rank) + + # Reassemble: both ranks → [keep, recv] + result = [] + offset = 0 + for i in range(len(local_cu_seqlens) - 1): + chunk_len = (local_cu_seqlens[i + 1] - local_cu_seqlens[i]) // 2 + result.append(keep_parts[i]) + result.append(recv_buf[offset : offset + chunk_len]) + offset += chunk_len + return torch.cat(result, dim=0) + + +def _sequential_to_zigzag(hidden_states, local_cu_seqlens, cp_group, cp_rank): + """Convert sequential CP layout back to zigzag via P2P exchange (CP=2 only). + + Rank 0 sequential: [sub_0, sub_1] → zigzag: [sub_0, sub_3] + Rank 1 sequential: [sub_2, sub_3] → zigzag: [sub_1, sub_2] + """ + keep_parts, send_parts = [], [] + for i in range(len(local_cu_seqlens) - 1): + start, end = local_cu_seqlens[i], local_cu_seqlens[i + 1] + mid = (start + end) // 2 + keep_parts.append(hidden_states[start:mid]) # first half + send_parts.append(hidden_states[mid:end]) # second half + + send_buf = torch.cat(send_parts, dim=0) + recv_buf = _ZigzagSequentialExchange.apply(send_buf, cp_group, cp_rank) + + # Rank 0: [keep (sub_0), recv (sub_3)] + # Rank 1: [recv (sub_1), keep (sub_2)] + result = [] + offset = 0 + for i in range(len(local_cu_seqlens) - 1): + chunk_len = (local_cu_seqlens[i + 1] - local_cu_seqlens[i]) // 2 + if cp_rank == 0: + result.append(keep_parts[i]) + result.append(recv_buf[offset : offset + chunk_len]) + else: + result.append(recv_buf[offset : offset + chunk_len]) + result.append(keep_parts[i]) + offset += chunk_len + return torch.cat(result, dim=0) + + class _AllGatherForDuplicatedComputation(torch.autograd.Function): """All-gather whose backward just returns the local gradient slice (no reduce). @@ -119,7 +218,17 @@ def forward( group=mpu.get_tensor_model_parallel_group(), ) - if mpu.get_context_parallel_world_size() > 1 and not self.hybrid_cp: + if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp: + cp_size = mpu.get_context_parallel_world_size() + local_cu_seqlens = cu_seqlens // cp_size + hidden_states = _zigzag_to_sequential( + hidden_states, + local_cu_seqlens, + mpu.get_context_parallel_group(), + mpu.get_context_parallel_rank(), + ) + + elif mpu.get_context_parallel_world_size() > 1: cp_size = mpu.get_context_parallel_world_size() # Use custom all-gather whose backward returns local gradient # instead of reduce-scatter, since the computation is duplicated. @@ -154,7 +263,15 @@ def forward( output = output.permute(1, 0, 2) # [seq_len, bsz, hidden_dim] - if mpu.get_context_parallel_world_size() > 1 and not self.hybrid_cp: + if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp: + output = _sequential_to_zigzag( + output, + local_cu_seqlens, + mpu.get_context_parallel_group(), + mpu.get_context_parallel_rank(), + ) + + elif mpu.get_context_parallel_world_size() > 1: cp_rank = mpu.get_context_parallel_rank() output_list = [] for i in range(len(cu_seqlens) - 1): From 77c0c6ada812cf534574ed47ad364a64e1ec241b Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 7 Apr 2026 23:11:40 +0000 Subject: [PATCH 3/3] feat: generalize zigzag-sequential exchange for any CP size - Replace CP=2-only P2P exchange with general implementation - Use batch_isend_irecv with group_peer for correct group-local ranks - Support arbitrary CP sizes via sub-chunk routing table - Update correctness test to pass global cu_seqlens --- miles_plugins/models/hf_attention.py | 221 +++++++++++------- .../precision/test_qwen3_5_cp_correctness.py | 6 +- 2 files changed, 143 insertions(+), 84 deletions(-) diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index 8b12cb78ea6..02593e9cac2 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -38,103 +38,160 @@ def _fix_dtype(d): return ns -class _ZigzagSequentialExchange(torch.autograd.Function): - """P2P exchange to convert between zigzag and sequential CP layouts. +def _sub_chunk_location(sub_id, cp_size): + """Return (zigzag_rank, half_index) for a given sub-chunk id. - For CP=2, zigzag rank 0 holds [sub_0, sub_3] and rank 1 holds [sub_1, sub_2]. - Sequential rank 0 needs [sub_0, sub_1] and rank 1 needs [sub_2, sub_3]. - This exchanges the misplaced sub-chunk between ranks via a single sendrecv. + In zigzag layout with N=cp_size ranks and 2N sub-chunks: + rank k holds [sub_k, sub_{2N-1-k}] + So sub_x lives on rank x (half 0) if x < N, else rank 2N-1-x (half 1). + """ + if sub_id < cp_size: + return sub_id, 0 + return 2 * cp_size - 1 - sub_id, 1 + + +def _p2p_exchange(send_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group): + """Exchange multiple buffers via batched async P2P, handling self-sends.""" + # Handle self-sends as local copies + for i, dst in enumerate(send_dsts): + if dst == cp_rank: + for j, src in enumerate(recv_srcs): + if src == cp_rank and recv_bufs[j] is not send_bufs[i]: + recv_bufs[j].copy_(send_bufs[i]) + + # Build P2P ops for remote exchanges (use group_peer for group-local ranks) + p2p_ops = [] + for j, src in enumerate(recv_srcs): + if src != cp_rank: + p2p_ops.append(dist.P2POp(dist.irecv, recv_bufs[j], group_peer=src, group=cp_group)) + for i, dst in enumerate(send_dsts): + if dst != cp_rank: + p2p_ops.append(dist.P2POp(dist.isend, send_bufs[i].contiguous(), group_peer=dst, group=cp_group)) + + if p2p_ops: + reqs = dist.batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + +class _ZigzagToSequential(torch.autograd.Function): + """Convert zigzag CP layout to sequential layout for any CP size. + + Zigzag rank k holds: [sub_k, sub_{2N-1-k}] + Sequential rank j needs: [sub_{2j}, sub_{2j+1}] """ @staticmethod - def forward(ctx, send_buf, cp_group, cp_rank): + def forward(ctx, hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size): ctx.cp_group = cp_group ctx.cp_rank = cp_rank - recv_buf = torch.empty_like(send_buf) - peer = 1 - cp_rank - if cp_rank == 0: - dist.send(send_buf.contiguous(), group_dst=peer, group=cp_group) - dist.recv(recv_buf, group_src=peer, group=cp_group) - else: - dist.recv(recv_buf, group_src=peer, group=cp_group) - dist.send(send_buf.contiguous(), group_dst=peer, group=cp_group) - return recv_buf + ctx.cp_size = cp_size + ctx.save_for_backward(local_cu_seqlens) + + N = cp_size + + # Split local data into first_half (sub_k) and second_half (sub_{2N-1-k}) + first_halves, second_halves = [], [] + for i in range(len(local_cu_seqlens) - 1): + start, end = local_cu_seqlens[i].item(), local_cu_seqlens[i + 1].item() + mid = (start + end) // 2 + first_halves.append(hidden_states[start:mid]) + second_halves.append(hidden_states[mid:end]) + + my_bufs = [torch.cat(first_halves, dim=0), torch.cat(second_halves, dim=0)] + my_sub_ids = [cp_rank, 2 * N - 1 - cp_rank] + send_dsts = [sid // 2 for sid in my_sub_ids] + + # What sequential rank cp_rank needs: sub_{2*cp_rank} and sub_{2*cp_rank+1} + need_ids = [2 * cp_rank, 2 * cp_rank + 1] + recv_srcs = [_sub_chunk_location(x, N)[0] for x in need_ids] + recv_bufs = [torch.empty_like(my_bufs[0]) for _ in range(2)] + + # Handle self-send: if I send to myself, point recv_buf to send_buf + for i, dst in enumerate(send_dsts): + if dst == cp_rank: + for j, src in enumerate(recv_srcs): + if src == cp_rank: + recv_bufs[j] = my_bufs[i] + + _p2p_exchange(my_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group) + + return torch.cat(recv_bufs, dim=0) @staticmethod - def backward(ctx, grad_recv): - grad_send = torch.empty_like(grad_recv) - peer = 1 - ctx.cp_rank - if ctx.cp_rank == 0: - dist.send(grad_recv.contiguous(), group_dst=peer, group=ctx.cp_group) - dist.recv(grad_send, group_src=peer, group=ctx.cp_group) - else: - dist.recv(grad_send, group_src=peer, group=ctx.cp_group) - dist.send(grad_recv.contiguous(), group_dst=peer, group=ctx.cp_group) - return grad_send, None, None - - -def _zigzag_to_sequential(hidden_states, local_cu_seqlens, cp_group, cp_rank): - """Convert zigzag CP layout to sequential via P2P exchange (CP=2 only). - - Rank 0 zigzag: [sub_0, sub_3] → sequential: [sub_0, sub_1] - Rank 1 zigzag: [sub_1, sub_2] → sequential: [sub_2, sub_3] - """ - # Split each sample into ascending (first half) and descending (second half) - keep_parts, send_parts = [], [] - for i in range(len(local_cu_seqlens) - 1): - start, end = local_cu_seqlens[i], local_cu_seqlens[i + 1] - mid = (start + end) // 2 - if cp_rank == 0: - keep_parts.append(hidden_states[start:mid]) # sub_0 - send_parts.append(hidden_states[mid:end]) # sub_3 → send to rank 1 - else: - send_parts.append(hidden_states[start:mid]) # sub_1 → send to rank 0 - keep_parts.append(hidden_states[mid:end]) # sub_2 - - send_buf = torch.cat(send_parts, dim=0) - recv_buf = _ZigzagSequentialExchange.apply(send_buf, cp_group, cp_rank) - - # Reassemble: both ranks → [keep, recv] + def backward(ctx, grad_output): + (local_cu_seqlens,) = ctx.saved_tensors + # Backward: sequential → zigzag (inverse permutation) + result = _sequential_to_zigzag_impl( + grad_output, local_cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size + ) + return result, None, None, None, None + + +def _sequential_to_zigzag_impl(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size): + """Core implementation for sequential → zigzag conversion.""" + N = cp_size + half_len = hidden_states.shape[0] // 2 + + seq_bufs = [hidden_states[:half_len], hidden_states[half_len:]] + my_seq_sub_ids = [2 * cp_rank, 2 * cp_rank + 1] + send_dsts = [_sub_chunk_location(x, N)[0] for x in my_seq_sub_ids] + + # Zigzag rank cp_rank needs sub_{cp_rank} and sub_{2N-1-cp_rank} + need_ids = [cp_rank, 2 * N - 1 - cp_rank] + recv_srcs = [nid // 2 for nid in need_ids] + recv_bufs = [torch.empty_like(seq_bufs[0]) for _ in range(2)] + + for i, dst in enumerate(send_dsts): + if dst == cp_rank: + for j, src in enumerate(recv_srcs): + if src == cp_rank: + recv_bufs[j] = seq_bufs[i] + + _p2p_exchange(seq_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group) + + # Reassemble zigzag: [first_half (sub_k), second_half (sub_{2N-1-k})] result = [] - offset = 0 + half_chunk = half_len // max(len(local_cu_seqlens) - 1, 1) + offset_0, offset_1 = 0, 0 for i in range(len(local_cu_seqlens) - 1): - chunk_len = (local_cu_seqlens[i + 1] - local_cu_seqlens[i]) // 2 - result.append(keep_parts[i]) - result.append(recv_buf[offset : offset + chunk_len]) - offset += chunk_len + chunk_len = (local_cu_seqlens[i + 1].item() - local_cu_seqlens[i].item()) // 2 + result.append(recv_bufs[0][offset_0 : offset_0 + chunk_len]) + result.append(recv_bufs[1][offset_1 : offset_1 + chunk_len]) + offset_0 += chunk_len + offset_1 += chunk_len return torch.cat(result, dim=0) -def _sequential_to_zigzag(hidden_states, local_cu_seqlens, cp_group, cp_rank): - """Convert sequential CP layout back to zigzag via P2P exchange (CP=2 only). +class _SequentialToZigzag(torch.autograd.Function): + """Convert sequential CP layout back to zigzag for any CP size.""" - Rank 0 sequential: [sub_0, sub_1] → zigzag: [sub_0, sub_3] - Rank 1 sequential: [sub_2, sub_3] → zigzag: [sub_1, sub_2] - """ - keep_parts, send_parts = [], [] - for i in range(len(local_cu_seqlens) - 1): - start, end = local_cu_seqlens[i], local_cu_seqlens[i + 1] - mid = (start + end) // 2 - keep_parts.append(hidden_states[start:mid]) # first half - send_parts.append(hidden_states[mid:end]) # second half + @staticmethod + def forward(ctx, hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size): + ctx.cp_group = cp_group + ctx.cp_rank = cp_rank + ctx.cp_size = cp_size + ctx.save_for_backward(local_cu_seqlens) + return _sequential_to_zigzag_impl(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size) - send_buf = torch.cat(send_parts, dim=0) - recv_buf = _ZigzagSequentialExchange.apply(send_buf, cp_group, cp_rank) + @staticmethod + def backward(ctx, grad_output): + (local_cu_seqlens,) = ctx.saved_tensors + # Backward: zigzag → sequential + result = _ZigzagToSequential.apply( + grad_output, local_cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size + ) + return result, None, None, None, None - # Rank 0: [keep (sub_0), recv (sub_3)] - # Rank 1: [recv (sub_1), keep (sub_2)] - result = [] - offset = 0 - for i in range(len(local_cu_seqlens) - 1): - chunk_len = (local_cu_seqlens[i + 1] - local_cu_seqlens[i]) // 2 - if cp_rank == 0: - result.append(keep_parts[i]) - result.append(recv_buf[offset : offset + chunk_len]) - else: - result.append(recv_buf[offset : offset + chunk_len]) - result.append(keep_parts[i]) - offset += chunk_len - return torch.cat(result, dim=0) + +def _zigzag_to_sequential(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size): + """Convert zigzag CP layout to sequential layout.""" + return _ZigzagToSequential.apply(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size) + + +def _sequential_to_zigzag(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size): + """Convert sequential CP layout back to zigzag layout.""" + return _SequentialToZigzag.apply(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size) class _AllGatherForDuplicatedComputation(torch.autograd.Function): @@ -226,6 +283,7 @@ def forward( local_cu_seqlens, mpu.get_context_parallel_group(), mpu.get_context_parallel_rank(), + cp_size, ) elif mpu.get_context_parallel_world_size() > 1: @@ -269,6 +327,7 @@ def forward( local_cu_seqlens, mpu.get_context_parallel_group(), mpu.get_context_parallel_rank(), + cp_size, ) elif mpu.get_context_parallel_world_size() > 1: diff --git a/tests/e2e/precision/test_qwen3_5_cp_correctness.py b/tests/e2e/precision/test_qwen3_5_cp_correctness.py index 3c897f2daa7..d0a2f3f32b6 100644 --- a/tests/e2e/precision/test_qwen3_5_cp_correctness.py +++ b/tests/e2e/precision/test_qwen3_5_cp_correctness.py @@ -86,11 +86,11 @@ def test_cp_forward_backward(rank, world_size): full_hidden_cp = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype) local_hidden = full_hidden_cp[:, start:end, :].clone().contiguous().requires_grad_(True) - # Local cu_seqlens for the chunk - local_cu = torch.tensor([0, local_seq_len], dtype=torch.int32, device=device) + # Global cu_seqlens (build_gdn_cp_context expects global boundaries) + global_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device) # Forward with CP - cp_out = model_cp(local_hidden, cu_seqlens=local_cu) + cp_out = model_cp(local_hidden, cu_seqlens=global_cu) cp_loss = cp_out.sum() # Reduce loss across ranks to match reference