Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / utils.py: 33%
105 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""PyTorch helpers used by context-parallel implementations."""
17# Context parallel is Torch-only after removal of the backend abstraction.
18# pylint: disable=forbidden-backend-import,missing-public-type-hints
20from __future__ import annotations
22__all__ = [
23 # Re-exported collectives.
24 "differentiable_all_gather_concat",
25 "differentiable_all_to_all_single",
26 "p2p_exchange",
27 # Async waits and the autograd functions backing them.
28 "differentiable_async_a2a_wait",
29 "differentiable_async_allgather_wait",
30 # Collective wrappers.
31 "all_to_all_single",
32 "all_gather_single",
33 "reduce_scatter_single",
34 # Hook registration helpers.
35 "register_forward_pre_hook",
36 "register_full_backward_pre_hook",
37 # Small aliases callers here and in the sibling modules rely on.
38 "Function",
39 "is_tensor",
40 "cat",
41]
43from typing import Sequence
45import torch
46import torch.distributed as dist
47from torch import Tensor
48from torch.nn import Module
50from hyper_parallel.core.utils.communication import (
51 differentiable_all_gather_concat,
52 differentiable_all_to_all_single,
53 p2p_exchange,
54)
56Function = torch.autograd.Function
59def _normalize_dim(dim: int, ndim: int) -> int:
60 """Normalize a possibly negative dimension index."""
61 return dim + ndim if dim < 0 else dim
64def _move_dim_to_front(tensor: Tensor, dim: int) -> Tensor:
65 """Move ``dim`` to the front while keeping the other dimensions ordered."""
66 dim = _normalize_dim(dim, tensor.dim())
67 if dim == 0:
68 return tensor.contiguous()
69 perm = [dim] + [index for index in range(tensor.dim()) if index != dim]
70 return tensor.permute(perm).contiguous()
73def _move_dim_from_front(tensor: Tensor, dim: int) -> Tensor:
74 """Move the leading dimension back to ``dim``."""
75 dim = _normalize_dim(dim, tensor.dim())
76 if dim == 0:
77 return tensor.contiguous()
78 perm = [dim] + [index for index in range(tensor.dim()) if index != dim]
79 inverse = [0] * len(perm)
80 for index, value in enumerate(perm):
81 inverse[value] = index
82 return tensor.permute(inverse).contiguous()
85def _a2a_reconstruct(out_perm: Tensor, concat_dim: int) -> Tensor:
86 """Reconstruct an all-to-all output from its leading-rank layout."""
87 chunk_in_perm = concat_dim + 1
88 recon_perm = list(range(1, chunk_in_perm)) + [0] + list(range(chunk_in_perm, out_perm.dim()))
89 reconstructed = out_perm.permute(recon_perm).contiguous()
90 shape = list(reconstructed.shape)
91 merged = shape[concat_dim] * shape[concat_dim + 1]
92 return reconstructed.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
95class _AsyncA2AWait(torch.autograd.Function):
96 """Wait for a pre-launched all-to-all and preserve its backward overlap."""
98 @staticmethod
99 def forward(ctx, tensor, work, out_perm, group, world_size, concat_dim, split_dim, handle_box):
100 """Wait and reconstruct the forward all-to-all output."""
101 ctx.group = group
102 ctx.world_size = world_size
103 ctx.concat_dim = concat_dim
104 ctx.split_dim = split_dim
105 ctx.handle_box = handle_box
106 ctx.input_shape = tensor.shape
107 work.wait()
108 return _a2a_reconstruct(out_perm, concat_dim)
110 @staticmethod
111 def backward(ctx, grad_output):
112 """Launch the reverse all-to-all when overlap was requested."""
113 if ctx.handle_box is not None:
114 grad_output = grad_output.contiguous()
115 shape = list(grad_output.shape)
116 seq_dim = ctx.concat_dim
117 full_size = shape[seq_dim]
118 ndim = len(shape) + 1
119 grad_perm = grad_output.reshape(
120 shape[:seq_dim]
121 + [ctx.world_size, full_size // ctx.world_size]
122 + shape[seq_dim + 1:]
123 ).permute(
124 [seq_dim] + list(range(seq_dim)) + list(range(seq_dim + 1, ndim))
125 ).contiguous()
126 out_perm = torch.empty_like(grad_perm)
127 work = dist.all_to_all_single(out_perm, grad_perm, group=ctx.group, async_op=True)
128 ctx.handle_box.append((work, out_perm))
129 return grad_output.new_zeros(ctx.input_shape), None, None, None, None, None, None, None
132class _AsyncAllGatherWait(torch.autograd.Function):
133 """Wait for a pre-launched all-gather and provide reduce-scatter backward."""
135 @staticmethod
136 def forward(ctx, tensor, work, out_perm, group, world_size, gather_dim, handle_box):
137 """Wait and reconstruct the gathered tensor."""
138 ctx.group = group
139 ctx.world_size = world_size
140 ctx.gather_dim = gather_dim
141 ctx.handle_box = handle_box
142 ctx.input_shape = tensor.shape
143 work.wait()
144 return _move_dim_from_front(out_perm, gather_dim)
146 @staticmethod
147 def backward(ctx, grad_output):
148 """Reduce-scatter the gathered gradient."""
149 grad_perm = _move_dim_to_front(grad_output.contiguous(), ctx.gather_dim)
150 output_shape = list(grad_perm.shape)
151 if output_shape[0] % ctx.world_size != 0:
152 raise ValueError(
153 "all_gather backward expected gathered dimension to be divisible by world_size, "
154 f"got {output_shape[0]} and {ctx.world_size}."
155 )
156 output_shape[0] //= ctx.world_size
157 output = torch.empty(output_shape, dtype=grad_perm.dtype, device=grad_perm.device)
158 work = dist.reduce_scatter_tensor(output, grad_perm, group=ctx.group, async_op=True)
159 if ctx.handle_box is not None:
160 ctx.handle_box.append((work, output, ctx.gather_dim))
161 return grad_output.new_zeros(ctx.input_shape), None, None, None, None, None, None
162 work.wait()
163 return _move_dim_from_front(output, ctx.gather_dim), None, None, None, None, None, None
166def is_tensor(value) -> bool:
167 """Return whether ``value`` is a PyTorch tensor."""
168 return isinstance(value, Tensor)
171def register_forward_pre_hook(module: Module, hook, *, with_kwargs: bool = False):
172 """Register a module forward pre-hook."""
173 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs)
176def register_full_backward_pre_hook(module: Module, hook):
177 """Register a module full backward pre-hook."""
178 return module.register_full_backward_pre_hook(hook)
181def all_to_all_single(input_tensor: Tensor, output_shape: Sequence[int], group, async_op: bool = False):
182 """Run a fixed-shape all-to-all collective."""
183 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
184 work = dist.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
185 return output, work
188def all_gather_single(input_tensor: Tensor, output_shape: Sequence[int], group, async_op: bool = False):
189 """Run an all-gather-into-tensor collective."""
190 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
191 work = dist.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
192 return output, work
195def reduce_scatter_single(input_tensor: Tensor, output_shape: Sequence[int], group, async_op: bool = False):
196 """Run a reduce-scatter-into-tensor collective."""
197 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
198 work = dist.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
199 return output, work
202def differentiable_async_a2a_wait(
203 tensor: Tensor, work, out_perm: Tensor, group, world_size: int, concat_dim: int, split_dim: int, handle_box=None
204) -> Tensor:
205 """Wait for an asynchronous all-to-all while preserving autograd."""
206 return _AsyncA2AWait.apply(tensor, work, out_perm, group, world_size, concat_dim, split_dim, handle_box)
209def differentiable_async_allgather_wait(
210 tensor: Tensor, work, out_perm: Tensor, group, world_size: int, gather_dim: int, handle_box=None
211) -> Tensor:
212 """Wait for an asynchronous all-gather while preserving autograd."""
213 return _AsyncAllGatherWait.apply(tensor, work, out_perm, group, world_size, gather_dim, handle_box)
216def cat(tensors, dim: int = 0) -> Tensor:
217 """Concatenate tensors along ``dim``."""
218 return torch.cat(tensors, dim=dim)