SubstanceSHIFT/SeedVR2-3B
0
1# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates2# //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 at6# //7# // http://www.apache.org/licenses/LICENSE-2.08# //9# // Unless required by applicable law or agreed to in writing, software10# // 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 and13# // limitations under the License.14 15"""16Distributed ops for supporting sequence parallel.17"""18 19from collections import defaultdict20from typing import Any, Callable, Dict, List, Optional, Tuple, Union21import torch22import torch.distributed as dist23from torch import Tensor24 25from common.cache import Cache26from common.distributed.advanced import (27 get_sequence_parallel_group,28 get_sequence_parallel_rank,29 get_sequence_parallel_world_size,30)31 32from .basic import get_device33 34_SEQ_DATA_BUF = defaultdict(lambda: [None, None, None])35_SEQ_DATA_META_SHAPES = defaultdict()36_SEQ_DATA_META_DTYPES = defaultdict()37_SEQ_DATA_ASYNC_COMMS = defaultdict(list)38_SYNC_BUFFER = defaultdict(dict)39 40 41def single_all_to_all(42 local_input: Tensor,43 scatter_dim: int,44 gather_dim: int,45 group: dist.ProcessGroup,46 async_op: bool = False,47):48 """49 A function to do all-to-all on a tensor50 """51 seq_world_size = dist.get_world_size(group)52 prev_scatter_dim = scatter_dim53 if scatter_dim != 0:54 local_input = local_input.transpose(0, scatter_dim)55 if gather_dim == 0:56 gather_dim = scatter_dim57 scatter_dim = 058 59 inp_shape = list(local_input.shape)60 inp_shape[scatter_dim] = inp_shape[scatter_dim] // seq_world_size61 input_t = local_input.reshape(62 [seq_world_size, inp_shape[scatter_dim]] + inp_shape[scatter_dim + 1 :]63 ).contiguous()64 output = torch.empty_like(input_t)65 comm = dist.all_to_all_single(output, input_t, group=group, async_op=async_op)66 if async_op:67 # let user's code transpose & reshape68 return output, comm, prev_scatter_dim69 70 # first dim is seq_world_size, so we can split it directly71 output = torch.cat(output.split(1), dim=gather_dim + 1).squeeze(0)72 if prev_scatter_dim:73 output = output.transpose(0, prev_scatter_dim).contiguous()74 return output75 76 77def _all_to_all(78 local_input: Tensor,79 scatter_dim: int,80 gather_dim: int,81 group: dist.ProcessGroup,82):83 seq_world_size = dist.get_world_size(group)84 input_list = [85 t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)86 ]87 output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]88 dist.all_to_all(output_list, input_list, group=group)89 return torch.cat(output_list, dim=gather_dim).contiguous()90 91 92class SeqAllToAll(torch.autograd.Function):93 @staticmethod94 def forward(95 ctx: Any,96 group: dist.ProcessGroup,97 local_input: Tensor,98 scatter_dim: int,99 gather_dim: int,100 async_op: bool,101 ) -> Tensor:102 ctx.group = group103 ctx.scatter_dim = scatter_dim104 ctx.gather_dim = gather_dim105 ctx.async_op = async_op106 if async_op:107 output, comm, prev_scatter_dim = single_all_to_all(108 local_input, scatter_dim, gather_dim, group, async_op=async_op109 )110 ctx.prev_scatter_dim = prev_scatter_dim111 return output, comm112 113 return _all_to_all(local_input, scatter_dim, gather_dim, group)114 115 @staticmethod116 def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:117 if ctx.async_op:118 input_t = torch.cat(grad_output[0].split(1), dim=ctx.gather_dim + 1).squeeze(0)119 if ctx.prev_scatter_dim:120 input_t = input_t.transpose(0, ctx.prev_scatter_dim)121 else:122 input_t = grad_output[0]123 return (124 None,125 _all_to_all(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group),126 None,127 None,128 None,129 )130 131 132class Slice(torch.autograd.Function):133 @staticmethod134 def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int) -> Tensor:135 ctx.group = group136 ctx.rank = dist.get_rank(group)137 seq_world_size = dist.get_world_size(group)138 ctx.seq_world_size = seq_world_size139 ctx.dim = dim140 dim_size = local_input.shape[dim]141 return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous()142 143 @staticmethod144 def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]:145 dim_size = list(grad_output.size())146 split_size = dim_size[0]147 dim_size[0] = dim_size[0] * ctx.seq_world_size148 output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device())149 dist._all_gather_base(output, grad_output, group=ctx.group)150 return (None, torch.cat(output.split(split_size), dim=ctx.dim), None)151 152 153class Gather(torch.autograd.Function):154 @staticmethod155 def forward(156 ctx: Any,157 group: dist.ProcessGroup,158 local_input: Tensor,159 dim: int,160 grad_scale: Optional[bool] = False,161 ) -> Tensor:162 ctx.group = group163 ctx.rank = dist.get_rank(group)164 ctx.dim = dim165 ctx.grad_scale = grad_scale166 seq_world_size = dist.get_world_size(group)167 ctx.seq_world_size = seq_world_size168 dim_size = list(local_input.size())169 split_size = dim_size[0]170 ctx.part_size = dim_size[dim]171 dim_size[0] = dim_size[0] * seq_world_size172 output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device())173 dist._all_gather_base(output, local_input.contiguous(), group=ctx.group)174 return torch.cat(output.split(split_size), dim=dim)175 176 @staticmethod177 def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]:178 if ctx.grad_scale:179 grad_output = grad_output * ctx.seq_world_size180 return (181 None,182 grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(),183 None,184 None,185 )186 187 188def gather_seq_scatter_heads_qkv(189 qkv_tensor: Tensor,190 *,191 seq_dim: int,192 qkv_shape: Optional[Tensor] = None,193 cache: Cache = Cache(disable=True),194 restore_shape: bool = True,195):196 """197 A func to sync splited qkv tensor198 qkv_tensor: the tensor we want to do alltoall with. The last dim must199 be the projection_idx, which we will split into 3 part. After200 spliting, the gather idx will be projecttion_idx + 1201 seq_dim: gather_dim for all2all comm202 restore_shape: if True, output will has the same shape length as input203 """204 group = get_sequence_parallel_group()205 if not group:206 return qkv_tensor207 world = get_sequence_parallel_world_size()208 orig_shape = qkv_tensor.shape209 scatter_dim = qkv_tensor.dim()210 bef_all2all_shape = list(orig_shape)211 qkv_proj_dim = bef_all2all_shape[-1]212 bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3]213 qkv_tensor = qkv_tensor.view(bef_all2all_shape)214 qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, False)215 if restore_shape:216 out_shape = list(orig_shape)217 out_shape[seq_dim] *= world218 out_shape[-1] = qkv_proj_dim // world219 qkv_tensor = qkv_tensor.view(out_shape)220 221 # remove padding222 if qkv_shape is not None:223 unpad_dim_size = cache(224 "unpad_dim_size", lambda: torch.sum(torch.prod(qkv_shape, dim=-1)).item()225 )226 if unpad_dim_size % world != 0:227 padding_size = qkv_tensor.size(seq_dim) - unpad_dim_size228 qkv_tensor = _unpad_tensor(qkv_tensor, seq_dim, padding_size)229 return qkv_tensor230 231 232def slice_inputs(x: Tensor, dim: int, padding: bool = True):233 """234 A func to slice the input sequence in sequence parallel235 """236 group = get_sequence_parallel_group()237 if group is None:238 return x239 sp_rank = get_sequence_parallel_rank()240 sp_world = get_sequence_parallel_world_size()241 dim_size = x.shape[dim]242 unit = (dim_size + sp_world - 1) // sp_world243 if padding and dim_size % sp_world:244 padding_size = sp_world - (dim_size % sp_world)245 x = _pad_tensor(x, dim, padding_size)246 slc = [slice(None)] * len(x.shape)247 slc[dim] = slice(unit * sp_rank, unit * (sp_rank + 1))248 return x[slc]249 250 251def remove_seqeunce_parallel_padding(x: Tensor, dim: int, unpad_dim_size: int):252 """253 A func to remove the padding part of the tensor based on its original shape254 """255 group = get_sequence_parallel_group()256 if group is None:257 return x258 sp_world = get_sequence_parallel_world_size()259 if unpad_dim_size % sp_world == 0:260 return x261 padding_size = sp_world - (unpad_dim_size % sp_world)262 assert (padding_size + unpad_dim_size) % sp_world == 0263 return _unpad_tensor(x, dim=dim, padding_size=padding_size)264 265 266def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor:267 """268 A func to sync attention result with alltoall in sequence parallel269 """270 group = get_sequence_parallel_group()271 if not group:272 return x273 dim_size = x.size(seq_dim)274 sp_world = get_sequence_parallel_world_size()275 if dim_size % sp_world != 0:276 padding_size = sp_world - (dim_size % sp_world)277 x = _pad_tensor(x, seq_dim, padding_size)278 return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)279 280 281def gather_seq_scatter_heads(x: Tensor, seq_dim: int, head_dim: int) -> Tensor:282 """283 A func to sync embedding input with alltoall in sequence parallel284 """285 group = get_sequence_parallel_group()286 if not group:287 return x288 return SeqAllToAll.apply(group, x, head_dim, seq_dim, False)289 290 291def scatter_heads(x: Tensor, dim: int) -> Tensor:292 """293 A func to split heads before attention in sequence parallel294 """295 group = get_sequence_parallel_group()296 if not group:297 return x298 return Slice.apply(group, x, dim)299 300 301def gather_heads(x: Tensor, dim: int, grad_scale: Optional[bool] = False) -> Tensor:302 """303 A func to gather heads for the attention result in sequence parallel304 """305 group = get_sequence_parallel_group()306 if not group:307 return x308 return Gather.apply(group, x, dim, grad_scale)309 310 311def gather_outputs(312 x: Tensor,313 *,314 gather_dim: int,315 padding_dim: Optional[int] = None,316 unpad_shape: Optional[Tensor] = None,317 cache: Cache = Cache(disable=True),318 scale_grad=True,319):320 """321 A func to gather the outputs for the model result in sequence parallel322 """323 group = get_sequence_parallel_group()324 if not group:325 return x326 x = Gather.apply(group, x, gather_dim, scale_grad)327 if padding_dim is not None:328 unpad_dim_size = cache(329 "unpad_dim_size", lambda: torch.sum(torch.prod(unpad_shape, dim=1)).item()330 )331 x = remove_seqeunce_parallel_padding(x, padding_dim, unpad_dim_size)332 return x333 334 335def _pad_tensor(x: Tensor, dim: int, padding_size: int):336 shape = list(x.shape)337 shape[dim] = padding_size338 pad = torch.zeros(shape, dtype=x.dtype, device=x.device)339 return torch.cat([x, pad], dim=dim)340 341 342def _unpad_tensor(x: Tensor, dim: int, padding_size):343 slc = [slice(None)] * len(x.shape)344 slc[dim] = slice(0, -padding_size)345 return x[slc]346 347 348def _broadcast_data(data, shape, dtype, src, group, async_op):349 comms = []350 if isinstance(data, (list, tuple)):351 for i, sub_shape in enumerate(shape):352 comms += _broadcast_data(data[i], sub_shape, dtype[i], src, group, async_op)353 elif isinstance(data, dict):354 for key, sub_data in data.items():355 comms += _broadcast_data(sub_data, shape[key], dtype[key], src, group, async_op)356 elif isinstance(data, Tensor):357 comms.append(dist.broadcast(data, src=src, group=group, async_op=async_op))358 return comms359 360 361def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]:362 if isinstance(data, (list, tuple)):363 return [_traverse(sub_data, op) for sub_data in data]364 elif isinstance(data, dict):365 return {key: _traverse(sub_data, op) for key, sub_data in data.items()}366 elif isinstance(data, Tensor):367 return op(data)368 else:369 return None370 371 372def _get_shapes(data):373 return _traverse(data, op=lambda x: x.shape)374 375 376def _get_dtypes(data):377 return _traverse(data, op=lambda x: x.dtype)378 379 380def _construct_broadcast_buffer(shapes, dtypes, device):381 if isinstance(shapes, torch.Size):382 return torch.empty(shapes, dtype=dtypes, device=device)383 384 if isinstance(shapes, (list, tuple)):385 buffer = []386 for i, sub_shape in enumerate(shapes):387 buffer.append(_construct_broadcast_buffer(sub_shape, dtypes[i], device))388 elif isinstance(shapes, dict):389 buffer = {}390 for key, sub_shape in shapes.items():391 buffer[key] = _construct_broadcast_buffer(sub_shape, dtypes[key], device)392 else:393 return None394 return buffer395 396 397class SPDistForward:398 """A forward tool to sync different result across sp group399 400 Args:401 module: a function or module to process users input402 sp_step: current training step to judge which rank to broadcast its result to all403 name: a distinct str to save meta and async comm404 comm_shape: if different ranks have different shape, mark this arg to True405 device: the device for current rank, can be empty406 """407 408 def __init__(409 self,410 name: str,411 comm_shape: bool,412 device: torch.device = None,413 ):414 self.name = name415 self.comm_shape = comm_shape416 if device:417 self.device = device418 else:419 self.device = get_device()420 421 def __call__(self, inputs) -> Any:422 group = get_sequence_parallel_group()423 if not group:424 yield inputs425 else:426 device = self.device427 sp_world = get_sequence_parallel_world_size()428 sp_rank = get_sequence_parallel_rank()429 for local_step in range(sp_world):430 src_rank = dist.get_global_rank(group, local_step)431 is_src = sp_rank == local_step432 local_shapes = []433 local_dtypes = []434 if local_step == 0:435 local_result = inputs436 _SEQ_DATA_BUF[self.name][-1] = local_result437 local_shapes = _get_shapes(local_result)438 local_dtypes = _get_dtypes(local_result)439 if self.comm_shape:440 group_shapes_lists = [None] * sp_world441 dist.all_gather_object(group_shapes_lists, local_shapes, group=group)442 _SEQ_DATA_META_SHAPES[self.name] = group_shapes_lists443 else:444 _SEQ_DATA_META_SHAPES[self.name] = [local_shapes] * sp_world445 _SEQ_DATA_META_DTYPES[self.name] = local_dtypes446 shapes = _SEQ_DATA_META_SHAPES[self.name][local_step]447 dtypes = _SEQ_DATA_META_DTYPES[self.name]448 buf_id = local_step % 2449 if local_step == 0:450 sync_data = (451 local_result452 if is_src453 else _construct_broadcast_buffer(shapes, dtypes, device)454 )455 _broadcast_data(sync_data, shapes, dtypes, src_rank, group, False)456 _SEQ_DATA_BUF[self.name][buf_id] = sync_data457 458 # wait for async comm ops459 if _SEQ_DATA_ASYNC_COMMS[self.name]:460 for comm in _SEQ_DATA_ASYNC_COMMS[self.name]:461 comm.wait()462 # before return the sync result, do async broadcast for next batch463 if local_step < sp_world - 1:464 next_buf_id = 1 - buf_id465 shapes = _SEQ_DATA_META_SHAPES[self.name][local_step + 1]466 src_rank = dist.get_global_rank(group, local_step + 1)467 is_src = sp_rank == local_step + 1468 next_sync_data = (469 _SEQ_DATA_BUF[self.name][-1]470 if is_src471 else _construct_broadcast_buffer(shapes, dtypes, device)472 )473 _SEQ_DATA_ASYNC_COMMS[self.name] = _broadcast_data(474 next_sync_data, shapes, dtypes, src_rank, group, True475 )476 _SEQ_DATA_BUF[self.name][next_buf_id] = next_sync_data477 yield _SEQ_DATA_BUF[self.name][buf_id]478 479 480sync_inputs = SPDistForward(name="bef_fwd", comm_shape=True)481 482 483def sync_data(data, sp_idx, name="tmp"):484 group = get_sequence_parallel_group()485 if group is None:486 return data487 # if sp_idx in _SYNC_BUFFER[name]:488 # return _SYNC_BUFFER[name][sp_idx]489 sp_rank = get_sequence_parallel_rank()490 src_rank = dist.get_global_rank(group, sp_idx)491 objects = [data] if sp_rank == sp_idx else [None]492 dist.broadcast_object_list(objects, src=src_rank, group=group)493 # _SYNC_BUFFER[name] = {sp_idx: objects[0]}494 return objects[0]495