Spaces:
Running on Zero
Running on Zero
File size: 4,996 Bytes
36c83de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """Data utilities for Context Parallel training.
Provides helpers for:
- Broadcasting tensors across CP ranks.
- Splitting temporal tensors by CP rank.
- Handling non-divisible temporal lengths via right-padding.
- Building frame-valid masks for padded temporal tails.
- Reducing loss scalars across CP ranks.
"""
from __future__ import annotations
import torch
import torch.distributed as dist
from torch import Tensor
from torch.distributed import ProcessGroup
def _cp_src_global_rank(group: ProcessGroup) -> int:
"""Global rank of CP-rank-0 in the given process group."""
return dist.get_global_rank(group, 0)
def cp_broadcast_tensor(
tensor: Tensor,
group: ProcessGroup,
) -> Tensor:
"""In-place broadcast *tensor* from CP-rank-0 to all ranks in *group*."""
src = _cp_src_global_rank(group)
dist.broadcast(tensor, src=src, group=group)
return tensor
def cp_split_temporal(
tensor: Tensor,
dim: int,
group: ProcessGroup,
) -> Tensor:
"""Slice *tensor* along *dim* to keep only this rank's temporal chunk."""
cp_rank = dist.get_rank(group)
cp_world = dist.get_world_size(group)
T = tensor.shape[dim]
assert T % cp_world == 0, f"Temporal size {T} (dim={dim}) must be divisible by cp_size={cp_world}"
chunk = T // cp_world
return tensor.narrow(dim, cp_rank * chunk, chunk).contiguous()
def cp_right_pad_size(length: int, multiple: int) -> int:
"""Return right-pad size needed to make ``length`` divisible by ``multiple``."""
if multiple <= 0:
raise ValueError(f"multiple must be > 0, got {multiple}")
return (-length) % multiple
def cp_right_pad_temporal(
tensor: Tensor,
dim: int,
pad_size: int,
value: float = 0.0,
) -> Tensor:
"""Right-pad ``tensor`` along temporal ``dim`` by ``pad_size``."""
if pad_size <= 0:
return tensor
if dim < 0:
dim = tensor.ndim + dim
if dim < 0 or dim >= tensor.ndim:
raise ValueError(f"Invalid dim={dim} for tensor with ndim={tensor.ndim}")
pad_shape = list(tensor.shape)
pad_shape[dim] = pad_size
pad_tensor = torch.full(
pad_shape,
fill_value=value,
dtype=tensor.dtype,
device=tensor.device,
)
return torch.cat([tensor, pad_tensor], dim=dim)
def cp_build_frame_valid_mask(clean_images: Tensor, pad_frames: int) -> Tensor:
"""Build ``(B, 1, T, 1, 1)`` frame-valid mask after temporal right-padding."""
if clean_images.ndim < 3:
raise ValueError(f"clean_images must have at least 3 dims (B, C, T, ...), got shape={list(clean_images.shape)}")
B = clean_images.shape[0]
T = clean_images.shape[2]
if pad_frames < 0 or pad_frames > T:
raise ValueError(f"pad_frames must satisfy 0 <= pad_frames <= T, got pad_frames={pad_frames}, T={T}")
mask = torch.ones((B, 1, T, 1, 1), device=clean_images.device, dtype=clean_images.dtype)
if pad_frames > 0:
mask[:, :, T - pad_frames :, :, :] = 0
return mask
def cp_reduce_loss(
loss: Tensor,
group: ProcessGroup,
num_valid_tokens: Tensor | int | float | None = None,
) -> Tensor:
"""Reduce CP-local loss to a global scalar with correct gradient scaling.
This function is autograd-safe for CP: it returns a forward value equal to
the CP-global reduced loss while preserving backward gradients scaled by the
local contribution ratio.
Args:
loss: Local scalar loss.
group: CP process group.
num_valid_tokens: Optional local token count for weighted reduction.
If omitted, all ranks are weighted equally.
"""
if num_valid_tokens is None:
loss_avg_detached = loss.detach().clone()
dist.all_reduce(loss_avg_detached, op=dist.ReduceOp.SUM, group=group)
loss_avg_detached = loss_avg_detached / dist.get_world_size(group)
# Keep local backward unchanged; only replace forward scalar for logging.
return loss + (loss_avg_detached - loss.detach())
if torch.is_tensor(num_valid_tokens):
local_tokens = num_valid_tokens.to(device=loss.device, dtype=loss.dtype)
else:
local_tokens = torch.tensor(float(num_valid_tokens), device=loss.device, dtype=loss.dtype)
world = dist.get_world_size(group)
total_tokens = local_tokens.detach().clone()
dist.all_reduce(total_tokens, op=dist.ReduceOp.SUM, group=group)
total_tokens = total_tokens.clamp_min(1.0)
# FSDP2 already averages grads across the CP-enabled sharding mesh.
# To obtain weighted global-token gradients, scale by n_i / mean(n).
mean_tokens = (total_tokens / world).clamp_min(1.0)
loss_for_backward = loss * (local_tokens / mean_tokens)
weighted_loss_detached = loss.detach() * local_tokens.detach()
dist.all_reduce(weighted_loss_detached, op=dist.ReduceOp.SUM, group=group)
loss_avg_detached = weighted_loss_detached / total_tokens
return loss_for_backward + (loss_avg_detached - loss_for_backward.detach())
|