wehub-resource-sync
222 lines
8.6 KiB
Python
222 lines
8.6 KiB
Python
"""Differentiable halo exchange for temporal convolutions under Context Parallel.
|
|
|
|
When the temporal sequence is sharded across CP ranks, causal convolutions
|
|
of kernel size K need K-1 frames of left context from the previous rank.
|
|
Bidirectional convolutions additionally need right context from the next rank.
|
|
|
|
This module provides ``cp_halo_exchange``, a differentiable primitive that
|
|
uses ``torch.distributed.batch_isend_irecv`` for P2P communication on the
|
|
CP process group. Gradients flow back correctly through the same channel.
|
|
|
|
Safety with FSDP2: FSDP2 uses stream-to-stream synchronization (not
|
|
``recordStream``), so P2P ops on a separate CP group are inherently safe
|
|
and will not cause deadlocks.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch import Tensor
|
|
from torch.distributed import P2POp, ProcessGroup
|
|
|
|
from diffusion.distributed.context_parallel.config import get_cp_halo_impl
|
|
|
|
|
|
def _to_global_rank(group: ProcessGroup, local_rank: int) -> int:
|
|
"""Convert a group-local rank to its global rank."""
|
|
return dist.get_global_rank(group, local_rank)
|
|
|
|
|
|
def _allgather_tensor(inp: Tensor, group: ProcessGroup) -> Tensor:
|
|
"""All-gather helper returning ``(world, *inp.shape)``."""
|
|
world = dist.get_world_size(group)
|
|
inp_contig = inp.contiguous()
|
|
flat_out = torch.empty(
|
|
(world * inp_contig.shape[0],) + tuple(inp_contig.shape[1:]),
|
|
dtype=inp_contig.dtype,
|
|
device=inp_contig.device,
|
|
)
|
|
dist.all_gather_into_tensor(flat_out, inp_contig, group=group)
|
|
return flat_out.reshape((world,) + tuple(inp_contig.shape))
|
|
|
|
|
|
class _CPHaloExchange(torch.autograd.Function):
|
|
"""Differentiable halo exchange via P2P send/recv.
|
|
|
|
Forward:
|
|
Rank r sends its first ``right_size`` slices to rank r-1 (as their
|
|
right halo) and its last ``left_size`` slices to rank r+1 (as their
|
|
left halo). Conversely, it receives left halo from rank r-1 and
|
|
right halo from rank r+1.
|
|
|
|
Backward:
|
|
Gradients are routed back via the reverse P2P direction.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx: object,
|
|
x: Tensor,
|
|
left_size: int,
|
|
right_size: int,
|
|
dim: int,
|
|
group: ProcessGroup,
|
|
) -> Tensor:
|
|
rank = dist.get_rank(group)
|
|
world = dist.get_world_size(group)
|
|
|
|
ctx.left_size = left_size
|
|
ctx.right_size = right_size
|
|
ctx.dim = dim
|
|
ctx.group = group
|
|
ctx.rank = rank
|
|
ctx.world = world
|
|
|
|
T = x.shape[dim]
|
|
|
|
left_recv = torch.zeros_like(x.narrow(dim, 0, left_size)) if left_size > 0 else None
|
|
right_recv = torch.zeros_like(x.narrow(dim, 0, right_size)) if right_size > 0 else None
|
|
|
|
halo_impl = get_cp_halo_impl()
|
|
if halo_impl == "p2p":
|
|
ops: list[P2POp] = []
|
|
|
|
if left_size > 0:
|
|
if rank > 0:
|
|
peer = _to_global_rank(group, rank - 1)
|
|
ops.append(P2POp(dist.irecv, left_recv, peer, group=group))
|
|
if rank < world - 1:
|
|
send_buf = x.narrow(dim, T - left_size, left_size).contiguous()
|
|
peer = _to_global_rank(group, rank + 1)
|
|
ops.append(P2POp(dist.isend, send_buf, peer, group=group))
|
|
|
|
if right_size > 0:
|
|
if rank < world - 1:
|
|
peer = _to_global_rank(group, rank + 1)
|
|
ops.append(P2POp(dist.irecv, right_recv, peer, group=group))
|
|
if rank > 0:
|
|
send_buf = x.narrow(dim, 0, right_size).contiguous()
|
|
peer = _to_global_rank(group, rank - 1)
|
|
ops.append(P2POp(dist.isend, send_buf, peer, group=group))
|
|
|
|
if ops:
|
|
reqs = dist.batch_isend_irecv(ops)
|
|
for req in reqs:
|
|
req.wait()
|
|
else:
|
|
# Deterministic collective path: all ranks always participate in
|
|
# the same collectives, which is safer under FSDP2 overlap.
|
|
if left_size > 0:
|
|
send_left = x.narrow(dim, T - left_size, left_size).contiguous()
|
|
gathered_left = _allgather_tensor(send_left, group)
|
|
left_recv = gathered_left[rank - 1].clone() if rank > 0 else torch.zeros_like(send_left)
|
|
if right_size > 0:
|
|
send_right = x.narrow(dim, 0, right_size).contiguous()
|
|
gathered_right = _allgather_tensor(send_right, group)
|
|
right_recv = gathered_right[rank + 1].clone() if rank < world - 1 else torch.zeros_like(send_right)
|
|
|
|
parts: list[Tensor] = []
|
|
if left_size > 0:
|
|
parts.append(left_recv)
|
|
parts.append(x)
|
|
if right_size > 0:
|
|
parts.append(right_recv)
|
|
|
|
out = torch.cat(parts, dim=dim)
|
|
return out
|
|
|
|
@staticmethod
|
|
def backward(ctx: object, grad_output: Tensor) -> tuple[Tensor | None, ...]:
|
|
left_size = ctx.left_size
|
|
right_size = ctx.right_size
|
|
dim = ctx.dim
|
|
group = ctx.group
|
|
rank = ctx.rank
|
|
world = ctx.world
|
|
|
|
T_with_halo = grad_output.shape[dim]
|
|
T_local = T_with_halo - left_size - right_size
|
|
|
|
grad_left = grad_output.narrow(dim, 0, left_size) if left_size > 0 else None
|
|
grad_local = grad_output.narrow(dim, left_size, T_local)
|
|
grad_right = grad_output.narrow(dim, left_size + T_local, right_size) if right_size > 0 else None
|
|
|
|
recv_from_left = (
|
|
torch.zeros_like(grad_local.narrow(dim, T_local - left_size, left_size)) if left_size > 0 else None
|
|
)
|
|
recv_from_right = torch.zeros_like(grad_local.narrow(dim, 0, right_size)) if right_size > 0 else None
|
|
|
|
halo_impl = get_cp_halo_impl()
|
|
if halo_impl == "p2p":
|
|
ops: list[P2POp] = []
|
|
|
|
if left_size > 0:
|
|
if rank > 0:
|
|
peer = _to_global_rank(group, rank - 1)
|
|
ops.append(P2POp(dist.isend, grad_left.contiguous(), peer, group=group))
|
|
if rank < world - 1:
|
|
peer = _to_global_rank(group, rank + 1)
|
|
ops.append(P2POp(dist.irecv, recv_from_left, peer, group=group))
|
|
|
|
if right_size > 0:
|
|
if rank < world - 1:
|
|
peer = _to_global_rank(group, rank + 1)
|
|
ops.append(P2POp(dist.isend, grad_right.contiguous(), peer, group=group))
|
|
if rank > 0:
|
|
peer = _to_global_rank(group, rank - 1)
|
|
ops.append(P2POp(dist.irecv, recv_from_right, peer, group=group))
|
|
|
|
if ops:
|
|
reqs = dist.batch_isend_irecv(ops)
|
|
for req in reqs:
|
|
req.wait()
|
|
else:
|
|
# Collective gradient routing mirrors forward neighbor selection.
|
|
if left_size > 0:
|
|
gathered_grad_left = _allgather_tensor(grad_left.contiguous(), group)
|
|
recv_from_left = (
|
|
gathered_grad_left[rank + 1].clone() if rank < world - 1 else torch.zeros_like(recv_from_left)
|
|
)
|
|
if right_size > 0:
|
|
gathered_grad_right = _allgather_tensor(grad_right.contiguous(), group)
|
|
recv_from_right = (
|
|
gathered_grad_right[rank - 1].clone() if rank > 0 else torch.zeros_like(recv_from_right)
|
|
)
|
|
|
|
grad_x = grad_local.clone()
|
|
if left_size > 0 and recv_from_left is not None:
|
|
grad_x.narrow(dim, T_local - left_size, left_size).add_(recv_from_left)
|
|
if right_size > 0 and recv_from_right is not None:
|
|
grad_x.narrow(dim, 0, right_size).add_(recv_from_right)
|
|
|
|
return grad_x, None, None, None, None
|
|
|
|
|
|
def cp_halo_exchange(
|
|
x: Tensor,
|
|
left_size: int,
|
|
right_size: int,
|
|
dim: int,
|
|
group: ProcessGroup,
|
|
) -> Tensor:
|
|
"""Exchange halo regions between CP ranks along the given dimension.
|
|
|
|
Args:
|
|
x: Local tensor shard.
|
|
left_size: Number of slices to receive from the left neighbor
|
|
(appended before ``x`` along ``dim``). Rank 0 gets zero-padding.
|
|
right_size: Number of slices to receive from the right neighbor
|
|
(appended after ``x`` along ``dim``). Last rank gets zero-padding.
|
|
dim: Dimension along which to exchange halos.
|
|
group: CP process group.
|
|
|
|
Returns:
|
|
Tensor with shape ``x.shape[dim] + left_size + right_size`` along
|
|
``dim``, where boundary halos are filled from neighbors (or zeros
|
|
for edge ranks).
|
|
"""
|
|
if left_size == 0 and right_size == 0:
|
|
return x
|
|
return _CPHaloExchange.apply(x, left_size, right_size, dim, group)
|