diff --git a/bmtrain/__init__.py b/bmtrain/__init__.py index 3e025846..25c72052 100644 --- a/bmtrain/__init__.py +++ b/bmtrain/__init__.py @@ -10,7 +10,7 @@ from .block_layer import CheckpointBlock, TransformerBlockList from .backward import optim_step from .wrapper import BMTrainModelWrapper - +from .pipe_layer import PipelineTransformerBlockList from . import debug from .store import save, load diff --git a/bmtrain/benchmark/__init__.py b/bmtrain/benchmark/__init__.py index 237cfca4..571d621f 100644 --- a/bmtrain/benchmark/__init__.py +++ b/bmtrain/benchmark/__init__.py @@ -1,2 +1,3 @@ from .all_gather import all_gather -from .reduce_scatter import reduce_scatter \ No newline at end of file +from .reduce_scatter import reduce_scatter +from .send_recv import send_recv \ No newline at end of file diff --git a/bmtrain/benchmark/send_recv.py b/bmtrain/benchmark/send_recv.py new file mode 100644 index 00000000..e3c971e4 --- /dev/null +++ b/bmtrain/benchmark/send_recv.py @@ -0,0 +1,31 @@ +from .. import nccl +from .shape import SHAPES +from ..global_var import config +from ..utils import print_rank +from .utils import format_size +import torch +def send_recv(): + current_stream = torch.cuda.current_stream() + for shape in SHAPES: + send_size = shape + + send_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" ) + recv_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" ) + + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + + current_stream.record_event(start_evt) + nccl.groupStart() + if config['rank'] in [0,2,4,6]: + nccl.send(send_buffer.storage(), config['rank']+1, config['comm']) + else: + nccl.recv(recv_buffer.storage(), config['rank']-1, config['comm']) + nccl.groupEnd() + current_stream.record_event(end_evt) + current_stream.synchronize() + time_usage = start_evt.elapsed_time(end_evt) + + bw = shape / 1024 / 1024 / 1024 * 1000 / time_usage + print_rank("Send Recv:\tsize {}\ttime: {:4.3f}\tbw: {:2.6f} GB/s".format(format_size(send_size), time_usage, bw)) + diff --git a/bmtrain/block_layer.py b/bmtrain/block_layer.py index 46c47572..138e5d08 100644 --- a/bmtrain/block_layer.py +++ b/bmtrain/block_layer.py @@ -1,6 +1,6 @@ -from typing import Dict, Iterable, Iterator, Union - +from typing import Dict, Iterable, Iterator, Union, List +from .utils import round_up from .global_var import config import torch from . import nccl @@ -10,8 +10,6 @@ from . import debug import copy -def round_up(x, d): - return (x + d - 1) // d * d # the flag is used to control the zero level , 0 means normal zero3 , 1 means forward without release parameter ,2 means backward without gather parameter class OpCheckpointBlock(torch.autograd.Function): @@ -120,16 +118,19 @@ def backward(ctx, *grad_outputs): return (None, None, None, None) + tuple(grads) class CheckpointBlockContext: - def __init__(self ,block : 'CheckpointBlock',ctx_dict : dict = None, flag : int = 0) -> None: + def __init__(self, block : 'CheckpointBlock', ctx_dict : dict = None, flag : int = 0, pipe = False) -> None: self.block = block - self.ctx_dict=ctx_dict + self.ctx_dict = ctx_dict self._param_buffer = {} self._grad_buffer = {} self._param_tensor = {} self._grad_tensor = {} self.flag = flag self._need_release = False - + if pipe: + self.comm = config["zero_comm"] + else: + self.comm = config["comm"] def enter(self): """ gather parameters @@ -146,15 +147,15 @@ def enter(self): assert self.block._storage_params[kw].is_cuda assert kw not in self._grad_buffer assert kw not in self._param_buffer - local_param = self.block._storage_params[kw] - + local_param = self.block._storage_params[kw] + storage_type = local_param.storage_type() if self.flag != 2: - self._param_buffer[kw] = storage_type(val["partition_size"] * config["world_size"]) + self._param_buffer[kw] = storage_type(val["partition_size"] * val["world_size"]) self._param_tensor[kw] = torch.tensor([], dtype=self._param_buffer[kw].dtype, device=self._param_buffer[kw].device).set_(self._param_buffer[kw]) if requires_grad and local_param.requires_grad: - self._grad_buffer[kw] = storage_type(val["partition_size"] * config["world_size"]) + self._grad_buffer[kw] = storage_type(val["partition_size"] * val["world_size"]) self._grad_tensor[kw] = torch.tensor([], dtype=self._grad_buffer[kw].dtype, device=self._grad_buffer[kw].device).set_(self._grad_buffer[kw]).zero_() if self.flag != 2: nccl.groupStart() @@ -162,7 +163,7 @@ def enter(self): nccl.allGather( self.block._storage_params[kw].storage(), self._param_buffer[kw], - config["comm"] + self.comm ) nccl.groupEnd() @@ -177,25 +178,25 @@ def enter(self): self._grad_tensor[kw].record_stream(current_stream) # update parameters in block - for param in self.block._param_info: - kw_name = param["kw_name"] - offset = param["offset"] - shape = param["shape"] + for param in self.block._param_info: + kw_name = param["kw_name"] + offset = param["offset"] + shape = param["shape"] - if self.flag != 2: - dtype = self._param_buffer[kw_name].dtype - device = self._param_buffer[kw_name].device - param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape) - else: - dtype = param["parameter"].data.dtype - device = param["parameter"].data.device - param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape) + if self.flag != 2: + dtype = self._param_buffer[kw_name].dtype + device = self._param_buffer[kw_name].device + param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape) + else: + dtype = param["parameter"].data.dtype + device = param["parameter"].data.device + param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape) - if requires_grad and kw_name in self._grad_buffer: - param["parameter"].requires_grad_(True) - param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape) - else: - param["parameter"].requires_grad_(False) + if requires_grad and kw_name in self._grad_buffer: + param["parameter"].requires_grad_(True) + param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape) + else: + param["parameter"].requires_grad_(False) def __enter__(self): self.enter() @@ -209,7 +210,6 @@ def exit(self): return self._need_release = False self.block._ready = False - requires_grad = torch.is_grad_enabled() if requires_grad: for kw, val in self.block._storage_info.items(): @@ -237,7 +237,7 @@ def exit(self): self._grad_buffer[kw], local_param.grad.storage(), "sum", - config["comm"] + self.comm ) nccl.groupEnd() @@ -355,7 +355,8 @@ def __init__(self, inner_module : torch.nn.Module): offsets = {} # intialize storage buffers for kw, val in self._storage_info.items(): - partition_size = round_up(val["total"], config['world_size']) // config['world_size'] + val["world_size"] = config["world_size"] + partition_size = round_up(val["total"], val["world_size"]) // val["world_size"] val["partition_size"] = partition_size val["begin"] = config['rank'] * partition_size val["end"] = (config['rank'] + 1) * partition_size @@ -538,7 +539,6 @@ def init_parameters(self): # initialzie here tmp_tensor = torch.empty(it["shape"], device=param.device, dtype=param.dtype) param._init_method(tmp_tensor) - param_st = it["offset"] param_end = it["offset"] + it["size"] kw_name = it["kw_name"] @@ -604,7 +604,6 @@ def named_modules(self, memo = None, prefix: str = '', remove_duplicate: bool = 1 -> ('0', Linear(in_features=2, out_features=2, bias=True)) """ - # print("here in named_modules hahaha") if memo is None: memo = set() @@ -654,7 +653,7 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s with torch.no_grad(): for i in range(len(self)): if save_list[i][0] == i: - layer_inputs.append(hidden_state) + layer_inputs.append(hidden_state.detach()) cuda_rng_state.append( torch.cuda.get_rng_state() ) if config['zero_level']==2: flag = 1 @@ -674,13 +673,20 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s ctx.layer_inspector = layer_inspector ctx.cuda_rng_state = cuda_rng_state - ctx.save_for_backward(*layer_inputs, *tensors) - return hidden_state + + if self.return_hidden_states: + middle_hiddens = layer_inputs + for mid in middle_hiddens: + mid.requires_grad_() + middle_hiddens = torch.stack(middle_hiddens, dim=0) + return hidden_state, middle_hiddens + else: + return hidden_state, None @staticmethod - def backward(ctx, grad_hidden_state : torch.Tensor): + def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle: List[torch.Tensor]): def exit_prev(prev_ctx, prev_grad): if prev_ctx is not None: if prev_grad: @@ -738,7 +744,7 @@ def exit_prev(prev_ctx, prev_grad): ctx.save_list[j+1][0] = j+1 torch.cuda.set_rng_state(ctx.cuda_rng_state[i]) - ipt = layer_inputs[ctx.save_list[i][1]].detach().requires_grad_() + ipt = layer_inputs[ctx.save_list[i][1]].requires_grad_() if config['zero_level'] == 2: flag = 2 else: @@ -759,12 +765,15 @@ def exit_prev(prev_ctx, prev_grad): assert it["group"] == ctx.layer_inspector[i][j]["group"], "Backward step changed" # change the tensor in placeholder + ctx.layer_inspector[i][j]["requires_grad"] = it["requires_grad"] ctx.layer_inspector[i][j]["tensor"] = it["tensor"] torch.autograd.backward( [output], [grad_hidden_state] ) grad_hidden_state = ipt.grad + if grad_middle is not None: + grad_hidden_state = grad_hidden_state + grad_middle[i] exit_prev(prev_ctx, prev_grad) @@ -836,6 +845,11 @@ def __iter__(self) -> Iterator[CheckpointBlock]: def __getitem__(self, index: Union[int, str]) -> CheckpointBlock: return self._modules[str(index)] - def forward(self, hidden_state, *args): + def forward(self, hidden_state, *args, return_hidden_states = False): + self.return_hidden_states = return_hidden_states placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled()) - return OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args) + last_hidden, middle_hiddens = OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args) + if return_hidden_states: + return last_hidden, middle_hiddens + else: + return last_hidden diff --git a/bmtrain/distributed/__init__.py b/bmtrain/distributed/__init__.py index 0cb1fd5a..9dc64bb8 100644 --- a/bmtrain/distributed/__init__.py +++ b/bmtrain/distributed/__init__.py @@ -1 +1 @@ -from .ops import all_gather, all_reduce +from .ops import all_gather, all_reduce, broadcast, recv_activations, send_activations diff --git a/bmtrain/distributed/ops.py b/bmtrain/distributed/ops.py index 9cfff1b9..6124622d 100644 --- a/bmtrain/distributed/ops.py +++ b/bmtrain/distributed/ops.py @@ -1,29 +1,82 @@ import torch from ..global_var import config -from ..nccl import allGather as ncclAllGather +from ..nccl import allGather as ncclAllGather, recv from ..nccl import allReduce as ncclAllReduce - +from ..nccl import broadcast as ncclBroadcast +from ..nccl import send as ncclSend +from ..nccl import recv as ncclRecv +from ..nccl import commCount,commRank,NCCLCommunicator +DTYPE_LIST = [ + torch.float64, + torch.float32, + torch.float16, + torch.int64, + torch.int32, + torch.int16, + torch.int8, + torch.bfloat16, + torch.bool +] +def send_activations(hidden_state, next_rank, comm): + send_meta(hidden_state, next_rank, comm) + ncclSend(hidden_state.storage(), next_rank, comm) +def recv_activations(prev_rank, comm): + dtype, shape = recv_meta(prev_rank, comm) + hidden_state = torch.empty(shape, dtype=dtype, device="cuda") + ncclRecv(hidden_state.storage(), prev_rank, comm) + return hidden_state +def send_meta(x, next_rank, comm): + meta = [len(x.size()), DTYPE_LIST.index(x.dtype)] + list(x.size()) + meta_data = torch.tensor(data=meta, device=x.device, dtype=torch.long) + ncclSend(meta_data.storage(), next_rank, comm) +def recv_meta(prev_rank, comm): + meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.long) + ncclRecv(meta_data.storage(), prev_rank, comm) + n_dims = meta_data[0].item() + dtype = DTYPE_LIST[meta_data[1].item()] + shape = meta_data[2:n_dims+2].tolist() + return dtype,shape +class OpBroadcast(torch.autograd.Function): + @staticmethod + def forward(ctx, src, root, comm = None): + if comm is None: + comm = config["comm"] + ctx.comm = comm + outputs = torch.empty_like(src, dtype = src.dtype, device = src.device) + ncclBroadcast(src.storage(), outputs.storage(), root, comm) + return outputs + @staticmethod + def backward(ctx, grad_output): + res = all_reduce(grad_output, "sum", ctx.comm) + return res, None, None +def broadcast(src, root, comm=None): + if not config["initialized"]: + raise RuntimeError("BMTrain is not initialized") + return OpBroadcast.apply(src, root, comm) class OpAllGather(torch.autograd.Function): @staticmethod - def forward(ctx, input : torch.Tensor): + def forward(ctx, input : torch.Tensor, comm = None): + if comm is None: + comm = config["comm"] + world_size = commCount(comm) if not input.is_contiguous(): input = input.contiguous() if input.storage_offset() != 0 or input.storage().size() != input.numel(): input = input.clone() - output = torch.empty( (config['world_size'],) + input.size(), dtype=input.dtype, device=input.device) - + output = torch.empty( (world_size,) + input.size(), dtype=input.dtype, device=input.device) + ctx.comm = comm ncclAllGather( input.storage(), output.storage(), - config['comm'] + comm ) return output @staticmethod def backward(ctx, grad_output): - return grad_output[config['rank']] + return grad_output[commRank(ctx.comm)], None -def all_gather(x : torch.Tensor): +def all_gather(x : torch.Tensor, comm = None): """Gathers the input tensor from all processes. Args: @@ -36,11 +89,14 @@ def all_gather(x : torch.Tensor): raise RuntimeError("BMTrain is not initialized") assert x.is_cuda - return OpAllGather.apply(x) + return OpAllGather.apply(x, comm) class OpAllReduce(torch.autograd.Function): @staticmethod - def forward(ctx, input : torch.Tensor, op : str): + def forward(ctx, input : torch.Tensor, op : str, comm : NCCLCommunicator = None): + if comm is None: + comm = config["comm"] + ctx.comm = comm if not input.is_contiguous(): input = input.contiguous() if input.storage_offset() != 0 or input.storage().size() != input.numel(): @@ -51,10 +107,10 @@ def forward(ctx, input : torch.Tensor, op : str): input.storage(), output.storage(), op, - config['comm'] + comm ) ctx.op = op - + if op in ["sum", "avg"]: pass elif op in ["max", "min"]: @@ -66,15 +122,15 @@ def forward(ctx, input : torch.Tensor, op : str): @staticmethod def backward(ctx, grad_output): if ctx.op == "sum": - return grad_output, None + return grad_output, None, None elif ctx.op == "avg": - return grad_output / config['world_size'], None + return grad_output / commCount(ctx.comm), None, None elif ctx.op in ["max", "min"]: - return torch.masked_fill(grad_output, ctx.saved_tensors[0], 0), None + return torch.masked_fill(grad_output, ctx.saved_tensors[0], 0), None, None else: - return grad_output * ctx.saved_tensors[0], None + return grad_output * ctx.saved_tensors[0], None, None -def all_reduce(x : torch.Tensor, op : str = "sum"): +def all_reduce(x : torch.Tensor, op : str = "sum", comm = None): """Reduces the input tensor from all processes. Args: @@ -89,7 +145,7 @@ def all_reduce(x : torch.Tensor, op : str = "sum"): raise RuntimeError("BMTrain is not initialized") assert x.is_cuda - return OpAllReduce.apply(x, op) + return OpAllReduce.apply(x, op, comm) diff --git a/bmtrain/global_var.py b/bmtrain/global_var.py index 67a54cf7..137fa9cd 100644 --- a/bmtrain/global_var.py +++ b/bmtrain/global_var.py @@ -6,13 +6,15 @@ class ConfigMap(TypedDict): world_size : int local_size : int zero_level : int + pipe_size : int + num_micro_batches : int calc_stream : torch.cuda.Stream load_stream : torch.cuda.Stream load_event : torch.cuda.Event barrier_stream : torch.cuda.Stream loss_scale_factor : float loss_scale_steps : int - + topology : 'topology' gradient_inspect : bool initialized : bool diff --git a/bmtrain/init.py b/bmtrain/init.py index 42186923..dfc87e14 100644 --- a/bmtrain/init.py +++ b/bmtrain/init.py @@ -13,6 +13,8 @@ def init_distributed( loss_scale_factor : float = 2, loss_scale_steps : int = 1024, zero_level: int = 3, + pipe_size: int = -1, + num_micro_batches: int = None, ): """Initialize distributed training. This function will initialize the distributed training, set the random seed and global configurations. @@ -40,7 +42,7 @@ def init_distributed( """ torch.backends.cudnn.enabled = False - + local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) @@ -52,11 +54,11 @@ def init_distributed( ) store, rank, world_size = next(rendezvous_iterator) store.set_timeout(timeout) - store = dist.PrefixStore("bmtrain", store) torch.cuda.set_device(local_rank) - config["initialized"] = True + config["pipe_size"] = pipe_size if pipe_size > 0 else 1 + config["pipe_enabled"] = pipe_size > 0 config["local_rank"] = local_rank config["local_size"] = local_size config["rank"] = rank @@ -68,7 +70,8 @@ def init_distributed( config["zero_level"] = zero_level config["loss_scale_factor"] = loss_scale_factor if loss_scale_factor > 1 else 1 / loss_scale_factor config["loss_scale_steps"] = loss_scale_steps - + config["topology"] = topology(config) + config["zero_rank"] = config["topology"].get_group_rank("zero") if pipe_size > 1 else config['rank'] cpus_this_worker = None all_available_cpus = sorted(list(os.sched_getaffinity(0))) @@ -97,7 +100,21 @@ def init_distributed( unique_id = bytes.fromhex(store.get("BMTRAIN_UNIQUE_ID").decode()) config['comm'] = nccl.commInitRank(unique_id, world_size, rank) - + if config['pipe_enabled']: + config["micros"] = num_micro_batches if num_micro_batches else config["pipe_size"] + topo = config['topology'] + if topo.stage_id == 0: + unique_id = nccl.getUniqueId() + store.set(f"PIPE_UNIQUE_ID{topo.pipe_idx}", unique_id.hex()) + unique_id = bytes.fromhex(store.get(f"PIPE_UNIQUE_ID{topo.pipe_idx}").decode()) + config ['pipe_comm'] = nccl.commInitRank(unique_id, pipe_size, topo.stage_id) + if topo.zero_id == 0: + unique_id = nccl.getUniqueId() + store.set(f"ZERO_UNIQUE_ID{topo.zero_idx}", unique_id.hex() ) + unique_id = bytes.fromhex(store.get(f"ZERO_UNIQUE_ID{topo.zero_idx}").decode()) + config ['zero_comm'] = nccl.commInitRank(unique_id, world_size//pipe_size, topo.zero_id) + else: + config['zero_comm'] = config['comm'] for i in range(world_size): if i == rank: print_dict("Initialization", { @@ -110,6 +127,38 @@ def init_distributed( "cpus": cpus_this_worker }) synchronize() +class topology: + def __init__(self,config): + # pipe_idx is the idx of the pipeline in the group + self.rank = config['rank'] + pp_size = config["pipe_size"] + world_size = config["world_size"] + assert world_size % pp_size == 0, "The nums of GPUs must be divisible by the pipeline parallel size" + + dp_size = world_size // pp_size + topo=torch.tensor(range(dp_size*pp_size),dtype=torch.int,device='cuda') + topo=topo.view(pp_size,dp_size) + self.pp_group=topo.transpose(0,1).reshape(-1,pp_size) + self.dp_group=topo + self.stage_id = (self.pp_group == self.rank).nonzero()[0,-1].item() + self.stages = config['pipe_size'] + self.pipe_idx = (self.pp_group == self.rank).nonzero()[0, 0].item() # x axes + self.zero_id = self.pipe_idx + self.zero_idx = self.stage_id + self.next_rank = self.stage_id+1 if self.stage_id < config['pipe_size'] - 1 else -1 + self.prev_rank = self.stage_id-1 if self.stage_id > 0 else -1 + self.tails = self.pp_group[self.pipe_idx, self.stage_id:].tolist() + self.heads = self.pp_group[self.pipe_idx, :self.stage_id + 1].tolist() + def get_group_id(self,group_name): + if group_name == "pipe": + return self.pipe_idx + elif group_name == "zero": + return self.zero_idx + def get_group_rank(self,group_name): + if group_name == "pipe": + return self.stage_id + elif group_name == "zero": + return self.zero_id def is_initialized() -> bool: return config["initialized"] diff --git a/bmtrain/inspect/model.py b/bmtrain/inspect/model.py index 6bbe57c1..158c2de4 100644 --- a/bmtrain/inspect/model.py +++ b/bmtrain/inspect/model.py @@ -1,4 +1,6 @@ import torch +from ..store import broadcast_object +from ..pipe_layer import PipelineTransformerBlockList from ..block_layer import CheckpointBlock from ..parameter import DistributedParameter from .. import nccl @@ -30,6 +32,91 @@ def _gather_value(value : torch.Tensor, partition_size, origin_size): return output_tensor +def inspect_pipeline_transformer_block_list(pipe_model: PipelineTransformerBlockList, param_name : str, _prefix : str = ''): + ret = [] + for name, model in pipe_model._modules.items(): + idx = int(name) + prefix = _prefix + name + '.' + + # fast check + pass_fast_check = False + for param in model._param_info: + abs_name = prefix + param["name"] + if fnmatch.fnmatch(abs_name, param_name): + pass_fast_check = True + break + if not pass_fast_check: + continue + + if idx in pipe_model.layer_ids: + _param_buffer = {} + _grad_buffer = {} + for kw, val in model._storage_info.items(): + storage_type = model._storage_params[kw].storage_type() + + _param_buffer[kw] = storage_type(val["partition_size"] * val['world_size']) + if model._storage_params[kw].grad is not None: + _grad_buffer[kw] = storage_type(val["partition_size"] * val['world_size']) + + nccl.groupStart() + for kw, val in model._storage_info.items(): + nccl.allGather( + model._storage_params[kw].storage(), + _param_buffer[kw], + config["zero_comm"] + ) + if model._storage_params[kw].grad is not None: + nccl.allGather( + model._storage_params[kw].grad.storage(), + _grad_buffer[kw], + config["zero_comm"] + ) + + nccl.groupEnd() + for param in model._param_info: + abs_name = prefix + param["name"] + if fnmatch.fnmatch(abs_name, param_name): + kw_name = param["kw_name"] + dtype = _param_buffer[kw_name].dtype + device = _param_buffer[kw_name].device + offset = param["offset"] + shape = param["shape"] + p = torch.tensor([], dtype=dtype, device=device).set_(_param_buffer[kw_name], offset, shape) + if kw_name in _grad_buffer: + g = torch.tensor([], dtype=dtype, device=device).set_(_grad_buffer[kw_name], offset, shape) + info = { + "name": abs_name, + "shape": tuple(shape), + "std": p.std().cpu().item(), + "mean": p.mean().cpu().item(), + "grad_std": g.std().cpu().item(), + "grad_mean": g.mean().cpu().item(), + "max": p.max().cpu().item(), + "min": p.min().cpu().item(), + } + else: + info = { + "name": abs_name, + "shape": tuple(shape), + "std": p.std().cpu().item(), + "mean": p.mean().cpu().item(), + "grad_std": None, + "grad_mean": None, + "max": p.max().cpu().item(), + "min": p.min().cpu().item(), + } + broadcast_object(info, config["pipe_comm"], pipe_model.get_stage_by_layer_id(idx)) + ret.append(info) + else: + for param in model._param_info: + abs_name = prefix + param["name"] + if fnmatch.fnmatch(abs_name, param_name): + info = broadcast_object({}, config["pipe_comm"], pipe_model.get_stage_by_layer_id(idx)) + ret.append(info) + + return ret + + def inspect_checkpoint_block(model : CheckpointBlock, param_name : str, prefix : str = ''): # fast check pass_fast_check = False @@ -121,7 +208,9 @@ def inspect_model(model : torch.nn.Module, param_name : str, prefix : str = ''): ... """ - if isinstance(model, CheckpointBlock): + if isinstance(model, PipelineTransformerBlockList): + return inspect_pipeline_transformer_block_list(model, param_name, prefix) + elif isinstance(model, CheckpointBlock): return inspect_checkpoint_block(model, param_name, prefix) else: ret = [] diff --git a/bmtrain/inspect/tensor.py b/bmtrain/inspect/tensor.py index 5d7030f3..7d56ad7f 100644 --- a/bmtrain/inspect/tensor.py +++ b/bmtrain/inspect/tensor.py @@ -3,6 +3,7 @@ from .. import debug from .. import nccl from ..global_var import config +from ..store import broadcast_object import math @@ -16,62 +17,88 @@ def __init__(self): self._summary = [] def _set_summary(self, summary): - self._summary = [] - group_cnt = {} - for item in summary: - group = item["group"] - name = item["name"] - if group not in group_cnt: - group_cnt[group] = {} - if name not in group_cnt[group]: - group_cnt[group][name] = 0 - group_cnt[group][name] += 1 - - group_idx = {} for item in summary: - group = item["group"] - name = item["name"] - if group not in group_idx: - group_idx[group] = {} - if name not in group_idx[group]: - group_idx[group][name] = 0 - - group_name_prefix = f"{group}." if group is not None else "" - if group_cnt[group][name] > 1: - item["name"] = f"{group_name_prefix}{group_idx[group][name]}.{name}" + item['prefix'] = "" if item["group"] is None else f'{item["group"]}.' + + self._summary = [] + + kw_cnt = {} + i = 0 + while i < len(summary): + item = summary[i] + if item["inside_pipe"] is not None: + assert item["inside_pipe"]["st"] + pipe_cnt = {} + j = i + while j < len(summary): + item = summary[j] + kw = f'{item["prefix"]}{item["name"]}' + + assert item["inside_pipe"] is not None + stage_id = item["inside_pipe"]["stage_id"] + stages = item["inside_pipe"]["stages"] + st = item["inside_pipe"]["st"] + ed = item["inside_pipe"]["ed"] + + if kw not in pipe_cnt: + pipe_cnt[kw] = 0 + pipe_cnt[kw] += 1 + + j += 1 + if ed: + break + + for stage in range(stages): + if stage_id == stage: + broadcast_object(pipe_cnt, config["pipe_comm"], src = stage) + for k in range(i, j): + item = summary[k] + kw = f'{item["prefix"]}{item["name"]}' + if kw not in kw_cnt: + kw_cnt[kw] = 0 + item["name"] = f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}' + item["shape"] = ((item["shape"][0] * config['micros'],) + item["shape"][1:]) + if item['tensor'].grad is not None: + grad = torch.cat([summary[k+m*(j-i)]['tensor'].grad for m in range(config['micros'])], dim=0) + else: + grad = None + item["tensor"] = torch.cat([summary[k+m*(j-i)]['tensor'] for m in range(config['micros'])], dim=0) + item["tensor"].grad = grad + self._summary.append(item) + kw_cnt[kw] += 1 + else: + cnt = broadcast_object({}, config["pipe_comm"], src = stage) + for kw, val in cnt.items(): + if kw not in kw_cnt: + kw_cnt[kw] = 0 + for _ in range(val): + self._summary.append({ + "name": f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}', + "group": None, + "requires_grad": None, + "min": None, + "max": None, + "mean": None, + "std": None, + "shape": None, + "grad_mean" : None, + "grad_std" : None, + "tensor": None, + "inside_pipe": {"stage_id": stage}, + }) + kw_cnt[kw] += 1 + + i = i + config['micros'] * (j - i) else: - item["name"] = f"{group_name_prefix}{name}" - if not item["requires_grad"]: - x = item["tensor"] - info = torch.empty(2, dtype=x.dtype, device=x.device) - info[0] = x.mean() - info[1] = x.var() - nccl.allReduce( - info.storage(), - info.storage(), - "avg", - config['comm'] - ) - x_mean = info[0].cpu().item() - x_std = math.sqrt(info[1].cpu().item()) - item["mean"] = x_mean - item["std"] = x_std + kw = f'{item["prefix"]}{item["name"]}' + if kw not in kw_cnt: + kw_cnt[kw] = 0 + item["name"] = f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}' + self._summary.append(item) + kw_cnt[kw] += 1 + i = i + 1 - info[0] = x.max() - info[1] = -x.min() - nccl.allReduce( - info.storage(), - info.storage(), - 'max', - config['comm'] - ) - x_max = info[0].cpu().item() - x_min = - info[1].cpu().item() - item["max"] = x_max - item["min"] = x_min - self._summary.append(item) - group_idx[group][name] += 1 def get_summary(self): """Get the summary of the tensors recorded by `record_tensor`. @@ -90,78 +117,71 @@ def get_summary(self): **Note:** This method must be called outside of the `with` block. """ - nw_summary = [] + ret = [] for item in self._summary: - if item["requires_grad"] and item["tensor"].grad is not None: - x = item["tensor"] - info = torch.empty(4, dtype=x.dtype, device=x.device) - info[0] = x.mean() - info[1] = x.var() - info[2] = x.grad.mean() - info[3] = x.grad.var() - nccl.allReduce( - info.storage(), - info.storage(), - "avg", - config['comm'] - ) - x_mean = info[0].cpu().item() - x_std = math.sqrt(info[1].cpu().item()) - grad_mean = info[2].cpu().item() - grad_std = math.sqrt(info[3].cpu().item()) - + comm = config["zero_comm"] if item["inside_pipe"] else config["comm"] + if item["tensor"] is not None: + if not item["requires_grad"] or item["tensor"].grad is None: + x = item["tensor"] + info = torch.empty(2, dtype=x.dtype, device=x.device) + info[0] = x.mean() + info[1] = x.var() + nccl.allReduce( + info.storage(), + info.storage(), + "avg", + comm + ) + x_mean = info[0].cpu().item() + x_std = math.sqrt(info[1].cpu().item()) + grad_mean = None + grad_std = None + else: + x = item["tensor"] + info = torch.empty(4, dtype=x.dtype, device=x.device) + info[0] = x.mean() + info[1] = x.var() + info[2] = x.grad.mean() + info[3] = x.grad.var() + nccl.allReduce( + info.storage(), + info.storage(), + "avg", + comm + ) + x_mean = info[0].cpu().item() + x_std = math.sqrt(info[1].cpu().item()) + grad_mean = info[2].cpu().item() + grad_std = math.sqrt(info[3].cpu().item()) + info[0] = x.max() info[1] = -x.min() nccl.allReduce( info.storage(), info.storage(), 'max', - config['comm'] + comm ) x_max = info[0].cpu().item() x_min = - info[1].cpu().item() - nw_summary.append({ + summary = { "name": item["name"], - "group": item["group"], - "requires_grad" : False, "min": x_min, "max": x_max, "mean": x_mean, "std": x_std, "shape": item["shape"], "grad_mean" : grad_mean, - "grad_std" : grad_std, - "tensor": item["tensor"] - }) - else: - nw_summary.append(item) - - ret = [] - self._summary = nw_summary - for item in self._summary: - if item["requires_grad"]: - ret.append({ - "name": item["name"], - "min": item["min"], - "max": item["max"], - "mean": item["mean"], - "std": item["std"], - "shape": item["shape"], - "grad_mean" : None, - "grad_std" : None - }) + "grad_std" : grad_std + } + + if item["inside_pipe"] is not None: + broadcast_object(summary, config["pipe_comm"], item["inside_pipe"]["stage_id"]) else: - ret.append({ - "name": item["name"], - "min": item["min"], - "max": item["max"], - "mean": item["mean"], - "std": item["std"], - "shape": item["shape"], - "grad_mean" : item["grad_mean"], - "grad_std" : item["grad_std"] - }) + summary = broadcast_object({}, config["pipe_comm"], item["inside_pipe"]["stage_id"]) + + ret.append(summary) return ret def get_tensor(self, name : str, group : Optional[str] = None, index : Optional[int] = None) -> torch.Tensor: @@ -268,5 +288,6 @@ def record_tensor(x : torch.Tensor, name : str, group = None, requires_grad = Tr "shape": x_shape, "grad_mean" : None, "grad_std" : None, - "tensor": x + "tensor": x, + "inside_pipe": None, }) diff --git a/bmtrain/nccl/__init__.py b/bmtrain/nccl/__init__.py index c3f0ae78..bccae34b 100644 --- a/bmtrain/nccl/__init__.py +++ b/bmtrain/nccl/__init__.py @@ -38,7 +38,8 @@ def dtype2nccl(dtype : torch.dtype) -> int: torch.float32 : ncclFloat32, torch.float : ncclFloat, torch.float64 : ncclFloat64, - torch.double : ncclDouble + torch.double : ncclDouble, + torch.bool : ncclBool } if dtype not in MAP: raise TypeError("Unsupport dtype %s" % dtype) @@ -83,9 +84,21 @@ def commDestroy(comm : NCCLCommunicator): """ C.ncclCommDestroy(comm.ptr) comm._destroy_ptr() +def commCount(comm : NCCLCommunicator): + """NCCL API: `ncclCommCount `_ + Args: + comm (NCCLCommunicator): NCCL communicator. + """ + return C.ncclCommCount(comm.ptr) ### collective +def commRank(comm : NCCLCommunicator): + """NCCL API: `ncclCommUserRank `_ + Args: + comm (NCCLCommunicator): NCCL communicator. + """ + return C.ncclCommUserRank(comm.ptr) def allReduce( src : torch.storage._StorageBase, dst : torch.storage._StorageBase, @@ -124,7 +137,45 @@ def allReduce( comm.ptr, torch.cuda.current_stream().cuda_stream ) +def send(src : torch.storage._StorageBase, + peer : int, + comm : NCCLCommunicator + ): + """NCCL API: `ncclsend `_ + Args: + src (torch.storage._StorageBase): Source buffer. + peer (int): rank peer needs to call ncclRecv + comm (NCCLCommunicator): NCCL communicator. + """ + + sendbuff = src.data_ptr() + count = src.size() + datatype = dtype2nccl(src.dtype) + C.ncclSend( + sendbuff, + count, + datatype, + peer, + comm.ptr, + torch.cuda.current_stream().cuda_stream + ) +def recv(dst : torch.storage._StorageBase, + peer : int, + comm : NCCLCommunicator + ): + recvbuff = dst.data_ptr() + count = dst.size() + datatype = dtype2nccl(dst.dtype) + C.ncclRecv( + recvbuff, + count, + datatype, + peer, + comm.ptr, + torch.cuda.current_stream().cuda_stream + ) + def broadcast( src : torch.storage._StorageBase, dst : torch.storage._StorageBase, @@ -221,7 +272,6 @@ def allGather( recvbuff = dst.data_ptr() sendcount = src.size() datatype = dtype2nccl(src.dtype) - assert dst.size() % sendcount == 0, "Buffer size not aligned" C.ncclAllGather( sendbuff, diff --git a/bmtrain/nccl/enums.py b/bmtrain/nccl/enums.py index 5e388728..dc2adbd4 100644 --- a/bmtrain/nccl/enums.py +++ b/bmtrain/nccl/enums.py @@ -3,6 +3,7 @@ ncclInt8 = 0 ncclChar = 0 +ncclBool = 0 ncclUint8 = 1 ncclInt32 = 2 ncclInt = 2 @@ -16,7 +17,6 @@ ncclFloat64 = 8 ncclDouble = 8 - ### ncclRedOp_t ncclSum = 0 diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py index 06ee195e..713d3008 100644 --- a/bmtrain/optim/adam.py +++ b/bmtrain/optim/adam.py @@ -152,7 +152,7 @@ def loss_scale(self, loss : torch.Tensor) -> torch.Tensor: """ Backward with loss scale. """ - return loss * (self.scale / config['world_size']) + return loss * (self.scale * config['pipe_size'] / config['world_size']) def load_state_dict(self, state_dict: dict) -> None: r"""Loads the optimizer state. diff --git a/bmtrain/optim/adam_offload.py b/bmtrain/optim/adam_offload.py index 4f106b07..7ae2f706 100644 --- a/bmtrain/optim/adam_offload.py +++ b/bmtrain/optim/adam_offload.py @@ -182,7 +182,7 @@ def loss_scale(self, loss : torch.Tensor) -> torch.Tensor: """ Backward with loss scale. """ - return loss * (self.scale / config['world_size']) + return loss * (self.scale * config['pipe_size'] / config['world_size']) def load_state_dict(self, state_dict: dict) -> None: r"""Loads the optimizer state. diff --git a/bmtrain/optim/clip_grad.py b/bmtrain/optim/clip_grad.py index dcf5f739..92d9a48e 100644 --- a/bmtrain/optim/clip_grad.py +++ b/bmtrain/optim/clip_grad.py @@ -23,18 +23,23 @@ def clip_grad_norm(param_groups, max_norm, scale, norm_type=2, eps=1e-6): >>> bmt.optim.clip_grad_norm(optimizer.param_groups, max_norm=1.0, scale=optimizer.scale, norm_type=2) """ - scale = scale / config['world_size'] - parameters = [p for group in param_groups for p in group['params'] if p.grad is not None] - + scale = scale * config['pipe_size'] / config['world_size'] + grads = [] + parameters = [p for group in param_groups for p in group['params']] + for p in parameters: + if p.grad is not None: + grads.append(p.grad.data) + else: + grads.append(torch.zeros_like(p.data)) if norm_type == 'inf': - total_norm_cuda = max(p.grad.data.abs().max() for p in parameters).detach() + total_norm_cuda = max(g.data.abs().max() for g in grads).detach() nccl.allReduce(total_norm_cuda.storage(), total_norm_cuda.storage(), "max", config["comm"]) total_norm = total_norm_cuda else: norm_type = float(norm_type) total_norm_cuda = torch.cuda.FloatTensor([0]) - for index, p in enumerate(parameters): - param_norm = p.grad.data.float().norm(norm_type) + for index, g in enumerate(grads): + param_norm = g.data.float().norm(norm_type) total_norm_cuda += param_norm ** norm_type nccl.allReduce(total_norm_cuda.storage(), total_norm_cuda.storage(), "sum", config["comm"]) total_norm = total_norm_cuda[0] ** (1. / norm_type) @@ -43,5 +48,6 @@ def clip_grad_norm(param_groups, max_norm, scale, norm_type=2, eps=1e-6): clip_coef = float(max_norm * scale) / (total_norm + eps) if clip_coef < 1: for p in parameters: - p.grad.data.mul_(clip_coef) + if p.grad is not None: + p.grad.data.mul_(clip_coef) return total_norm / scale \ No newline at end of file diff --git a/bmtrain/param_init.py b/bmtrain/param_init.py index 49f00346..c95f90f8 100644 --- a/bmtrain/param_init.py +++ b/bmtrain/param_init.py @@ -1,6 +1,5 @@ from typing import Generator, Iterable, List, Tuple import torch - from .block_layer import CheckpointBlock from .parameter import DistributedParameter from .global_var import config @@ -23,6 +22,7 @@ def init_distributed_parameter(params : Iterable[torch.nn.Parameter]): param._init_method(tmp_tensor) # Pytorch 1.11 changed the API of storage.__getitem__ + # use zero_rank to support pipeline torch.tensor([], dtype=param.dtype, device=param.device).set_(param.storage())[:] = \ torch.tensor([], dtype=param.dtype, device=param.device).set_(tmp_storage)[partition_size * config['rank'] : partition_size * (config['rank'] + 1)] # param.storage().copy_(tmp_storage[partition_size * config['rank'] : partition_size * (config['rank'] + 1)]) diff --git a/bmtrain/pipe_layer.py b/bmtrain/pipe_layer.py new file mode 100644 index 00000000..a408c954 --- /dev/null +++ b/bmtrain/pipe_layer.py @@ -0,0 +1,484 @@ +from collections import OrderedDict +import copy +import torch +import copy +from typing import Dict, Iterable, Iterator, Tuple, Union, List +import torch + +from .distributed import all_gather, broadcast, all_reduce, send_activations, recv_activations +from .global_var import config +from . import nccl +from .checkpointing import ScopedTensorInspectorContext +from . import debug +from .block_layer import CheckpointBlockContext, CheckpointBlock, round_up, _get_param_kw + +class OpMicroForward(torch.autograd.Function): + @staticmethod + def forward(ctx, placeholder, self : 'PipelineTransformerBlockList', layers_dict, save_list, hidden_state, *args): + with PipeContext(self, hidden_state) as pipe_input: + hidden_state = pipe_input[0].detach() + tensors = [arg if torch.is_tensor(arg) else None for arg in args] + others = [arg if not torch.is_tensor(arg) else None for arg in args] + ctx.nontensor_inputs = others + ctx.self = self + ctx.save_list = copy.deepcopy(save_list) + ctx.num_save_needed = save_list[-1][1]+1 + ctx.layers_dict = layers_dict + layer_inputs = [] + layer_inspector = [] + cuda_rng_state = [] + with torch.no_grad(): + for idx,layer_id in enumerate(self.layer_ids): + if save_list[idx][0] == idx: + layer_inputs.append(hidden_state.detach()) + cuda_rng_state.append( torch.cuda.get_rng_state() ) + if not ctx.layers_dict[idx]: + block_ctx = CheckpointBlockContext(self._modules[str(layer_id)], ctx.layers_dict[idx], 1, pipe=True) + else: + block_ctx = CheckpointBlockContext(self._modules[str(layer_id)], ctx.layers_dict[idx], 2, pipe=True) + # gather parameter on load stream + block_ctx.enter() + # call inner module directly + with ScopedTensorInspectorContext() as inspector: + hidden_state = self._modules[str(layer_id)]._module._call_impl(hidden_state, *args) + for ith, it in enumerate(inspector.hidden_states): + it["shape"] = ((it["shape"][0] // config['pipe_size'],) + it["shape"][1:]) + it["inside_pipe"] = { + "stage_id": self.stage_id, + "stages": self.stages, + "st": (layer_id==self.layer_ids[0] and ith==0), + "ed": (layer_id==self.layer_ids[-1] and ith==len(inspector.hidden_states)-1), + } + debug.append("_inspect_hidden_states", it) + layer_inspector.append(inspector.hidden_states) + block_ctx.exit() + + ctx.layer_inspector = layer_inspector + ctx.cuda_rng_state = cuda_rng_state + + ctx.save_for_backward(*layer_inputs, *tensors) + pipe_input[0] = hidden_state + if self.return_hidden_states: + middle_hiddens = layer_inputs + for mid in middle_hiddens: + mid.requires_grad_() + middle_hiddens = torch.stack(middle_hiddens, dim=0) + return pipe_input[0], middle_hiddens + else: + return pipe_input[0], None + + @staticmethod + def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle : torch.Tensor): + def exit_prev(prev_ctx, prev_grad): + if prev_ctx is not None: + if prev_grad: + with torch.enable_grad(): + prev_ctx.exit() + config["load_stream"].record_event(config["load_event"]) + else: + with torch.no_grad(): + prev_ctx.exit() + config["load_stream"].record_event(config["load_event"]) + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError( + "Checkpointing is not compatible with .grad() or when an `inputs` parameter" + " is passed to .backward(). Please use .backward() and do not pass its `inputs`" + " argument.") + all_inputs = [] + input_requires_grad = [] + + layer_inputs = ctx.saved_tensors[:ctx.num_save_needed] + save_args = ctx.saved_tensors[ctx.num_save_needed:] + for tensor, other in zip(save_args, ctx.nontensor_inputs): + if tensor is None: + all_inputs.append(other) + input_requires_grad.append(False) + else: + # detach for tensor inputs + input_requires_grad.append( tensor.requires_grad ) + nw_tensor = tensor.detach() + nw_tensor.requires_grad = tensor.requires_grad + all_inputs.append(nw_tensor) + with PipeContext(ctx.self, grad_hidden_state, backward=True) as pipe_input: + grad_hidden_state = pipe_input[0] + with torch.random.fork_rng(devices=[torch.cuda.current_device()], enabled=True): + with torch.enable_grad(): + # overlap load and scatter here + prev_ctx = None + prev_grad = False + for idx,layer_id in list(enumerate(ctx.self.layer_ids))[::-1]: + torch.cuda.set_rng_state(ctx.cuda_rng_state[idx]) + ipt = layer_inputs[ctx.save_list[idx][1]].requires_grad_() + block_ctx = CheckpointBlockContext(ctx.self._modules[str(layer_id)], ctx.layers_dict[idx], 2, pipe=True) + block_ctx.enter() + exit_prev(prev_ctx, prev_grad) + prev_ctx = block_ctx + prev_grad = True + + with ScopedTensorInspectorContext() as inspector: + output = ctx.self._modules[str(layer_id)]._module._call_impl(ipt, *all_inputs) + + assert len(ctx.layer_inspector[idx]) == len(inspector.hidden_states), "Backward step changed" + for j, it in enumerate(inspector.hidden_states): + it["shape"] = ((it["shape"][0] // config['pipe_size'],) + it["shape"][1:]) + assert it["name"] == ctx.layer_inspector[idx][j]["name"], "Backward step changed" + assert it["shape"] == ctx.layer_inspector[idx][j]["shape"], "Backward step changed" + assert it["group"] == ctx.layer_inspector[idx][j]["group"], "Backward step changed" + + # change the tensor in placeholder + ctx.layer_inspector[idx][j]["requires_grad"] = it["requires_grad"] + ctx.layer_inspector[idx][j]["tensor"] = it["tensor"] + torch.autograd.backward( + [output], + [grad_hidden_state] + ) + grad_hidden_state = ipt.grad + if grad_middle is not None: + grad_hidden_state = grad_hidden_state + grad_middle[idx] + exit_prev(prev_ctx, prev_grad) + + pipe_input[0] = grad_hidden_state + grads = [] + for inp, requires_grad in zip(all_inputs, input_requires_grad): + if requires_grad: + grads.append(inp.grad) + else: + grads.append(None) + return (None, None, None, None, pipe_input[0]) + tuple(grads) + +class OpPipeTransformerBlockList(torch.autograd.Function): + @staticmethod + def forward(ctx, placeholder, self : 'PipelineTransformerBlockList', save_list, hidden_state, *args): + num_micros = config["micros"] + ctx.self = self + ctx.num_micros = num_micros + layers_dict = [{} for _ in range(len(self))] + args_list = [[] for _ in range(num_micros)] + batch_related = args[-1] + batch_related_origin = [True if i in args[-1] else False for i in range(len(args[:-1]))] + batch_related_rule = [] + args = args[:-1] + batch_size = hidden_state.shape[0] + assert (batch_size * config["pipe_size"]) % num_micros == 0, f'The batch size {(batch_size * config["pipe_size"])} must be divisible by the number of micro_batch {num_micros}' + input_requires_grad = [] + with torch.enable_grad(): + for arg in args: + if torch.is_tensor(arg): + arg_all = all_gather(arg, config['pipe_comm']) + if arg.shape[0] == batch_size: + batch_related_rule.append(True) + arg_all = arg_all.flatten(0, 1).chunk(num_micros, dim=0) + arg_all = [tensor.detach().requires_grad_(arg.requires_grad) for tensor in arg_all] + else: + batch_related_rule.append(False) + # assert num_micros % self.stages == 0, "batch unrelated only support num_micros % stages == 0" + # arg_all = [arg_all[i // (num_micros // self.stages)].detach().requires_grad_(arg.requires_grad) for i in range(num_micros)] + arg_all = [arg_all[0].detach().requires_grad_(arg.requires_grad) for i in range(num_micros)] + input_requires_grad.append(arg.requires_grad) + else: + batch_related_rule.append(False) + arg_all = [arg for _ in range(num_micros)] + input_requires_grad.append(False) + for i in range(num_micros): + args_list[i].append(arg_all[i]) + outputs = [] + if self.return_hidden_states: + middles = [] + hidden_state_list = all_gather(hidden_state, config["pipe_comm"]).flatten(0, 1).detach().requires_grad_() + ctx.hidden_state_list = hidden_state_list + hidden_state_list = hidden_state_list.chunk(num_micros, dim=0) + for micro_idx, (hidden_state, arg) in enumerate(zip(hidden_state_list, args_list)): + placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled()) + output, middle = OpMicroForward.apply(placeholder, self, layers_dict, save_list, hidden_state, *arg) + outputs.append(output) + if self.return_hidden_states: + middles.append(middle) + if len(batch_related) == 0: + ctx.batch_related = batch_related_rule + else: + ctx.batch_related = batch_related_origin + ctx.args_list = args_list + ctx.input_requires_grad = input_requires_grad + ctx.output_list = outputs + if self.return_hidden_states: + ctx.middle_list = middles + + with torch.enable_grad(): + last_hidden = torch.cat(outputs, dim=0) + last_hidden_shape = last_hidden.shape + last_hidden = broadcast(last_hidden, config["pipe_size"] - 1, config["pipe_comm"]) + last_hidden = last_hidden.chunk(self.stages, dim=0) + last_hidden = last_hidden[self.stage_id].clone() + + if self.return_hidden_states: + middle_hiddens = [] + with torch.enable_grad(): + for stage_id in range(self.stages): + if self.stage_id == stage_id: + middle_hidden = torch.cat(middles, dim=1) # [(layers, micro_batch, ...), ] -> (layers, full_batch, ...) + else: + middle_shape = (self.get_part_len_by_stage_id(stage_id),)+last_hidden_shape + middle_hidden = torch.zeros(middle_shape, device=last_hidden.device, dtype=last_hidden.dtype) + middle_hidden = broadcast(middle_hidden, stage_id, config["pipe_comm"]) + middle_hidden = middle_hidden.chunk(self.stages, dim=1) + middle_hidden = middle_hidden[self.stage_id].clone() + middle_hiddens.append(middle_hidden) + middle_hiddens = torch.cat(middle_hiddens, dim=0) + return last_hidden, middle_hiddens + else: + return last_hidden, None + + + @staticmethod + def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle : torch.Tensor): + ipt = ctx.hidden_state_list + args_list = ctx.args_list + input_requires_grad = ctx.input_requires_grad + grad_hidden_state_list = all_gather(grad_hidden_state, config["pipe_comm"]).flatten(start_dim=0, end_dim=1).chunk(ctx.num_micros, dim=0) + if ctx.self.return_hidden_states: + for stage_id in range(ctx.self.stages): + layer_range = ctx.self.get_range_by_stage_id(stage_id) + grad_middle_state = grad_middle[layer_range] + grad_middle_state = all_gather(grad_middle_state.transpose(0,1), config["pipe_comm"]).flatten(start_dim=0, end_dim=1).transpose(0, 1).chunk(ctx.num_micros, dim=1) # (layer, micro_batch, ...) + if ctx.self.stage_id == stage_id: + grad_middle_state_list = grad_middle_state + + for m in range(ctx.num_micros): + output = ctx.output_list[m] + middle = ctx.middle_list[m] + grad_hidden_state = grad_hidden_state_list[m] + grad_middle_state = grad_middle_state_list[m] + torch.autograd.backward( + [output, middle], + [grad_hidden_state, grad_middle_state], + ) + else: + for m in range(ctx.num_micros): + output = ctx.output_list[m] + grad_hidden_state = grad_hidden_state_list[m] + torch.autograd.backward( + [output], + [grad_hidden_state], + ) + grads = [] + for idx,requires_grad in enumerate(input_requires_grad): + if requires_grad: + grad = torch.cat([args_list[m][idx].grad for m in range(ctx.num_micros)], dim=0) + grad = all_reduce(grad, "sum", config["pipe_comm"]) + split_size = ctx.self.stages if ctx.batch_related[idx] else ctx.num_micros + grad = grad.chunk(split_size) + if ctx.batch_related[idx]: + grads.append(grad[ctx.self.stage_id]) + else: + grads.append(grad[0]) + else: + grads.append(None) + grad = broadcast(ipt.grad, 0, config["pipe_comm"]).chunk(ctx.self.stages) + grad = grad[ctx.self.stage_id] + + return (None, None, None, grad) + tuple(grads) + (None,) + +class PipelineTransformerBlockList(torch.nn.Module): + r""" + TransformerBlockList is a list of CheckpointBlocks. + + This is designed to reduce the communication overhead by overlapping the computation and reduce_scatter operation during backward pass. + + It is similar to `torch.nn.ModuleList` but with the difference when calling .forward() and .backward(). + + Example: + >>> module_list = [ ... ] + >>> normal_module_list = torch.nn.ModuleList(module_list) + >>> transformer_module_list = TransformerBlockList(module_list) + >>> # Calling normal module list + >>> for layer in normal_module_list: + >>> hidden_state = layer.forward(hidden_state, ...) + >>> # Calling transformer module list + >>> hidden_state = transformer_module_list(hidden_state, ...) + + """ + _modules: Dict[str, CheckpointBlock] + + def __init__(self, modules: Iterable[CheckpointBlock]) -> None: + super().__init__() + + self._modules = {} + rank = config['rank'] + topo = config['topology'] + self.layer_ids = [] + pipe_group = topo.pp_group + self.stages = topo.stages + self.stage_id = topo.stage_id + self.pipe_idx = topo.pipe_idx + for idx, module in enumerate(modules): + if not isinstance(module, CheckpointBlock): + module = CheckpointBlock(module) + self._modules[str(idx)] = module + + self.layer_ids = self.get_range_by_stage_id(self.stage_id) + self.partition_modules(self.layer_ids) + self.next_rank = pipe_group[self.pipe_idx, self.stage_id + 1] if self.stage_id < config['pipe_size'] - 1 else -1 + self.prev_rank = pipe_group[self.pipe_idx, self.stage_id - 1] if self.stage_id > 0 else -1 + # self.micro_batches = config['num_micro_batches'] + + self.save_list = [(i, i) for i in range(len(self.layer_ids))] + + def __len__(self) -> int: + return len(self._modules) + + def __iter__(self) -> Iterator[CheckpointBlock]: + return iter(self._modules.values()) + + def __getitem__(self, index: Union[int, str]) -> CheckpointBlock: + return self._modules[str(index)] + + def forward(self, hidden_state, *args, batch_related=[], return_hidden_states=False): + self.return_hidden_states = return_hidden_states + placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled()) + args = list(args) + args.append(batch_related) + hidden_state, middle_states = OpPipeTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args) + if return_hidden_states: + return hidden_state, middle_states + else: + return hidden_state + + def get_range_by_stage_id(self, stage_id : int) -> List[int]: + part_lens = [0]+[self.get_part_len_by_stage_id(i) for i in range(stage_id+1)] + start = sum(part_lens[:stage_id+1]) + end = start + part_lens[stage_id+1] + return range(start, end) + + def get_part_len_by_stage_id(self, stage_id : int) -> int: + return len(self) // self.stages + (stage_id < (len(self) % self.stages)) + + def get_stage_by_layer_id(self, layer_id : int) -> int: + part_len = len(self) // self.stages + rest = len(self) % self.stages + if layer_id // (part_len + 1) < rest: + return layer_id // (part_len + 1) + else: + return rest + (layer_id - rest * (part_len+1)) // part_len + + def partition_modules(self, idxs) -> None: + for i in range(len(self)): + contiguous_params = {} + for kw, val in self[i]._storage_info.items(): + storage_type = val["storage_type"] + contiguous_params[kw] = storage_type(round_up(val["total"], config["world_size"] // config["pipe_size"])) + nccl.allGather( + self[i]._storage_params[kw].storage(), + contiguous_params[kw], + config["comm"] + ) + + if i not in idxs: + for name, param in self[i]._module.named_parameters(): + param.data = torch.tensor([], dtype = param.dtype, device = param.device) + for kw, val in self[i]._storage_info.items(): + val["begin"] = self.stage_id + val["end"] = self.stage_id + 1 + val["partition_size"] = 1 + val["total"] = val["world_size"] + dtype = self[i]._storage_params[kw].dtype + device = self[i]._storage_params[kw].device + self[i]._storage_params[kw] = \ + torch.nn.Parameter(torch.tensor([0], dtype = dtype, device=device)) + else: + for kw, val in self[i]._storage_info.items(): + storage_type = val["storage_type"] + val["world_size"] = config["world_size"] // config["pipe_size"] + partition_size = round_up(val["total"], val["world_size"]) // val["world_size"] + val["partition_size"] = partition_size + val["begin"] = config['zero_rank'] * partition_size + val["end"] = (config['zero_rank'] + 1) * partition_size + storage_param_buffer = storage_type(partition_size) + dtype = storage_param_buffer.dtype + device = storage_param_buffer.device + self[i]._storage_params[kw] = torch.nn.Parameter( + torch.tensor([], dtype=dtype, device=device).set_(storage_param_buffer) + ) + if val["requires_grad"]: + self[i]._storage_params[kw].requires_grad_(True) + else: + self[i]._storage_params[kw].requires_grad_(False) + ordered_parameters = list(self[i]._module.named_parameters()) + for idx, named_param in enumerate(ordered_parameters): + name, param = named_param + param_info = self[i]._param_info[idx] + kw_name = _get_param_kw(param) + storage_info = self[i]._storage_info[kw_name] + storage_st = storage_info["begin"] + storage_end = storage_info["end"] + param_st = param_info["offset"] + param_end = param_st + param_info["size"] + if not (param_st >= storage_end or param_end <= storage_st): + # copy offset in parameter storage + offset_st = max(storage_st - param_st, 0) + offset_end = min(storage_end - param_st, param_info["size"]) + assert offset_st < offset_end + to_offset_st = offset_st + param_st - storage_st + to_offset_end = offset_end + param_st - storage_st + d_dtype = self[i]._storage_params[kw_name].dtype + d_device = self[i]._storage_params[kw_name].device + param.data = torch.tensor([], dtype=param.dtype, device=param.device).set_(self[i]._storage_params[kw_name].storage(), to_offset_st, (to_offset_end - to_offset_st,)) + param_info["begin"] = to_offset_st + param_info["end"] = (to_offset_end - to_offset_st,) + param.data[:] = \ + torch.tensor([], dtype=d_dtype, device=d_device).set_(contiguous_params[kw], storage_st+to_offset_st, (to_offset_end - to_offset_st,))[:] + else: + param.data = torch.tensor([], dtype=param.dtype, device=param.device) + del contiguous_params + + def _save_to_state_dict(self, destination, prefix, keep_vars): + for name, module in self._modules.items(): + idx = int(name) + name = prefix + name + '.' + + dst = OrderedDict() # creates an temporary ordered dict + dst._metadata = OrderedDict() + + if idx in self.layer_ids: + with torch.no_grad(): + with CheckpointBlockContext(module, pipe=True): + module._module.state_dict(destination=dst, prefix=name, keep_vars=False) + if config["zero_rank"] == 0: + if config["rank"] == 0: + destination.update(dst) + else: + assert list(dst.keys()) == [name+n for n, parameter in module._module.named_parameters()] + for key, tensor in dst.items(): + send_activations(tensor.cuda(), 0, config['pipe_comm']) + if config['rank'] == 0 and idx not in self.layer_ids: + for n, parameter in module._module.named_parameters(): + destination[name+n] = recv_activations(self.get_stage_by_layer_id(idx), config['pipe_comm']) + +class PipeContext: + def __init__(self, module, hidden_state, backward=False): + self.module = module + self.stage_id = module.stage_id + self.stages = module.stages + self.next_rank = module.next_rank + self.prev_rank = module.prev_rank + self.hidden_state = [hidden_state] + self.backward = backward + self.send_buffer = {} + def enter(self): + if self.backward: + if self.stage_id != self.stages -1: + self.hidden_state[0] = recv_activations(self.stage_id + 1, config['pipe_comm']) + else: + if self.stage_id != 0: + self.hidden_state[0] = recv_activations(self.stage_id - 1, config['pipe_comm']) + return self.hidden_state + def exit(self): + if self.backward: + if self.stage_id != 0: + send_activations(self.hidden_state[0], self.stage_id - 1, config['pipe_comm']) + else: + if self.stage_id != self.stages - 1: + send_activations(self.hidden_state[0], self.stage_id + 1, config['pipe_comm']) + def __enter__(self): + return self.enter() + def __exit__(self, exc_type, exc_val, exc_tb): + self.exit() \ No newline at end of file diff --git a/bmtrain/store.py b/bmtrain/store.py index bf1c00f7..c0914f9a 100644 --- a/bmtrain/store.py +++ b/bmtrain/store.py @@ -1,6 +1,8 @@ from collections import OrderedDict from typing import Dict import torch + +from .pipe_layer import PipelineTransformerBlockList from .global_var import config from .block_layer import CheckpointBlock from . import nccl @@ -24,15 +26,19 @@ def _save_to_rank0(model : torch.nn.Module, destination=None, prefix=''): destination = OrderedDict() destination._metadata = OrderedDict() destination._metadata[prefix[:-1]] = local_metadata = dict(version=model._version) - _save_to_state_dict(model, destination, prefix) - for name, module in model._modules.items(): - if module is not None: - _save_to_rank0(module, destination, prefix + name + '.') - for hook in model._state_dict_hooks.values(): - hook_result = hook(model, destination, prefix, local_metadata) - if hook_result is not None: - destination = hook_result + if not isinstance(model, PipelineTransformerBlockList): + _save_to_state_dict(model, destination, prefix) + for name, module in model._modules.items(): + if module is not None: + _save_to_rank0(module, destination, prefix + name + '.') + for hook in model._state_dict_hooks.values(): + hook_result = hook(model, destination, prefix, local_metadata) + if hook_result is not None: + destination = hook_result + else: + model._save_to_state_dict(destination, prefix, False) return destination + def save(model : torch.nn.Module, file_name : str): @@ -66,9 +72,28 @@ def save(model : torch.nn.Module, file_name : str): _pickler = pickle.Pickler _unpickler = pickle.Unpickler - -def broadcast_object(obj): - if config['rank'] == 0: +def allgather_object(obj, comm): + f = io.BytesIO() + _pickler(f).dump(obj) + byte_storage = torch.ByteStorage.from_buffer(f.getvalue()) + # Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor and specifying dtype. + # Otherwise, it will casue 100X slowdown. + # See: + byte_tensor = torch.ByteTensor(byte_storage).cuda() + all_bytes_tensors = torch.empty(byte_tensor.numel() * nccl.commCount(comm), dtype=torch.uint8, device="cuda") + nccl.allGather( + byte_tensor.storage(), + all_bytes_tensors.storage(), + comm + ) + obj_list = [] + for i in range(nccl.commCount(comm)): + buf = all_bytes_tensors[i*byte_tensor.numel():(i+1)*byte_tensor.numel()].cpu().numpy().tobytes() + obj = _unpickler(io.BytesIO(buf)).load() + obj_list.append(obj) + return obj_list +def broadcast_object(obj, comm, src = 0): + if nccl.commRank(comm) == src: f = io.BytesIO() _pickler(f).dump(obj) byte_storage = torch.ByteStorage.from_buffer(f.getvalue()) @@ -81,30 +106,30 @@ def broadcast_object(obj): nccl.broadcast( local_size.storage(), local_size.storage(), - 0, - config['comm'] + src, + comm ) nccl.broadcast( byte_tensor.storage(), byte_tensor.storage(), - 0, - config['comm'] + src, + comm ) else: local_size = torch.LongTensor([0]).cuda() nccl.broadcast( local_size.storage(), local_size.storage(), - 0, - config['comm'] + src, + comm ) byte_tensor_size = local_size[0].item() byte_tensor = torch.empty(int(byte_tensor_size), dtype=torch.uint8, device="cuda") nccl.broadcast( byte_tensor.storage(), byte_tensor.storage(), - 0, - config['comm'] + src, + comm ) buf = byte_tensor.cpu().numpy().tobytes() obj = _unpickler(io.BytesIO(buf)).load() @@ -114,7 +139,7 @@ def broadcast_object(obj): class DistributedStateDictWrapper(Mapping): def __init__(self, state_dict : Dict) -> None: self._state_dict = state_dict - self._metadata = broadcast_object(getattr(state_dict, "_metadata", None)) + self._metadata = broadcast_object(getattr(state_dict, "_metadata", None), config["comm"]) def __getitem__(self, key : str): tmp_shape = torch.zeros(32, device="cuda", dtype=torch.int32) @@ -169,13 +194,13 @@ def copy(self): return self def __len__(self): - return broadcast_object(len(self._state_dict)) + return broadcast_object(len(self._state_dict), config["comm"]) def __contains__(self, key : str): - return broadcast_object(key in self._state_dict) + return broadcast_object(key in self._state_dict, config["comm"]) def keys(self): - return broadcast_object(list(self._state_dict.keys())) + return broadcast_object(list(self._state_dict.keys()),config["comm"]) def __iter__(self): # pytorch 1.12.0 updated the load_state_dict method, which needs the state_dict to be a `Mapping`. diff --git a/csrc/include/nccl.h b/csrc/include/nccl.h index fc6247b1..4bdda76a 100644 --- a/csrc/include/nccl.h +++ b/csrc/include/nccl.h @@ -351,6 +351,7 @@ ncclResult_t pncclGroupStart(); ncclResult_t ncclGroupEnd(); ncclResult_t pncclGroupEnd(); + #ifdef __cplusplus } // end extern "C" #endif diff --git a/csrc/nccl.cpp b/csrc/nccl.cpp index 5a4a275a..2d0ac8de 100644 --- a/csrc/nccl.cpp +++ b/csrc/nccl.cpp @@ -129,7 +129,40 @@ void pyNCCLReduceScatter( reinterpret_cast(stream) )); } - +void pyNCCLSend( + std::uintptr_t sendbuff, + size_t sendcount, + int data_type, + int peer, + std::uintptr_t comm, + std::uintptr_t stream +) { + checkNCCLStatus(ncclSend( + reinterpret_cast(sendbuff), + sendcount, + static_cast(data_type), + peer, + reinterpret_cast(comm), + reinterpret_cast(stream) + )); +} +void pyNCCLRecv( + std::uintptr_t recvbuff, + size_t recvcount, + int data_type, + int peer, + std::uintptr_t comm, + std::uintptr_t stream +) { + checkNCCLStatus(ncclRecv( + reinterpret_cast(recvbuff), + recvcount, + static_cast(data_type), + peer, + reinterpret_cast(comm), + reinterpret_cast(stream) + )); +} void pyNCCLGroupStart() { checkNCCLStatus(ncclGroupStart()); } @@ -137,7 +170,20 @@ void pyNCCLGroupStart() { void pyNCCLGroupEnd() { checkNCCLStatus(ncclGroupEnd()); } - +int pyNCCLCommCount( + std::uintptr_t comm +){ + int res; + checkNCCLStatus(ncclCommCount(reinterpret_cast(comm),&res)); + return res; +} +int pyNCCLCommUserRank( + std::uintptr_t comm +){ + int rank; + checkNCCLStatus(ncclCommUserRank(reinterpret_cast(comm),&rank)); + return rank; +} PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("ncclGetUniqueId", &pyNCCLGetUniqueID, "nccl get unique ID"); m.def("ncclCommInitRank", &pyNCCLCommInitRank, "nccl init rank"); @@ -149,4 +195,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("ncclReduceScatter", &pyNCCLReduceScatter, "nccl reduce scatter"); m.def("ncclGroupStart", &pyNCCLGroupStart, "nccl group start"); m.def("ncclGroupEnd", &pyNCCLGroupEnd, "nccl group end"); + m.def("ncclSend",&pyNCCLSend,"nccl send"); + m.def("ncclRecv",&pyNCCLRecv,"nccl recv"); + m.def("ncclCommCount",&pyNCCLCommCount,"nccl comm count"); + m.def("ncclCommUserRank",&pyNCCLCommUserRank,"nccl comm user rank"); } diff --git a/example/models/gpt.py b/example/models/gpt.py index 65779e5c..78d77a7d 100644 --- a/example/models/gpt.py +++ b/example/models/gpt.py @@ -36,9 +36,10 @@ def forward(self, mask_2d = mask[:, None, :] & mask[:, :, None] # (batch, seq_len, seq_len) mask_2d = mask_2d & (pos[:, None, :] >= pos[:, :, None]) - input_emb = self.pos_emb(pos) + self.word_emb(input) + out = self.pos_emb(pos) + self.word_emb(input) - out = self.transformers(input_emb, mask_2d, None) + # for layer in self.transformers: + out = self.transformers(out, mask_2d, None) out = self.layernorm(out) logits = self.word_emb(out, projection=True) diff --git a/example/run.sh b/example/run.sh index c0f5a47e..542e5252 100644 --- a/example/run.sh +++ b/example/run.sh @@ -1 +1,3 @@ -torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost train.py \ No newline at end of file +export NCCL_P2P_DISABLE=1 +export CUDA_LAUNCH_BLOCKING=1 +torchrun --nnodes=1 --nproc_per_node=4 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost train.py diff --git a/tests/test_middle_hidden.py b/tests/test_middle_hidden.py new file mode 100644 index 00000000..ca0dc9b3 --- /dev/null +++ b/tests/test_middle_hidden.py @@ -0,0 +1,188 @@ +import bmtrain as bmt +import random +import torch +from bmtrain import config +from bmtrain.block_layer import CheckpointBlock, TransformerBlockList +from bmtrain.pipe_layer import PipelineTransformerBlockList +import torch.nn.functional as F + +class Linear(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, init_weight = None, init_bias = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.out = {} + if init_weight: + self.weight = bmt.DistributedParameter(torch.tensor(init_weight, dtype=torch.float, device="cuda").reshape(out_features, in_features)) + else: + self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.xavier_normal_) + + if init_bias: + self.bias = bmt.DistributedParameter(torch.tensor(init_bias, dtype=torch.float, device="cuda").reshape(out_features,)) + else: + self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.zeros_) + + def forward(self, input): + ret = F.linear(input, self.weight, self.bias) + return ret + +class Model_ZERO(torch.nn.Module): + def __init__(self, pre, ms, post) -> None: + super().__init__() + self.pre = pre + self.ms = TransformerBlockList([ + CheckpointBlock(m) + for m in ms + ]) + self.post = post + + def forward(self, x, return_hidden_states=False): + x = self.pre(x) + if return_hidden_states: + x, o = self.ms(x, return_hidden_states=return_hidden_states) + return self.post(x), o + else: + x = self.ms(x, return_hidden_states=return_hidden_states) + return self.post(x) + +class Model_PIPE(torch.nn.Module): + def __init__(self, pre, ms, post) -> None: + super().__init__() + self.pre = pre + self.ms = PipelineTransformerBlockList([ + CheckpointBlock(m) + for m in ms + ]) + self.post = post + + def forward(self, x, return_hidden_states=False): + x = self.pre(x) + if return_hidden_states: + x, o = self.ms(x, return_hidden_states=return_hidden_states) + return self.post(x), o + else: + x = self.ms(x, return_hidden_states=return_hidden_states) + return self.post(x) + +class Model_BLOCK(torch.nn.Module): + def __init__(self, pre, ms, post) -> None: + super().__init__() + self.pre = pre + self.ms = torch.nn.ModuleList([ + CheckpointBlock(m) + for m in ms + ]) + self.post = post + + def forward(self, x, return_hidden_states=False): + x = self.pre(x) + o = [] + y = x + for m in self.ms: + o.append(y) + y = m(y) + if return_hidden_states: + return self.post(y), o + else: + return self.post(y) + +class Model_NORMAL(torch.nn.Module): + def __init__(self, pre, ms, post) -> None: + super().__init__() + self.pre = pre + self.ms = torch.nn.ModuleList(ms) + self.post = post + + def forward(self, x, return_hidden_states=False): + x = self.pre(x) + o = [] + y = x + for m in self.ms: + o.append(y) + y = m(y) + if return_hidden_states: + return self.post(y), o + else: + return self.post(y) + +def manual_seed(seed=33): + torch.manual_seed(seed) + random.seed(seed) + try: + import numpy as np + np.random.seed(seed) + except ModuleNotFoundError: + pass + +def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False): + manual_seed() + + pre = Linear(dim, dim) + post = Linear(dim, dim) + ms = [Linear(dim, dim) for i in range(num_layer)] + + inp = torch.randn((batch, seq_len, dim)).cuda() + last_weight = torch.randn((batch, seq_len, dim)).cuda() + middle_weight = [ + torch.randn((batch, seq_len, dim)).cuda() + for i in range(len(ms)) + ] + + bmt.init_parameters(pre) + bmt.init_parameters(post) + for m in ms: + bmt.init_parameters(m) + m = cls(pre, [m for m in ms], post) + + if only_last: + logits = m(inp) + loss = (logits * last_weight).sum() + loss.backward() + bmt.print_rank(f"========================{name}:only last========================") + bmt.print_rank( + bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') + ) + ) + if only_middle: + logits, hidden_states = m(inp, return_hidden_states=True) + loss = sum([ + (hidden_state * middle_weight[i]).sum() + for i, hidden_state in enumerate(hidden_states) + ]) + loss.backward() + bmt.print_rank(f"========================{name}:only middle========================") + bmt.print_rank( + bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') + ) + ) + if mix_test: + logits, hidden_states = m(inp, return_hidden_states=True) + loss = sum([ + (hidden_state * middle_weight[i]).sum() + for i, hidden_state in enumerate(hidden_states) + ]) + (logits * last_weight).sum() + loss.backward() + bmt.print_rank(f"========================{name}:mix========================") + bmt.print_rank( + bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') + ) + ) + +def test(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256): + sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True) + bmt.synchronize() + sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True) + bmt.synchronize() + sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True) + bmt.synchronize() + +bmt.init_distributed(pipe_size=4) + +test("normal", Model_NORMAL) +test("block", Model_BLOCK) +test("zero", Model_ZERO) +test("pipe", Model_PIPE) \ No newline at end of file diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py new file mode 100644 index 00000000..f969d661 --- /dev/null +++ b/tests/test_send_recv.py @@ -0,0 +1,27 @@ +import bmtrain as bmt +import torch +from bmtrain.global_var import config +from bmtrain.pipe_comm import send_activations, recv_activations, gather_input +from bmtrain import nccl +from torch.distributed.distributed_c10d import recv +from time import sleep +def test_send_tensor(): + bmt.init_distributed() + current_stream = torch.cuda.current_stream() + groups = [0,2] + rank = config['rank'] +def test_gather_input(): + bmt.init_distributed(pipe_size=2) + if config['topology'].get_group_id("pipe") == 0: + a = torch.ones((2,1)) + else: + a = torch.zeros((2,1)) + res = gather_input(a.cuda(),config['pipe_comm']) + if config['topology'].get_group_rank("pipe") == 0: + print(res) + +def main(): + test_send_tensor() + +if __name__ == '__main__': + main() \ No newline at end of file