| import os |
| import math |
| import random |
| import argparse |
| import datetime |
| import logging |
| import inspect |
| import subprocess |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.distributed.device_mesh import init_device_mesh |
| from einops import rearrange, repeat |
|
|
|
|
| dp_size = None |
| cp_size = None |
| dp_group = None |
| cp_group = None |
| cp_stream = None |
| dp_ranks = None |
| cp_ranks = None |
| dp_rank = None |
| cp_rank = None |
|
|
|
|
| def init_context_parallel(context_parallel_size: int = 1, |
| global_rank: int = 1, |
| world_size: int = 1,): |
|
|
| global dp_size |
| global cp_size |
| global dp_group |
| global cp_group |
| global dp_ranks |
| global cp_ranks |
| global dp_rank |
| global cp_rank |
|
|
|
|
| if world_size%context_parallel_size != 0: |
| raise RuntimeError(f'world_size {world_size} must be multiple of context_parallel_size {context_parallel_size}!!!') |
|
|
|
|
| cp_size = context_parallel_size |
| dp_size = world_size//context_parallel_size |
|
|
|
|
| print(f'[rank {global_rank}] init_device_mesh [dp_size x cp_size]: [{dp_size} x {cp_size}]') |
|
|
| mesh_2d = init_device_mesh("cuda", (dp_size, cp_size), mesh_dim_names=("dp", "cp")) |
|
|
| print(f'[rank {global_rank}] mesh_2d: {mesh_2d}') |
|
|
| dp_group = mesh_2d.get_group(mesh_dim="dp") |
| cp_group = mesh_2d.get_group(mesh_dim="cp") |
|
|
| dp_ranks = torch.distributed.get_process_group_ranks(dp_group) |
| cp_ranks = torch.distributed.get_process_group_ranks(cp_group) |
|
|
| dp_rank = dist.get_rank(group=dp_group) |
| cp_rank = dist.get_rank(group=cp_group) |
|
|
| global_rank_1 = torch.distributed.get_rank() |
| print(f'[rank {global_rank_1}] [dp_rank, cp_rank]: [{dp_rank}, {cp_rank}], dp_ranks: {dp_ranks}, cp_ranks: {cp_ranks}') |
|
|
|
|
| def get_cp_size(): |
|
|
| global cp_size |
| return cp_size |
|
|
| def get_dp_size(): |
|
|
| global dp_size |
| return dp_size |
|
|
| def get_cp_stream(): |
|
|
| global cp_stream |
| if cp_stream == None: |
| cp_stream = torch.cuda.Stream() |
| |
| return cp_stream |
|
|
| def get_dp_group(): |
|
|
| global dp_group |
| return dp_group |
|
|
| def get_cp_group(): |
|
|
| global cp_group |
| return cp_group |
|
|
|
|
| def get_dp_rank(): |
|
|
| global dp_rank |
| global cp_rank |
|
|
| return dp_rank |
|
|
|
|
| def get_cp_rank(): |
|
|
| global dp_rank |
| global cp_rank |
|
|
| return cp_rank |
|
|
|
|
|
|
| def get_cp_rank_list(): |
| |
| global cp_ranks |
| if cp_ranks == None: |
| cp_ranks = torch.distributed.get_process_group_ranks(cp_group) |
| return cp_ranks |
|
|
|
|
| def cp_broadcast(tensor, cp_index=0): |
|
|
| global dp_group |
| global cp_group |
|
|
| cp_ranks = get_cp_rank_list() |
|
|
| torch.distributed.broadcast(tensor, cp_ranks[cp_index], group=cp_group) |
|
|
|
|
|
|
|
|
| def cp_broadcast_objects(tensor): |
|
|
| global dp_group |
| global cp_group |
|
|
| raise NotImplementedError("cp_broadcast_objects method is not yet implemented!!!") |
|
|
|
|
|
|
|
|
| def split_tensor_in_cp(input, seq_dim): |
|
|
| global cp_size |
|
|
| seq_size = input.shape[seq_dim] |
|
|
| if seq_size%cp_size != 0: |
| raise RuntimeError(f'seq_length {seq_size} in dim {seq_dim} must be multiple of cp_size {cp_size}!!!') |
|
|
| split_seq_size = seq_size//cp_size |
|
|
| tensor_splits = input.split(split_seq_size, dim=seq_dim) |
|
|
| cp_rank = get_cp_rank() |
|
|
| split_tensor = tensor_splits[cp_rank] |
|
|
| return split_tensor |
|
|
|
|
|
|
|
|
|
|
| class GatherFunction(torch.autograd.Function): |
|
|
| @staticmethod |
| def forward(ctx, input, process_group, seq_dim, frames): |
| ctx.cp_group = process_group |
| ctx.seq_dim = seq_dim |
| ctx.frames = frames |
| ctx.cp_size = get_cp_size() |
|
|
| input = rearrange(input, "B (T S) C -> B T S C", T=frames) |
|
|
| with torch.no_grad(): |
|
|
| input = input.contiguous() |
|
|
| output_tensors = [torch.zeros_like(input) for _ in range(ctx.cp_size)] |
|
|
| dist.all_gather(output_tensors, input, group=ctx.cp_group) |
|
|
| output_tensor = torch.cat(output_tensors, dim=seq_dim) |
|
|
|
|
|
|
| output_tensor = rearrange(output_tensor, "B T S C -> B (T S) C", T=frames) |
|
|
|
|
| return output_tensor |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| |
|
|
| with torch.no_grad(): |
| |
| grad_output = grad_output * ctx.cp_size |
|
|
| grad_output = rearrange(grad_output, "B (T S) C -> B T S C", T=ctx.frames) |
|
|
| grad_input = split_tensor_in_cp(grad_output, ctx.seq_dim) |
|
|
| grad_input = rearrange(grad_input, "B T S C -> B (T S) C", T=ctx.frames) |
| |
|
|
| return grad_input, None, None, None |
|
|
|
|
|
|
|
|
| class SplitFunction(torch.autograd.Function): |
|
|
| @staticmethod |
| def forward(ctx, input, process_group, seq_dim): |
| ctx.cp_group = process_group |
| ctx.seq_dim = seq_dim |
| ctx.cp_size = get_cp_size() |
|
|
| output_tensor = split_tensor_in_cp(input, ctx.seq_dim) |
|
|
| return output_tensor |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| |
|
|
| with torch.no_grad(): |
|
|
|
|
| grad_output = grad_output / ctx.cp_size |
|
|
| output_tensors = [torch.zeros_like(grad_output) for _ in range(ctx.cp_size)] |
|
|
| dist.all_gather(output_tensors, grad_output, group=ctx.cp_group) |
|
|
| grad_input = torch.cat(output_tensors, dim=ctx.seq_dim) |
|
|
|
|
| return grad_input, None, None |
|
|
|
|
|
|
| def gather_cp(input, frames): |
|
|
| cp_process_group = get_cp_group() |
| |
| output_tensor = GatherFunction.apply(input, cp_process_group, 2, frames) |
|
|
| return output_tensor |
|
|
|
|
| def split_cp(input, seq_dim): |
|
|
| cp_process_group = get_cp_group() |
| |
| output_tensor = SplitFunction.apply(input, cp_process_group, seq_dim) |
|
|
| return output_tensor |
|
|
|
|
|
|
|
|
| class ReduceFunction(torch.autograd.Function): |
|
|
| @staticmethod |
| def forward(ctx, input, process_group): |
| ctx.cp_group = process_group |
|
|
| output = input.detach().clone() |
|
|
| dist.all_reduce(output, group=ctx.cp_group) |
|
|
| return output |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| |
| grad_input = grad_output.detach().clone() |
|
|
| return grad_input, None |
| |
|
|
|
|
| class ReplicateFunction(torch.autograd.Function): |
|
|
| @staticmethod |
| def forward(ctx, input, process_group): |
| ctx.cp_group = process_group |
|
|
| output = input.detach().clone() |
|
|
|
|
| return output |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| |
| grad_input = grad_output.detach().clone() |
|
|
| dist.all_reduce(grad_input, group=ctx.cp_group) |
|
|
|
|
| return grad_input, None |
|
|
|
|
| def reduce_cp(partial_sum, partial_square_sum): |
|
|
| cp_process_group = get_cp_group() |
| |
| all_sum = ReduceFunction.apply(partial_sum, cp_process_group) |
| all_square_sum = ReduceFunction.apply(partial_square_sum, cp_process_group) |
|
|
| return all_sum, all_square_sum |
|
|
|
|
| def replicate_cp(all_mean, all_var): |
|
|
| cp_process_group = get_cp_group() |
| |
| all_mean = ReplicateFunction.apply(all_mean, cp_process_group) |
| all_var = ReplicateFunction.apply(all_var, cp_process_group) |
|
|
| return all_mean, all_var |
|
|
|
|
|
|
| def _all_to_all_func(input_, world_size, group, scatter_dim, gather_dim): |
| input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)] |
| output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] |
| dist.all_to_all(output_list, input_list, group=group) |
| return torch.cat(output_list, dim=gather_dim).contiguous() |
|
|
|
|
| class _AllToAll(torch.autograd.Function): |
| """All-to-all communication. |
| |
| Args: |
| input_: input matrix |
| process_group: communication group |
| scatter_dim: scatter dimension |
| gather_dim: gather dimension |
| """ |
|
|
| @staticmethod |
| def forward(ctx, input_, process_group, scatter_dim, gather_dim): |
| ctx.process_group = process_group |
| ctx.scatter_dim = scatter_dim |
| ctx.gather_dim = gather_dim |
| world_size = dist.get_world_size(process_group) |
|
|
| return _all_to_all_func(input_, world_size, process_group, scatter_dim, gather_dim) |
|
|
| @staticmethod |
| def backward(ctx, *grad_output): |
| process_group = ctx.process_group |
| scatter_dim = ctx.gather_dim |
| gather_dim = ctx.scatter_dim |
| return_grad = _AllToAll.apply(*grad_output, process_group, scatter_dim, gather_dim) |
| return (return_grad, None, None, None) |
|
|
|
|
| def all_to_all_with_pad( |
| input_: torch.Tensor, |
| process_group: dist.ProcessGroup, |
| scatter_dim: int = 2, |
| gather_dim: int = 1, |
| scatter_pad: int = 0, |
| gather_pad: int = 0, |
| ): |
| if scatter_pad > 0: |
| pad_shape = list(input_.shape) |
| pad_shape[scatter_dim] = scatter_pad |
| pad_tensor = torch.zeros(pad_shape, device=input_.device, dtype=input_.dtype) |
| input_ = torch.cat([input_, pad_tensor], dim=scatter_dim) |
|
|
| assert ( |
| input_.shape[scatter_dim] % dist.get_world_size(process_group) == 0 |
| ), f"Dimension to scatter ({input_.shape[scatter_dim]}) is not divisible by world size ({dist.get_world_size(process_group)})" |
| input_ = _AllToAll.apply(input_, process_group, scatter_dim, gather_dim) |
|
|
| if gather_pad > 0: |
| input_ = input_.narrow(gather_dim, 0, input_.size(gather_dim) - gather_pad) |
|
|
| return input_ |
|
|
|
|
| def dynamic_switch(x, scatter_dim, gather_dim): |
|
|
| scatter_pad = 0 |
| gather_pad = 0 |
| cp_process_group = get_cp_group() |
|
|
| x = all_to_all_with_pad( |
| x, |
| cp_process_group, |
| scatter_dim=scatter_dim, |
| gather_dim=gather_dim, |
| scatter_pad=scatter_pad, |
| gather_pad=gather_pad, |
| ) |
| return x |
|
|
