From f32f948f1ba34c7fbf6c0b1f86f999a3ba7eb4d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Thu, 20 Jun 2024 17:06:40 +0800 Subject: [PATCH 1/6] first zeropp commit --- configs/7B_llama2_zeropp.py | 211 +++++ internlm/core/context/parallel_context.py | 6 +- internlm/core/scheduler/__init__.py | 3 +- internlm/core/scheduler/pipeline_scheduler.py | 775 +++++++++++++++++- internlm/core/trainer.py | 3 +- internlm/initialize/initialize_trainer.py | 37 +- internlm/model/modules/linear.py | 90 +- .../solver/optimizer/hybrid_zero_optim.py | 13 +- internlm/utils/zeropp_manager.py | 47 ++ 9 files changed, 1143 insertions(+), 42 deletions(-) create mode 100644 configs/7B_llama2_zeropp.py create mode 100644 internlm/utils/zeropp_manager.py diff --git a/configs/7B_llama2_zeropp.py b/configs/7B_llama2_zeropp.py new file mode 100644 index 000000000..1652743de --- /dev/null +++ b/configs/7B_llama2_zeropp.py @@ -0,0 +1,211 @@ +JOB_NAME = "7b_llama2_train" +model_type = "LLAMA2" +DO_ALERT = False + +VOCAB_SIZE = 32000 +SEQ_LEN = 2048 +HIDDEN_SIZE = 4096 +NUM_ATTENTION_HEAD = 32 +NUM_KV_ATTENTION_HEAD = 8 +MLP_RATIO = 3.5 +NUM_LAYER = 32 + + +MODEL_ONLY_FOLDER = "local:llm_ckpts/xxxx" +# Ckpt folder format: +# fs: 'local:/mnt/nfs/XXX' +SAVE_CKPT_FOLDER = "local:llm_ckpts" +LOAD_CKPT_FOLDER = "local:llm_ckpts/49" + +# boto3 Ckpt folder format: +# import os +# BOTO3_IP = os.environ["BOTO3_IP"] # boto3 bucket endpoint +# SAVE_CKPT_FOLDER = f"boto3:s3://model_weights.{BOTO3_IP}/internlm" +# LOAD_CKPT_FOLDER = f"boto3:s3://model_weights.{BOTO3_IP}/internlm/snapshot/1/" +CHECKPOINT_EVERY = 50 +ckpt = dict( + enable_save_ckpt=False, # enable ckpt save. + save_ckpt_folder=SAVE_CKPT_FOLDER, # Path to save training ckpt. + # 'auto_resume' is designed to automatically load the latest checkpoint from 'save_ckpt_folder' when encountering + # training interruptions/hangs caused by hardware failures, using a scheduling system (such as k8s/slurm) + # with an automatic restart mechanism upon training reboot. + # Please be aware that if `auto_resume` is not set (its default value is True), it will not load the checkpoint + # path specified in `load_ckpt_info` by default. + # If you want to initialize your model weights from another model, you must set `auto_resume` to False. + # If you want to train from scratch, please set `auto_resume` to False and 'load_ckpt_info' to None. + auto_resume=False, + checkpoint_every=CHECKPOINT_EVERY, + async_upload=True, # async ckpt upload. (only work for boto3 ckpt) + async_upload_tmp_folder="/dev/shm/internlm_tmp_ckpt/", # path for temporarily files during asynchronous upload. + oss_snapshot_freq=int(CHECKPOINT_EVERY / 2), # snapshot ckpt save frequency. +) + +TRAIN_FOLDER = None +VALID_FOLDER = None # "/path/to/dataset" +data = dict( + seq_len=SEQ_LEN, + # micro_num means the number of micro_batch contained in one gradient update + micro_num=4, + # unit_schedule_size is used with ZeroPP and means the number of micro_batch contained in one scheduling unit. + # If unit_schedule_size is not set, ZeroPP will use Interleaved 1F1B schedule. Otherwise it will use interleaved GPipe. + unit_schedule_size=4, + # packed_length = micro_bsz * SEQ_LEN + micro_bsz=1, + # defaults to the value of micro_num + valid_micro_num=4, + # defaults to 0, means disable evaluate + valid_every=0, + pack_sample_into_one=False, + total_steps=20, + skip_batches="", + # rampup_batch_size (str): A string with three space-separated integers representing the + # starting batch size, the increment, and the number of steps between + # each increment. For example, "192 24 8" means that the batch size (micro_num) + # starts at 192 and increases by 24 every 8 steps. Defaults to None. + # (IMPORTANT): The interval step size is 'micro_bsz'. + rampup_batch_size="", + # Datasets with less than 50 rows will be discarded + min_length=50, + train_folder=TRAIN_FOLDER, + valid_folder=VALID_FOLDER, + empty_cache_and_diag_interval=200, + diag_outlier_ratio=1.1, +) + +grad_scaler = dict( + fp16=dict( + # the initial loss scale, defaults to 2**16 + initial_scale=2**16, + # the minimum loss scale, defaults to None + min_scale=1, + # the number of steps to increase loss scale when no overflow occurs + growth_interval=1000, + ), + # the multiplication factor for increasing loss scale, defaults to 2 + growth_factor=2, + # the multiplication factor for decreasing loss scale, defaults to 0.5 + backoff_factor=0.5, + # the maximum loss scale, defaults to None + max_scale=2**24, + # the number of overflows before decreasing loss scale, defaults to 2 + hysteresis=2, +) + +hybrid_zero_optimizer = dict( + # Enable low_level_optimzer overlap_communication + overlap_sync_grad=False, + overlap_sync_param=False, + # bucket size for nccl communication params + reduce_bucket_size=512 * 1024 * 1024, + # grad clipping + clip_grad_norm=1.0, +) + +loss = dict( + label_smoothing=0, +) + +adam = dict( + lr=1e-4, + adam_beta1=0.9, + adam_beta2=0.95, + adam_beta2_c=0, + adam_eps=1e-8, + weight_decay=0.01, +) + +lr_scheduler = dict( + total_steps=data["total_steps"], + init_steps=0, # optimizer_warmup_step + warmup_ratio=0.01, + eta_min=1e-5, + last_epoch=-1, +) + +beta2_scheduler = dict( + init_beta2=adam["adam_beta2"], + c=adam["adam_beta2_c"], + cur_iter=-1, +) + +use_fp32_norm = False +model = dict( + checkpoint=False, + num_chunks=4, + num_attention_heads=NUM_ATTENTION_HEAD, + embed_split_hidden=True, + vocab_size=VOCAB_SIZE, + embed_grad_scale=1, + parallel_output=True, + hidden_size=HIDDEN_SIZE, + num_layers=NUM_LAYER, + no_bias=True, + mlp_ratio=MLP_RATIO, + apply_post_layer_norm=False, + dtype="torch.bfloat16", + norm_type="rmsnorm", + layer_norm_epsilon=1e-5, + num_kv_attention_heads=NUM_KV_ATTENTION_HEAD, + use_flash_attn=True, + # Whether the odd and even columns of the query and key in the model are normally interleaved. + # If it's True, the model's odd and even columns are normally ordered; if it's False, + # it means that the model has prematurely concatenated all odd columns and even columns in front + # and back, in order to improve the RoPE's computational efficiency. + # Example: + # qk_interleaved = True: q[-1] = [q1,q2,q3,q4,q5,q6,...], k[-1] = [k1,k2,k3,k4,k5,k6,...] + # qk_interleaved = False: q[-1] = [q1,q3,q5,...,q2,q4,q6,...], k[-1] = [k1,k3,k5,...,k2,k4,k6,...] + qk_interleaved=False, +) + +""" +zero1 parallel (dict): + 1. size: int + * if size <= 0, the size of the zero process group is equal to the size of the dp process group, + so parameters will be divided within the range of dp. + * if size == 1, zero is not used, and all dp groups retain the full amount of model parameters. + * if size > 1 and size <= dp world size, the world size of zero is a subset of dp world size. + For smaller models, it is usually a better choice to split the parameters within nodes with a setting <= 8. + 2. fsdp: bool, enable/disable torch's fully sharded data parallel, defaults to False. +tensor parallel (dict): + 1. size: int, the size of tensor parallel. + 2. mode: str, the tensor parallel mode, should be in ['mtp', 'msp', 'fsp', 'isp'], + defaults to 'mtp', means the pure megatron tensor parallel without sequence parallel. + msp: megatron tensor parallel with sequence parallel, sequence parallel size = tensor parallel size. + fsp: tensor parallel by flash-attn with sequence parallel, sequence parallel size = tensor parallel size. + isp: customed intern sequence parallel without tensor parallel, can be used with weight parallel. +pipeline parallel (dict): + 1. size: int, the size of pipeline parallel. + 2. interleaved_overlap: bool, enable/disable communication overlap when using interleaved pipeline scheduler, + defaults to False. +weight parallel (dict): + 1. size: int, the size of weight parallel. + 2. overlap: bool, enable/disable all_gather/reduce_scatter communication overlap, defaults to False. + 3. memory_pool: bool, enable/disable memory pool, defaults to False. +""" +parallel = dict( + zero1=dict(size=-1), + tensor=dict(size=1, mode="isp"), + #pipeline=dict(size=4, interleaved_overlap=True, use_zeropp=True, decouple_grad=True), + pipeline=dict(size=4, interleaved_overlap=True, use_zeropp=True, decouple_grad=True), + weight=dict(size=8, overlap=True, memory_pool=False), +) + +cudnn_deterministic = False +cudnn_benchmark = False + +monitor = dict( + # feishu alert configs + alert=dict( + enable_feishu_alert=DO_ALERT, + feishu_alert_address=None, # feishu webhook to send alert message + light_monitor_address=None, # light_monitor address to send heartbeat + alert_file_path=f"llm_alter/{JOB_NAME}_alert.log", + ), + tensorboard=dict( + queue_max_length=10, + ), +) + +# metric_dtype can be "fp32" or other string +# only when set to "fp32" will use fp32 to calc in metrics +# metric_dtype = "fp32" diff --git a/internlm/core/context/parallel_context.py b/internlm/core/context/parallel_context.py index 6b23fdae6..0ce567b43 100644 --- a/internlm/core/context/parallel_context.py +++ b/internlm/core/context/parallel_context.py @@ -19,6 +19,7 @@ from internlm.utils.common import SingletonMeta from internlm.utils.logger import get_logger from internlm.utils.timeout import LLM_NCCL_TIMEOUT +from internlm.utils.zeropp_manager import ZeroppManager from . import process_group_initializer as pgroup_initializer from .process_group_initializer import ParallelMode @@ -666,7 +667,10 @@ def set_seed(self, seed: int, dpseed_with_tpoffset: bool = False): def set_virtual_pipeline_parallel_size(self, size): self.virtual_pipeline_parallel_size = size - def set_virtual_pipeline_parallel_rank(self, rank): + def set_virtual_pipeline_parallel_rank(self, rank, clear_zeropp_cache=True): + if rank != self.virtual_pipeline_parallel_rank: + if clear_zeropp_cache: + ZeroppManager.clear_cached_wp_parameters() self.virtual_pipeline_parallel_rank = rank diff --git a/internlm/core/scheduler/__init__.py b/internlm/core/scheduler/__init__.py index b65794737..a86dee753 100644 --- a/internlm/core/scheduler/__init__.py +++ b/internlm/core/scheduler/__init__.py @@ -1,10 +1,11 @@ from .base_scheduler import BaseScheduler from .no_pipeline_scheduler import NonPipelineScheduler -from .pipeline_scheduler import InterleavedPipelineScheduler, PipelineScheduler +from .pipeline_scheduler import InterleavedPipelineScheduler, PipelineScheduler, ZeroPPScheduler __all__ = [ "BaseScheduler", "NonPipelineScheduler", "InterleavedPipelineScheduler", "PipelineScheduler", + "ZeroPPScheduler", ] diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py index 269ddb966..3324629f8 100644 --- a/internlm/core/scheduler/pipeline_scheduler.py +++ b/internlm/core/scheduler/pipeline_scheduler.py @@ -22,6 +22,7 @@ ) from internlm.utils.logger import get_logger from internlm.utils.timeout import llm_timeout +from internlm.utils.zeropp_manager import ZeroppManager from .base_scheduler import BaseScheduler @@ -92,10 +93,10 @@ def pack_return_tensors(return_tensors): def switch_virtual_pipeline_parallel_rank(rank): prev_rank = gpc.virtual_pipeline_parallel_rank try: - gpc.set_virtual_pipeline_parallel_rank(rank) + gpc.set_virtual_pipeline_parallel_rank(rank, clear_zeropp_cache=False) yield finally: - gpc.set_virtual_pipeline_parallel_rank(prev_rank) + gpc.set_virtual_pipeline_parallel_rank(prev_rank, clear_zeropp_cache=False) @contextmanager @@ -1417,3 +1418,773 @@ def forward_backward_step(self, engine, data_iter, forward_only=False, return_lo return output, label, accum_loss, accum_moe_loss else: return output, label, accum_loss + + + + +class ZeroPPScheduler(PipelineScheduler): + """ + ZeroPP Pipeline Scheduler. + """ + + def __init__( + self, + num_microbatches: int, + num_chunks: int, + dtype: torch.dtype = torch.float, + data_process_func: Callable = None, + tensor_shape: Union[torch.Size, List[int], Tuple[int]] = None, + scatter_gather_tensors: bool = False, + scheduler_hooks: Optional[List[SchedulerHook]] = None, + communication_overlap: bool = False, + unit_schedule_size: Optional[int] = None, + ): + """A helper schedule class for pipeline parallelism running environment. + It uses ZeroPP strategy. Other properties are similar as + :class:`PipelineSchedule`. + + .. ZeroPP: https://arxiv.org/abs/2402.03791 + + Args: + num_microbatches (int): The number of microbatches. + num_chunks (int): The number of model chunks. + dtype (torch.dtype, optional): The data type of the tensors. Default is torch.float. + data_process_func (Callable, optional): + The preprocessing function which receives a batch of data, and it will be executed in `load_batch`. + tensor_shape (torch.Size, optional): Specified shape in pipeline communication. + scatter_gather_tensors (bool, optional): + If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization. + scheduler_hooks (List[SchedulerHook], optional): List of scheduler hooks. Default is None. + communication_overlap (bool, optional): Whether to enable communication overlap. Default is False. + unit_schedule_size: (int, optional): Unit size if using interleaved GPipe schedule. Default is None and interleaved 1F1B schedule is used. + """ + assert ( + isinstance(num_chunks, int) and num_chunks > 0 + ), f"expected num_chunks to be an integer and larger than 0, but got {num_chunks}" + + super().__init__( + num_microbatches, + dtype=dtype, + data_process_func=data_process_func, + tensor_shape=tensor_shape, + scatter_gather_tensors=scatter_gather_tensors, + scheduler_hooks=scheduler_hooks, + ) + + gpc.set_virtual_pipeline_parallel_size(num_chunks) + gpc.set_virtual_pipeline_parallel_rank(0) + + self._num_chunks = num_chunks + self._communication_overlap = communication_overlap + # switch 1f1b loop runner function according to communication overlap + self._run_1f1b_loop = self._run_1f1b_loop_with_overlap + + # states + self._pp_size = gpc.get_world_size(ParallelMode.PIPELINE) + self._pp_rank = gpc.get_local_rank(ParallelMode.PIPELINE) + + self._accum_loss = None + self._accum_moe_loss = None + self._return_tensors = None + self._input_objs = [[] for _ in range(num_chunks)] + self._input_objs_for_backward = [[] for _ in range(num_chunks)] + self._output_objs = [[] for _ in range(num_chunks)] + self._output_obj_grads = [[] for _ in range(num_chunks)] + self._moe_losses = [[] for _ in range(num_chunks)] + + self._input_obj_shapes = [self.tensor_shape for _ in range(num_chunks)] + self._output_obj_shapes = [None for _ in range(num_chunks)] + self._send_tensor_shape_flags = [self.tensor_shape is None for _ in range(num_chunks)] + + if unit_schedule_size is not None and unit_schedule_size > 0: + self._unit_schedule_size = unit_schedule_size + self._forward_backward_step = self._forward_backward_step_gpipe + self._get_current_microbatch_id = self._get_chunk_by_microbatch_gpipe + self._get_chunk_by_microbatch = self._get_chunk_by_microbatch_gpipe + else: + self._unit_schedule_size = None + self._forward_backward_step = self._forward_backward_step_1f1b + self._get_current_microbatch_id = self._get_chunk_by_microbatch_1f1b + self._get_chunk_by_microbatch = self._get_chunk_by_microbatch_1f1b + + self._cur_unit_schedule_size = self._unit_schedule_size + if hasattr(gpc.config.parallel.pipeline, "decouple_grad") and gpc.config.parallel.pipeline.decouple_grad is True: + self.decouple_grad = True + else: + self.decouple_grad = False + + @property + def tensor_shape(self) -> torch.Size: + return self._tensor_shape + + @tensor_shape.setter + def tensor_shape(self, tensor_shape: torch.Size): + self._tensor_shape = tensor_shape + self._input_obj_shapes = [self._tensor_shape for _ in range(self._num_chunks)] + self._send_tensor_shape_flags = [self._tensor_shape is None for _ in range(self._num_chunks)] + + def _clear_state(self) -> None: + self._accum_loss = None + self._accum_moe_loss = None + self._return_tensors = None + self._input_objs = [[] for _ in range(self._num_chunks)] + self._input_objs_for_backward = [[] for _ in range(self._num_chunks)] + self._output_objs = [[] for _ in range(self._num_chunks)] + self._output_obj_grads = [[] for _ in range(self._num_chunks)] + self._moe_losses = [[] for _ in range(self._num_chunks)] + + self._input_obj_shapes = [self.tensor_shape for _ in range(self._num_chunks)] + self._output_obj_shapes = [None for _ in range(self._num_chunks)] + self._send_tensor_shape_flags = [self.tensor_shape is None for _ in range(self._num_chunks)] + + def _clear_state_unit(self) -> None: + self._input_objs = [[] for _ in range(self._num_chunks)] + self._input_objs_for_backward = [[] for _ in range(self._num_chunks)] + self._output_objs = [[] for _ in range(self._num_chunks)] + self._output_obj_grads = [[] for _ in range(self._num_chunks)] + self._moe_losses = [[] for _ in range(self._num_chunks)] + + self._input_obj_shapes = [self.tensor_shape for _ in range(self._num_chunks)] + self._output_obj_shapes = [None for _ in range(self._num_chunks)] + self._send_tensor_shape_flags = [self.tensor_shape is None for _ in range(self._num_chunks)] + + def load_batch(self, engine, data_iter): + super().load_batch(engine, data_iter) + # overwrite microbatch_offset, since model chunks load the same microbatch, and should tract the offset + self.microbatch_offset = [0 for _ in range(self._num_chunks)] + + def load_micro_batch(self, model_chunk_id): + micro_batch_data, micro_batch_label = self._load_micro_batch( + data=self.batch_data, + label=self.batch_label, + offset=self.microbatch_offset[model_chunk_id], + bsz_stride=self.bsz_stride, + ) + if self.data_process_func: + micro_batch_data, micro_batch_label = self.data_process_func(micro_batch_data, micro_batch_label) + + micro_batch_data["label"] = micro_batch_label + self.microbatch_offset[model_chunk_id] += self.bsz_stride + return move_to_device(micro_batch_data) + + def _forward_step(self, engine, chunk_id): + """Forward step for passed-in model. If it is the first stage, the input tensor + is obtained from data_iterator, otherwise the passed-in input_obj is used. + Returns output tensor. This is a helper function and can be ignored by users. + + Args: + engine (colossalai.engine.Engine): Colossalai engine for training and inference. + chunk_id (int): The id of model chunks. + Returns: + Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: output or the loss value of the current + pipeline stage. + """ + + + gpc.set_virtual_pipeline_parallel_rank(chunk_id) + + if gpc.is_pipeline_first_stage() and (len(self._input_objs[chunk_id]) + len(self._input_objs_for_backward[chunk_id])) == len(self._output_objs[chunk_id]): + self._input_objs[chunk_id].append(None) + + input_obj = self._input_objs[chunk_id].pop(0) + self._input_objs_for_backward[chunk_id].append(input_obj) + + micro_batch_data = self.load_micro_batch(chunk_id) + data, label = self._get_data_label_for_current_step(input_obj, micro_batch_data) + + self._call_hooks("before_forward", data) + if hasattr(gpc.config.model, "num_experts"): + output_obj, moe_losses = self._call_engine(engine.model[chunk_id], data) + else: + output_obj = self._call_engine(engine.model[chunk_id], data) + # Convert output_obj to fp32 when last model chunk of last stage + if gpc.is_pipeline_last_stage(ignore_virtual=False) and isinstance(engine.model[chunk_id], NaiveAMPModel): + output_obj = engine.model[chunk_id].convert_to_fp32(output_obj) + self._call_hooks("after_forward", output_obj) + + if gpc.is_pipeline_last_stage(): + self._call_hooks("post_helper_func", output_obj, label) + + if self._return_tensors is not None: + self._return_tensors.append((output_obj, label)) + if self._accum_loss is not None: + self._call_hooks("before_criterion", output_obj, label) + loss = self._call_engine_criterion(engine, output_obj, label) + self._call_hooks("after_criterion", loss) + + loss_reduced = loss / self.num_microbatches + self._accum_loss.add_(loss_reduced.detach()) + output_obj = loss_reduced + + moe_loss = ( + sum(moe_losses) * gpc.config.loss.moe_loss_coeff + if hasattr(gpc.config.model, "num_experts") and gpc.config.model.num_experts > 1 + else torch.tensor(0.0, device=get_current_device(), dtype=gpc.config.model.get("dtype")) + ) + # the moe_loss is computed among the "tensor" group if sequence parallel is enabled, so we need to do allreduce + if gpc.config.parallel.sequence_parallel: + dist.all_reduce(moe_loss, op=dist.ReduceOp.AVG, group=gpc.get_group(ParallelMode.TENSOR)) + moe_loss /= self.num_microbatches + + if self._accum_moe_loss is not None: + self._accum_moe_loss.add_(moe_loss.detach()) + + self._output_objs[chunk_id].append(output_obj) + self._moe_losses[chunk_id].append(moe_loss) + + return output_obj + + def _backward_step(self, engine, chunk_id, step_id): + """ + Backward step for passed-in model. If it is the last stage, the input tensor + is obtained from the previous forward step, otherwise the passed-in input_obj is used. + Returns input tensor gradient. This is a helper function and can be ignored by users. + + Args: + engine (colossalai.engine.Engine): Colossalai engine for training and inference. + chunk_id (int): The id of model chunks. + step_id (int): The current step id. + + Returns: + Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: input tensor gradient. + """ + + gpc.set_virtual_pipeline_parallel_rank(chunk_id) + + if gpc.is_pipeline_last_stage() and len(self._output_obj_grads[chunk_id]) == 0: + self._output_obj_grads[chunk_id].append(None) + + input_obj = self._input_objs_for_backward[chunk_id].pop(0) + output_obj = self._output_objs[chunk_id].pop(0) + output_obj_grad = self._output_obj_grads[chunk_id].pop(0) + moe_loss = self._moe_losses[chunk_id].pop(0) + + input_obj_grad = super()._backward_step(engine, step_id, input_obj, output_obj, output_obj_grad, moe_loss) + + return input_obj_grad + + def _get_chunk_by_microbatch_gpipe(self, step_id: int, backward: bool = False) -> int: + """Helper method to get the model chunk ID given the iteration number.""" + chunk_id = step_id // self._cur_unit_schedule_size + if backward: + chunk_id = self._num_chunks - chunk_id - 1 + return chunk_id + + def _get_current_microbatch_id_gpipe(self, step_id: int) -> int: + microbatch_id = step_id % self._cur_unit_schedule_size + return microbatch_id + + def _get_chunk_by_microbatch_1f1b(self, step_id: int, backward: bool = False) -> int: + """Helper method to get the model chunk ID given the iteration number.""" + microbatch_id_in_group = step_id % (self._pp_size * self._num_chunks) + chunk_id = microbatch_id_in_group // self._pp_size + if backward: + chunk_id = self._num_chunks - chunk_id - 1 + return chunk_id + + def _get_current_microbatch_id_1f1b(self, step_id: int) -> int: + num_microbatch_group = step_id // (self._pp_size * self._num_chunks) + step_id_in_group = step_id % (self._pp_size * self._num_chunks) + microbatch_id = num_microbatch_group * self._pp_size + step_id_in_group % self._pp_size + return microbatch_id + + def _run_warmup_loop( + self, + engine: Engine, + num_microsteps: int, + num_warmup_microsteps: int, + receive_extra_backward: bool = False, + forward_only: bool = False, + ) -> None: + """ + Run the warm-up loop and prepare data for the 1F1B stage. + + During the warm-up process, for each execution, it first performs a forward computation, + and then sends the computation result to the next stage. + It also receives data for the next forward computation. + Since the input for the first forward computation is not considered initially, + it needs to receive data once at the beginning. + + After the warm-up is completed, we need to prepare data for the 1F1B stage. + The data preparation process should be consistent with the communication method of the 1F1B stage. + + Args: + engine (Engine): The engine to run the warm-up loop. + num_microsteps (int): The total number of microsteps. + num_warmup_microsteps (int): The number of warm-up microsteps. + receive_extra_backward (bool, optional): Whether to receive extra backward input for the 1F1B stage. + Default is False. + forward_only (bool, optional): Whether to only perform forward pass. Default is False. + """ + if not gpc.is_pipeline_first_stage(): + if self._input_obj_shapes[0] is None: + self._input_obj_shapes[0] = comm.recv_obj_meta() + self._input_objs[0].append( + comm.recv_forward( + self._input_obj_shapes[0], + dtype=self.dtype, + scatter_gather_tensors=self.scatter_gather_tensors, + ) + ) + else: + self._input_objs[0].append(None) + + for k in range(num_warmup_microsteps): + chunk_id = self._get_chunk_by_microbatch(k) + + output_obj = self._forward_step(engine, chunk_id) + + if forward_only: + # when forward-only, no need to save tensors for a backward pass + self._input_objs[chunk_id].pop() + self._output_objs[chunk_id].pop() + self._moe_losses[chunk_id].pop() + + if not gpc.is_pipeline_last_stage(): + if isinstance(output_obj, torch.Tensor): + self._output_obj_shapes[chunk_id] = output_obj.shape + else: + self._output_obj_shapes[chunk_id] = [out_tensor.shape for out_tensor in output_obj] + + if self._send_tensor_shape_flags[chunk_id]: + comm.send_obj_meta(output_obj) + self._send_tensor_shape_flags[chunk_id] = False # send only once for each chunk. + + if gpc.is_pipeline_first_stage(ignore_virtual=True): + next_forward_chunk_id = self._get_chunk_by_microbatch(k - (self._pp_size - 1)) + next_forward_chunk_id += 1 + else: + next_forward_chunk_id = self._get_chunk_by_microbatch(k + 1) + next_forward_chunk_id = next_forward_chunk_id % self._num_chunks + + with switch_virtual_pipeline_parallel_rank(next_forward_chunk_id): + if not gpc.is_pipeline_first_stage() and self._input_obj_shapes[next_forward_chunk_id] is None: + self._input_obj_shapes[next_forward_chunk_id] = comm.recv_obj_meta() + if k == (num_microsteps - 1) or gpc.is_pipeline_first_stage(): + input_shape = None + else: + input_shape = self._input_obj_shapes[next_forward_chunk_id] + + # Don't send tensor downstream if on last stage. + if gpc.is_pipeline_last_stage(): + output_obj = None + + assert output_obj is None or output_obj.dtype == self.dtype + + # Send and receive tensors as appropriate (send tensors computed + # in this iteration; receive tensors for next iteration). + if k != (num_warmup_microsteps - 1) or not receive_extra_backward: + # Normal warm-up communication process, or no need to prepare backward input for the 1F1B stage + input_obj = comm.send_forward_recv_forward( + output_obj, + input_shape, + dtype=self.dtype, + scatter_gather_tensors=self.scatter_gather_tensors, + ) + else: + # In this case, we should handle forward and backward communication separately, consistent with the + # overlap version of the 1F1B stage + input_obj = comm.send_forward_recv_forward( + output_obj, + input_shape, + dtype=self.dtype, + scatter_gather_tensors=self.scatter_gather_tensors, + ) + output_obj_grad = comm.send_backward_recv_backward( + None, # nothing to send + self._output_obj_shapes[self._num_chunks - 1], + dtype=self.dtype, + scatter_gather_tensors=self.scatter_gather_tensors, + ) + self._output_obj_grads[self._num_chunks - 1].append(output_obj_grad) + + self._input_objs[next_forward_chunk_id].append(input_obj) + + + def pop_decoupled_grads(self, engine): + params = ZeroppManager.pop() + for param in params: + engine.optimizer._wait_reduce_scatter_and_accumulate_grads(param, skip_decoupled_grad_accum=False) + self._call_hooks("after_backward", None) + + + def _run_1f1b_loop_with_overlap( + self, + engine: Engine, + num_warmup_microsteps: int, + num_1f1b_micropairs: int, + all_warmup_microsteps: bool = False, + ) -> None: + """ + Args: + engine (Engine): The engine to run the 1F1B loop. + num_warmup_microsteps (int): The number of warm-up microsteps. + num_1f1b_micropairs (int): The number of 1F1B micropairs. + all_warmup_microsteps (bool, optional): Whether to run all warm-up microsteps. Default is False. + """ + + forward_async_communicator = None + backward_async_communicator = None + + # Run 1F1B in steady state. + for k in range(num_1f1b_micropairs): + forward_microstep_id = k + num_warmup_microsteps + backward_microstep_id = k + forward_chunk_id = self._get_chunk_by_microbatch(forward_microstep_id) + backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id, backward=True) + + if forward_async_communicator is not None: + input_obj = forward_async_communicator.wait_and_receive() + if forward_async_communicator.need_receive: + self._input_objs[next_forward_chunk_id].append(input_obj) + + # 1. Forward pass. + output_obj = self._forward_step(engine, forward_chunk_id) + + # 3. Send the forward outputs and receive the forward inputs from the previous rank. + + # Check if it is the last model chunk of the last pipeline stage, no need to send forward output. + gpc.set_virtual_pipeline_parallel_rank(forward_chunk_id) + if gpc.is_pipeline_last_stage(): + output_obj = None + + if gpc.is_pipeline_first_stage(ignore_virtual=True): + next_forward_chunk_id = self._get_chunk_by_microbatch(forward_microstep_id - (self._pp_size - 1)) + next_forward_chunk_id += 1 + else: + next_forward_chunk_id = self._get_chunk_by_microbatch(forward_microstep_id + 1) + next_forward_chunk_id = next_forward_chunk_id % self._num_chunks + + with switch_virtual_pipeline_parallel_rank(next_forward_chunk_id): + if gpc.is_pipeline_first_stage() or k == num_1f1b_micropairs - 1: + input_obj_shape = None + else: + input_obj_shape = self._input_obj_shapes[next_forward_chunk_id] + + assert output_obj is None or output_obj.dtype == self.dtype + + forward_async_communicator = comm.AsynCommunicator( + output_obj, + input_obj_shape, + self.dtype, + self.scatter_gather_tensors, + forward=True, + ) + forward_async_communicator.start() + + # 2. Check if the backward input is ready. + if backward_async_communicator is not None: + output_obj_grad = backward_async_communicator.wait_and_receive() + + if backward_async_communicator.need_receive: + self._output_obj_grads[next_backward_chunk_id].append(output_obj_grad) + + # 5. Backward pass. + + input_obj_grad = self._backward_step(engine, backward_chunk_id, backward_microstep_id) + ZeroppManager.flush() + + # 6. Send the backward output and receive the backward input for the next iteration. + gpc.set_virtual_pipeline_parallel_rank(backward_chunk_id) + if gpc.is_pipeline_first_stage(): + input_obj_grad = None + + #next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id + 1, backward=True) + + if gpc.is_pipeline_last_stage(ignore_virtual=True): + next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id - (self._pp_size - 1), backward=True) + next_backward_chunk_id -= 1 + else: + next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id + 1, backward=True) + next_backward_chunk_id = next_backward_chunk_id % self._num_chunks + + with switch_virtual_pipeline_parallel_rank(next_backward_chunk_id): + if gpc.is_pipeline_last_stage(): + output_obj_shape = None + else: + output_obj_shape = self._output_obj_shapes[next_backward_chunk_id] + + backward_async_communicator = comm.AsynCommunicator( + input_obj_grad, + output_obj_shape, + self.dtype, + self.scatter_gather_tensors, + forward=False, + ) + backward_async_communicator.start() + + if self.decouple_grad: + if backward_chunk_id == self._num_chunks - 1: + if backward_microstep_id > self._pp_rank*2: + self.pop_decoupled_grads(engine) + + if backward_microstep_id == self._cur_unit_schedule_size - 1: + for _ in range((self._pp_size - self._pp_rank - 1 )*2): + if not ZeroppManager.empty(): + self.pop_decoupled_grads(engine) + else: + self.pop_decoupled_grads(engine) + else: + self.pop_decoupled_grads(engine) + + if all_warmup_microsteps: + if not gpc.is_pipeline_last_stage(): + self._output_obj_grads[self._num_chunks - 1].append( + comm.recv_backward( + self._output_obj_shapes[self._num_chunks - 1], + dtype=self.dtype, + scatter_gather_tensors=self.scatter_gather_tensors, + ) + ) + else: + self._output_obj_grads[self._num_chunks - 1].append(None) + else: + output_obj_grad = backward_async_communicator.wait_and_receive() + if backward_async_communicator.need_receive: + self._output_obj_grads[next_backward_chunk_id].append(output_obj_grad) + + def _run_cooldown_loop(self, engine: Engine, num_microsteps: int, num_1f1b_micropairs: int) -> None: + """ + Run the cooldown loop. + + The cooldown loop consists of the following steps: + 1. Perform the backward step. + 2. Send the backward output to the next stage and receive inputs for next backward. + + Args: + engine (Engine): The engine to use for computation. + num_microsteps (int): The total number of microsteps. + num_1f1b_micropairs (int): The number of 1F1B micro-pairs. + """ + #if gpc.get_local_rank(ParallelMode.WEIGHT) == 0: + # logger.info(f"pp_rank: {gpc.get_local_rank(ParallelMode.PIPELINE)}, _run_cooldown_loop") + for k in range(num_1f1b_micropairs, num_microsteps): + chunk_id = self._get_chunk_by_microbatch(k, backward=True) + + input_obj_grad = self._backward_step(engine, chunk_id, k) + ZeroppManager.flush() + + #next_backward_chunk_id = self._get_chunk_by_microbatch(k + 1, backward=True) + if gpc.is_pipeline_last_stage(ignore_virtual=True): + next_backward_chunk_id = self._get_chunk_by_microbatch(k - (self._pp_size - 1), backward=True) + next_backward_chunk_id -= 1 + else: + next_backward_chunk_id = self._get_chunk_by_microbatch(k + 1, backward=True) + next_backward_chunk_id = next_backward_chunk_id % self._num_chunks + + if k != (num_microsteps - 1) and not ( + gpc.is_pipeline_last_stage(ignore_virtual=True) and next_backward_chunk_id == (self._num_chunks - 1) + ): + output_shape = self._output_obj_shapes[next_backward_chunk_id] + else: + output_shape = None + + backward_async_communicator = comm.AsynCommunicator( + input_obj_grad, + output_shape, + self.dtype, + self.scatter_gather_tensors, + forward=False, + ) + backward_async_communicator.start() + + self.pop_decoupled_grads(engine) + + output_obj_grad = backward_async_communicator.wait_and_receive() + if backward_async_communicator.need_receive: + self._output_obj_grads[next_backward_chunk_id].append(output_obj_grad) + else: + self._output_obj_grads[next_backward_chunk_id].append(None) + + + while not ZeroppManager.empty(): + self.pop_decoupled_grads(engine) + + + def _forward_only_step(self, engine: Engine): + num_microsteps = self.num_microbatches * self._num_chunks + num_warmup_microsteps = num_microsteps + + self._run_warmup_loop( + engine, + num_microsteps, + num_warmup_microsteps, + receive_extra_backward=False, + forward_only=True, + ) + + + def _unit_schedule_step(self, engine: Engine, unit_schedule_size): + gpc.set_virtual_pipeline_parallel_rank(0) + self._cur_unit_schedule_size = unit_schedule_size + # Compute number of warmup and remaining microbatches. + + + unit_num_microsteps = unit_schedule_size * self._num_chunks + + # Run all forward passes and then all backward passes if number of + # microbatches is just the number of pipeline stages. + # Otherwise, perform (num_chunks-1)*pipeline_parallel_size on + # all workers, followed by more microbatches after depending on + # stage ID (more forward passes for earlier stages, later stages can + # immediately start with 1F1B). + + + num_warmup_steps = (self._pp_size - self._pp_rank - 1) * 2 + num_warmup_steps += (self._num_chunks - 1) * unit_schedule_size + num_warmup_steps = min(num_warmup_steps, unit_num_microsteps) + num_1f1b_micropairs = unit_num_microsteps - num_warmup_steps + + all_warmup_microsteps = (num_1f1b_micropairs == 0) + + # We usually need to prepare an extra backward data for the 1F1B stage when the WarmUp stage ends, + # because the 1F1B stage typically performs one forward and backward pass together, + # except in the following cases: + receive_extra_backward = not ( + all_warmup_microsteps # Only warmup microsteps + or gpc.is_pipeline_last_stage(ignore_virtual=True) # The rank is the last pipeline stage + ) + + #if gpc.get_local_rank(ParallelMode.WEIGHT) == 0: + # logger.info(f"pp_rank: {gpc.get_local_rank(ParallelMode.PIPELINE)}, receive_extra_backward: {receive_extra_backward}, num_warmup_steps: {num_warmup_steps}") + + # 1. Warmup + self._run_warmup_loop( + engine, + unit_num_microsteps, + num_warmup_steps, + receive_extra_backward=receive_extra_backward, + ) + + # 2. 1F1B + self._run_1f1b_loop( + engine, + num_warmup_steps, + num_1f1b_micropairs=num_1f1b_micropairs, + all_warmup_microsteps=all_warmup_microsteps, + ) + + # 3. Cooldown + self._run_cooldown_loop(engine, unit_num_microsteps, num_1f1b_micropairs=num_1f1b_micropairs) + self._clear_state_unit() + + + def _forward_backward_step_gpipe(self, engine: Engine): + + last_unit_schedule_size = self.num_microbatches % self._unit_schedule_size + num_unit_schedules = self.num_microbatches // self._unit_schedule_size + if last_unit_schedule_size > 0: + num_unit_schedules += 1 + else: + last_unit_schedule_size = self._unit_schedule_size + + for i in range(num_unit_schedules - 1): + self._unit_schedule_step(engine, self._unit_schedule_size) + self._unit_schedule_step(engine, last_unit_schedule_size) + + + def _forward_backward_step_1f1b(self, engine: Engine): + # Compute number of warmup and remaining microbatches. + all_warmup_microsteps = False + num_microsteps = self.num_microbatches * self._num_chunks + + # Run all forward passes and then all backward passes if number of + # microbatches is just the number of pipeline stages. + # Otherwise, perform (num_chunks-1)*pipeline_parallel_size on + # all workers, followed by more microbatches after depending on + # stage ID (more forward passes for earlier stages, later stages can + # immediately start with 1F1B). + if self.num_microbatches == self._pp_size: + num_warmup_steps = num_microsteps + all_warmup_microsteps = True + else: + num_warmup_steps = (self._pp_size - self._pp_rank - 1) * 2 + num_warmup_steps += (self._num_chunks - 1) * self._pp_size + num_warmup_steps = min(num_warmup_steps, num_microsteps) + num_1f1b_micropairs = num_microsteps - num_warmup_steps + + # We usually need to prepare an extra backward data for the 1F1B stage when the WarmUp stage ends, + # because the 1F1B stage typically performs one forward and backward pass together, + # except in the following cases: + receive_extra_backward = not ( + all_warmup_microsteps # Only warmup microsteps + or gpc.is_pipeline_last_stage(ignore_virtual=True) # The rank is the last pipeline stage + ) + + # 1. Warmup + self._run_warmup_loop( + engine, + num_microsteps, + num_warmup_steps, + receive_extra_backward=receive_extra_backward, + ) + + # 2. 1F1B + self._run_1f1b_loop( + engine, + num_warmup_steps, + num_1f1b_micropairs=num_1f1b_micropairs, + all_warmup_microsteps=all_warmup_microsteps, + ) + + # 3. Cooldown + self._run_cooldown_loop(engine, num_microsteps, num_1f1b_micropairs=num_1f1b_micropairs) + + + @llm_timeout(func_name="zeropp_forward_backward_step") + def forward_backward_step(self, engine, data_iter, forward_only=False, return_loss=True, return_output_label=True): + """Run interleaved 1F1B schedule (model split into model chunks), with + communication between pipeline stages as needed. + + Args: + engine (colossalai.engine.Engine): Colossalai engine for training and inference. + data_iter (Iterable): Dataloader as the form of an iterator, obtained by calling iter(dataloader). + forward_only (bool, optional): + Whether run forward step only. Default is false. If true, no backward will be run. + return_loss (bool, optional): Whether returns the loss value. Default is true. + return_output_label (bool, optional): If False, the output and label won't be returned. + + Returns: + Tuple[:class:`torch.Tensor`]: A tuple of (output, label, loss, moe_loss), loss and label could be None. + The loss would be returned only in the last stage. And the moe_loss is accumulated from all stages. + """ + assert ( + forward_only or return_loss + ), "The argument 'return_loss' has to be True when 'forward_only' is False, but got False." + + gpc.set_virtual_pipeline_parallel_rank(0) + + self.load_batch(engine, data_iter) + + if return_loss and gpc.is_pipeline_last_stage(ignore_virtual=True): + self._accum_loss = torch.zeros(1, device=get_current_device()) + self._accum_moe_loss = torch.zeros(1, device=get_current_device()) + + if return_output_label: + self._return_tensors = [] + + if forward_only: + self._forward_only_step(engine) + else: + self._forward_backward_step(engine) + + if return_output_label and len(self._return_tensors) > 0: + output, label = pack_return_tensors(self._return_tensors) + else: + output, label = (None, None) + + if hasattr(gpc.config.model, "num_experts") and gpc.config.model.num_experts > 1: + dist.all_reduce(self._accum_moe_loss, group=gpc.get_group(ParallelMode.PIPELINE)) + accum_moe_loss = self._accum_moe_loss + + accum_loss = self._accum_loss + if accum_loss is not None: + accum_loss += self._accum_moe_loss + + self._clear_state() + + # Compatible for non-moe + if hasattr(gpc.config.model, "num_experts"): + return output, label, accum_loss, accum_moe_loss + else: + return output, label, accum_loss diff --git a/internlm/core/trainer.py b/internlm/core/trainer.py index 121501579..995a9f725 100644 --- a/internlm/core/trainer.py +++ b/internlm/core/trainer.py @@ -14,6 +14,7 @@ InterleavedPipelineScheduler, NonPipelineScheduler, PipelineScheduler, + ZeroPPScheduler, ) @@ -181,7 +182,7 @@ def schedule(self): @property def uses_pipeline(self): """Returns whether the pipeline parallel is used or not.""" - return isinstance(self._schedule, (PipelineScheduler, InterleavedPipelineScheduler)) + return isinstance(self._schedule, (PipelineScheduler, InterleavedPipelineScheduler, ZeroPPScheduler)) def train(self): """Sets the model to training mode.""" diff --git a/internlm/initialize/initialize_trainer.py b/internlm/initialize/initialize_trainer.py index b90a25e9f..3402e56d9 100644 --- a/internlm/initialize/initialize_trainer.py +++ b/internlm/initialize/initialize_trainer.py @@ -19,6 +19,7 @@ InterleavedPipelineScheduler, NonPipelineScheduler, PipelineScheduler, + ZeroPPScheduler, ) from internlm.core.scheduler.pipeline_scheduler import get_tensor_shape from internlm.core.trainer import Trainer @@ -94,16 +95,32 @@ def initialize_trainer( model = nn.ModuleList([model]) communication_overlap = gpc.config.parallel["pipeline"].get("interleaved_overlap", False) - scheduler = InterleavedPipelineScheduler( - data_process_func=data_fn, - num_microbatches=gpc.config.NUM_MICRO_BATCHES, - num_chunks=gpc.config.model.num_chunks, - dtype=gpc.config.model["dtype"], - tensor_shape=tensor_shape, - scatter_gather_tensors=scatter_gather, - scheduler_hooks=scheduler_hooks, - communication_overlap=communication_overlap, - ) + use_zeropp = hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp + if use_zeropp: + unit_schedule_size = gpc.config.data["unit_schedule_size"] if(hasattr(gpc.config, "data") and hasattr(gpc.config.data, "unit_schedule_size")) \ + else None + scheduler = ZeroPPScheduler( + data_process_func=data_fn, + num_microbatches=gpc.config.NUM_MICRO_BATCHES, + num_chunks=gpc.config.model.num_chunks, + dtype=gpc.config.model["dtype"], + tensor_shape=tensor_shape, + scatter_gather_tensors=scatter_gather, + scheduler_hooks=scheduler_hooks, + communication_overlap=communication_overlap, + unit_schedule_size=unit_schedule_size + ) + else: + scheduler = InterleavedPipelineScheduler( + data_process_func=data_fn, + num_microbatches=gpc.config.NUM_MICRO_BATCHES, + num_chunks=gpc.config.model.num_chunks, + dtype=gpc.config.model["dtype"], + tensor_shape=tensor_shape, + scatter_gather_tensors=scatter_gather, + scheduler_hooks=scheduler_hooks, + communication_overlap=communication_overlap, + ) else: scheduler = PipelineScheduler( data_process_func=data_fn, diff --git a/internlm/model/modules/linear.py b/internlm/model/modules/linear.py index 0d8c4bf82..1e8c829c7 100644 --- a/internlm/model/modules/linear.py +++ b/internlm/model/modules/linear.py @@ -20,6 +20,7 @@ ) from internlm.model.ops.linear import linear_backward_op, linear_forward_op from internlm.utils.logger import get_logger +from internlm.utils.zeropp_manager import ZeroppManager if TYPE_CHECKING: from internlm.core.parallel.comm.isp import WPCommunicator @@ -32,6 +33,22 @@ custom_fwd = internlm_accelerator.return_custom_fwd() +def decoupled_grad_func(input, grad_output, module, communicator): + grad_weight, grad_bias = linear_backward_op( + input, + grad_output, + module.bias is not None and module.bias.requires_grad, + ) + communicator.grad_hook( + grad_weight, async_op=True, module=module, is_bias=False + ) + if grad_bias is not None: + communicator.grad_hook( + grad_bias, async_op=True, module=module, is_bias=True + ) + return [module.weight, module.bias] + return [module.weight,] + # adpated from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/fused_dense.py class SPFusedDenseFunc(torch.autograd.Function): "FusedDenseFunc for tensor parallel in flash-attn implementation." @@ -170,8 +187,22 @@ def forward( x = x.to(dtype=torch.get_autocast_gpu_dtype()) x = x.contiguous() - total_weight = communicator.weight_hook(weight, module=module) - total_bias = bias if bias is None else communicator.weight_hook(bias, module=module, is_bias=True) + use_zeropp = hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp + if use_zeropp: + total_weight = ZeroppManager.retrieve_full_wp_parameters(module.weight) + if total_weight is None: + total_weight = communicator.weight_hook(weight, module=module) + ZeroppManager.cache_full_wp_parameters(module.weight, total_weight) + if bias is None: + total_bias = bias + else: + total_bias = ZeroppManager.retrieve_full_wp_parameters(module.bias) + if total_bias is None: + total_bias = communicator.weight_hook(bias, module=module) + ZeroppManager.cache_full_wp_parameters(module.bias, total_bias) + else: + total_weight = communicator.weight_hook(weight, module=module) + total_bias = bias if bias is None else communicator.weight_hook(bias, module=module, is_bias=True) if torch.is_autocast_enabled(): total_weight = total_weight.to(dtype=torch.get_autocast_gpu_dtype()) @@ -187,9 +218,10 @@ def forward( output = linear_forward_op(x, total_weight, total_bias) - # release memory - del total_weight - del total_bias + if not use_zeropp: + # release memory + del total_weight + del total_bias saved_x = None if ctx.compute_weight_gradient is False else x ctx.save_for_backward(saved_x, weight) @@ -212,24 +244,36 @@ def backward(ctx, grad_output, *args): batch_dim = batch_shape.numel() grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) - total_weight = communicator.weight_hook(weight, module=module) + use_zeropp = hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp + if use_zeropp: + total_weight = ZeroppManager.retrieve_full_wp_parameters(module.weight) + if total_weight is None: + total_weight = communicator.weight_hook(weight, module=module) + ZeroppManager.cache_full_wp_parameters(module.weight, total_weight) + else: + total_weight = communicator.weight_hook(weight, module=module) # compute weight grad if ctx.needs_input_grad[1]: assert ctx.compute_weight_gradient - grad_weight, grad_bias = linear_backward_op( - x.reshape(batch_dim, x.shape[-1]), - grad_output, - ctx.needs_input_grad[2], - ) + if use_zeropp: + grad_weight = None + grad_bias = None + ZeroppManager.put(x.reshape(batch_dim, x.shape[-1]), grad_output, module, communicator, decoupled_grad_func) + else: + grad_weight, grad_bias = linear_backward_op( + x.reshape(batch_dim, x.shape[-1]), + grad_output, + ctx.needs_input_grad[2], + ) - grad_weight, grad_weight_sync = communicator.grad_hook( - grad_weight, async_op=True, module=module, is_bias=False - ) - if grad_bias is not None: - grad_bias, grad_bias_sync = communicator.grad_hook( - grad_bias, async_op=True, module=module, is_bias=True + grad_weight, grad_weight_sync = communicator.grad_hook( + grad_weight, async_op=True, module=module, is_bias=False ) + if grad_bias is not None: + grad_bias, grad_bias_sync = communicator.grad_hook( + grad_bias, async_op=True, module=module, is_bias=True + ) else: grad_weight = None grad_bias = grad_output if ctx.needs_input_grad[2] else None @@ -247,12 +291,12 @@ def backward(ctx, grad_output, *args): else: grad_input = None - del total_weight - - if ctx.needs_input_grad[1]: - grad_weight_sync.wait() - if grad_bias is not None: - grad_bias_sync.wait() + if not use_zeropp: + del total_weight + if ctx.needs_input_grad[1]: + grad_weight_sync.wait() + if grad_bias is not None: + grad_bias_sync.wait() return grad_input, grad_weight, grad_bias, None, None, None, None diff --git a/internlm/solver/optimizer/hybrid_zero_optim.py b/internlm/solver/optimizer/hybrid_zero_optim.py index 5461f9228..65e2f6b17 100644 --- a/internlm/solver/optimizer/hybrid_zero_optim.py +++ b/internlm/solver/optimizer/hybrid_zero_optim.py @@ -43,6 +43,7 @@ from internlm.utils.megatron_timers import megatron_timer as timer from internlm.utils.parallel import is_using_isp, is_using_sequence_parallel from internlm.utils.timeout import llm_timeout +from internlm.utils.zeropp_manager import ZeroppManager from .base_optimizer import BaseOptimizer from .utils import compute_norm @@ -401,7 +402,10 @@ def _accum_grads_store_in_bucket(self, bucket: BucketStore, reduce_rank: Optiona _key = getattr(_param, "isp_reduce_scatter_name") _grad, _comm_handle = self._isp_communicator.reduce_scatter_handlers[_key] _comm_handle.wait() - _param.grad.add_(_grad) + if _param.grad is None: + _param.grad = _grad.clone() + else: + _param.grad.add_(_grad) # release cuda memory. if self._isp_communicator.enable_memory_pool: @@ -413,7 +417,7 @@ def _accum_grads_store_in_bucket(self, bucket: BucketStore, reduce_rank: Optiona bucket.reset_by_rank(reduce_rank) - def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional[int] = None): + def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional[int] = None, skip_decoupled_grad_accum = True): param_size = param.numel() group_id = getattr(param, "group_id") @@ -426,8 +430,9 @@ def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional self._accum_grads_store_in_bucket(current_bucket, reduce_rank) # otherwise, add the parameter into bucket. - current_bucket.add_num_elements_in_bucket(param_size, reduce_rank) - current_bucket.add_param(param, reduce_rank) + if not ZeroppManager.check_postpond_grad_accum(param) or not skip_decoupled_grad_accum: + current_bucket.add_num_elements_in_bucket(param_size, reduce_rank) + current_bucket.add_param(param, reduce_rank) def _store_and_try_reduce_grads_by_bucket(self, param, reduce_rank=None): param_size = param.numel() diff --git a/internlm/utils/zeropp_manager.py b/internlm/utils/zeropp_manager.py new file mode 100644 index 000000000..58f06fe4b --- /dev/null +++ b/internlm/utils/zeropp_manager.py @@ -0,0 +1,47 @@ +import queue + +class ZeroppManager: + cache = [] + cached_wp_params = {} + grad_queue = queue.LifoQueue() + + @classmethod + def put(cls, total_input, grad_output, mod, communicator, func): + # Store the weight gradient computation of linear layers. + cls.cache.append((total_input, grad_output, mod, communicator, func)) + + @classmethod + def flush(cls): + # Collect all stored computations during backward as a W pass. + cls.grad_queue.put(cls.cache) + cls.cache = [] + + @classmethod + def empty(cls): + return cls.grad_queue.empty() + + @classmethod + def pop(cls): + item = cls.grad_queue.get() + params = [] + for total_input, grad_output, mod, communicator, func in item: + params.extend(func(total_input, grad_output, mod, communicator)) + return params + + @classmethod + def cache_full_wp_parameters(cls, shard_param, full_param): + cls.cached_wp_params[shard_param] = full_param + + @classmethod + def retrieve_full_wp_parameters(cls, shard_param): + return cls.cached_wp_params.get(shard_param, None) + + @classmethod + def clear_cached_wp_parameters(cls): + for shard_param, full_param in cls.cached_wp_params.items(): + del full_param + cls.cached_wp_params[shard_param] = None + + @classmethod + def check_postpond_grad_accum(cls, param): + return param in cls.cached_wp_params From 12e73ce4f07c4062e560f0a785a408620e4829f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Thu, 20 Jun 2024 17:31:43 +0800 Subject: [PATCH 2/6] fix recomp --- internlm/core/parallel/shard.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/internlm/core/parallel/shard.py b/internlm/core/parallel/shard.py index 33c187ec5..f12d444f2 100644 --- a/internlm/core/parallel/shard.py +++ b/internlm/core/parallel/shard.py @@ -95,14 +95,19 @@ def pipeline_parallel_sharding_wrapper( models = [] - for start, end in parts: + use_checkpoint = kwargs["checkpoint"] + for part_idx, (start, end) in enumerate(parts): kwargs["num_layers"] = end - start kwargs["first"] = start == 0 # If there is no content in the final layer, assign the last layer. kwargs["last"] = end == num_layers and len(all_parts[-1]) != 0 kwargs["device"] = device kwargs["start_layer_idx"] = start - + if hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp: + if part_idx == num_chunks - 1 and gpc.is_pipeline_last_stage(ignore_virtual=True): + kwargs["checkpoint"] = False + else: + kwargs["checkpoint"] = use_checkpoint chunk = model_builder(**kwargs).to(device) setattr(chunk, "first_layer", start) setattr(chunk, "last_layer", end) From ebcba6c20ef832f2637739aaad3f356e6feb8ac9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Thu, 20 Jun 2024 20:15:48 +0800 Subject: [PATCH 3/6] fix lint --- internlm/core/scheduler/pipeline_scheduler.py | 56 ++++++------------- internlm/model/modules/linear.py | 8 ++- .../solver/optimizer/hybrid_zero_optim.py | 3 +- internlm/utils/zeropp_manager.py | 13 +++-- 4 files changed, 30 insertions(+), 50 deletions(-) diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py index 3324629f8..326fce451 100644 --- a/internlm/core/scheduler/pipeline_scheduler.py +++ b/internlm/core/scheduler/pipeline_scheduler.py @@ -1420,8 +1420,6 @@ def forward_backward_step(self, engine, data_iter, forward_only=False, return_lo return output, label, accum_loss - - class ZeroPPScheduler(PipelineScheduler): """ ZeroPP Pipeline Scheduler. @@ -1456,7 +1454,8 @@ def __init__( If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization. scheduler_hooks (List[SchedulerHook], optional): List of scheduler hooks. Default is None. communication_overlap (bool, optional): Whether to enable communication overlap. Default is False. - unit_schedule_size: (int, optional): Unit size if using interleaved GPipe schedule. Default is None and interleaved 1F1B schedule is used. + unit_schedule_size: (int, optional): + Unit size if using interleaved GPipe schedule. Default is None and interleaved 1F1B schedule is used. """ assert ( isinstance(num_chunks, int) and num_chunks > 0 @@ -1506,9 +1505,10 @@ def __init__( self._forward_backward_step = self._forward_backward_step_1f1b self._get_current_microbatch_id = self._get_chunk_by_microbatch_1f1b self._get_chunk_by_microbatch = self._get_chunk_by_microbatch_1f1b - + self._cur_unit_schedule_size = self._unit_schedule_size - if hasattr(gpc.config.parallel.pipeline, "decouple_grad") and gpc.config.parallel.pipeline.decouple_grad is True: + if hasattr(gpc.config.parallel.pipeline, "decouple_grad") \ + and gpc.config.parallel.pipeline.decouple_grad is True: self.decouple_grad = True else: self.decouple_grad = False @@ -1536,7 +1536,7 @@ def _clear_state(self) -> None: self._input_obj_shapes = [self.tensor_shape for _ in range(self._num_chunks)] self._output_obj_shapes = [None for _ in range(self._num_chunks)] self._send_tensor_shape_flags = [self.tensor_shape is None for _ in range(self._num_chunks)] - + def _clear_state_unit(self) -> None: self._input_objs = [[] for _ in range(self._num_chunks)] self._input_objs_for_backward = [[] for _ in range(self._num_chunks)] @@ -1579,12 +1579,12 @@ def _forward_step(self, engine, chunk_id): Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: output or the loss value of the current pipeline stage. """ - - gpc.set_virtual_pipeline_parallel_rank(chunk_id) - if gpc.is_pipeline_first_stage() and (len(self._input_objs[chunk_id]) + len(self._input_objs_for_backward[chunk_id])) == len(self._output_objs[chunk_id]): - self._input_objs[chunk_id].append(None) + if gpc.is_pipeline_first_stage(): + if (len(self._input_objs[chunk_id]) + len(self._input_objs_for_backward[chunk_id])) \ + == len(self._output_objs[chunk_id]): + self._input_objs[chunk_id].append(None) input_obj = self._input_objs[chunk_id].pop(0) self._input_objs_for_backward[chunk_id].append(input_obj) @@ -1648,7 +1648,6 @@ def _backward_step(self, engine, chunk_id, step_id): Returns: Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: input tensor gradient. """ - gpc.set_virtual_pipeline_parallel_rank(chunk_id) if gpc.is_pipeline_last_stage() and len(self._output_obj_grads[chunk_id]) == 0: @@ -1800,14 +1799,12 @@ def _run_warmup_loop( self._input_objs[next_forward_chunk_id].append(input_obj) - def pop_decoupled_grads(self, engine): params = ZeroppManager.pop() for param in params: engine.optimizer._wait_reduce_scatter_and_accumulate_grads(param, skip_decoupled_grad_accum=False) self._call_hooks("after_backward", None) - def _run_1f1b_loop_with_overlap( self, engine: Engine, @@ -1889,10 +1886,9 @@ def _run_1f1b_loop_with_overlap( if gpc.is_pipeline_first_stage(): input_obj_grad = None - #next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id + 1, backward=True) - if gpc.is_pipeline_last_stage(ignore_virtual=True): - next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id - (self._pp_size - 1), backward=True) + next_backward_chunk_id = \ + self._get_chunk_by_microbatch(backward_microstep_id - (self._pp_size - 1), backward=True) next_backward_chunk_id -= 1 else: next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id + 1, backward=True) @@ -1903,7 +1899,7 @@ def _run_1f1b_loop_with_overlap( output_obj_shape = None else: output_obj_shape = self._output_obj_shapes[next_backward_chunk_id] - + backward_async_communicator = comm.AsynCommunicator( input_obj_grad, output_obj_shape, @@ -1915,11 +1911,11 @@ def _run_1f1b_loop_with_overlap( if self.decouple_grad: if backward_chunk_id == self._num_chunks - 1: - if backward_microstep_id > self._pp_rank*2: + if backward_microstep_id > self._pp_rank * 2: self.pop_decoupled_grads(engine) if backward_microstep_id == self._cur_unit_schedule_size - 1: - for _ in range((self._pp_size - self._pp_rank - 1 )*2): + for _ in range((self._pp_size - self._pp_rank - 1) * 2): if not ZeroppManager.empty(): self.pop_decoupled_grads(engine) else: @@ -1956,15 +1952,12 @@ def _run_cooldown_loop(self, engine: Engine, num_microsteps: int, num_1f1b_micro num_microsteps (int): The total number of microsteps. num_1f1b_micropairs (int): The number of 1F1B micro-pairs. """ - #if gpc.get_local_rank(ParallelMode.WEIGHT) == 0: - # logger.info(f"pp_rank: {gpc.get_local_rank(ParallelMode.PIPELINE)}, _run_cooldown_loop") for k in range(num_1f1b_micropairs, num_microsteps): chunk_id = self._get_chunk_by_microbatch(k, backward=True) input_obj_grad = self._backward_step(engine, chunk_id, k) ZeroppManager.flush() - #next_backward_chunk_id = self._get_chunk_by_microbatch(k + 1, backward=True) if gpc.is_pipeline_last_stage(ignore_virtual=True): next_backward_chunk_id = self._get_chunk_by_microbatch(k - (self._pp_size - 1), backward=True) next_backward_chunk_id -= 1 @@ -1987,20 +1980,16 @@ def _run_cooldown_loop(self, engine: Engine, num_microsteps: int, num_1f1b_micro forward=False, ) backward_async_communicator.start() - self.pop_decoupled_grads(engine) - output_obj_grad = backward_async_communicator.wait_and_receive() if backward_async_communicator.need_receive: self._output_obj_grads[next_backward_chunk_id].append(output_obj_grad) else: self._output_obj_grads[next_backward_chunk_id].append(None) - while not ZeroppManager.empty(): self.pop_decoupled_grads(engine) - def _forward_only_step(self, engine: Engine): num_microsteps = self.num_microbatches * self._num_chunks num_warmup_microsteps = num_microsteps @@ -2013,13 +2002,10 @@ def _forward_only_step(self, engine: Engine): forward_only=True, ) - def _unit_schedule_step(self, engine: Engine, unit_schedule_size): gpc.set_virtual_pipeline_parallel_rank(0) self._cur_unit_schedule_size = unit_schedule_size # Compute number of warmup and remaining microbatches. - - unit_num_microsteps = unit_schedule_size * self._num_chunks # Run all forward passes and then all backward passes if number of @@ -2028,14 +2014,11 @@ def _unit_schedule_step(self, engine: Engine, unit_schedule_size): # all workers, followed by more microbatches after depending on # stage ID (more forward passes for earlier stages, later stages can # immediately start with 1F1B). - - num_warmup_steps = (self._pp_size - self._pp_rank - 1) * 2 num_warmup_steps += (self._num_chunks - 1) * unit_schedule_size num_warmup_steps = min(num_warmup_steps, unit_num_microsteps) num_1f1b_micropairs = unit_num_microsteps - num_warmup_steps - - all_warmup_microsteps = (num_1f1b_micropairs == 0) + all_warmup_microsteps = num_1f1b_micropairs == 0 # We usually need to prepare an extra backward data for the 1F1B stage when the WarmUp stage ends, # because the 1F1B stage typically performs one forward and backward pass together, @@ -2045,9 +2028,6 @@ def _unit_schedule_step(self, engine: Engine, unit_schedule_size): or gpc.is_pipeline_last_stage(ignore_virtual=True) # The rank is the last pipeline stage ) - #if gpc.get_local_rank(ParallelMode.WEIGHT) == 0: - # logger.info(f"pp_rank: {gpc.get_local_rank(ParallelMode.PIPELINE)}, receive_extra_backward: {receive_extra_backward}, num_warmup_steps: {num_warmup_steps}") - # 1. Warmup self._run_warmup_loop( engine, @@ -2068,9 +2048,7 @@ def _unit_schedule_step(self, engine: Engine, unit_schedule_size): self._run_cooldown_loop(engine, unit_num_microsteps, num_1f1b_micropairs=num_1f1b_micropairs) self._clear_state_unit() - def _forward_backward_step_gpipe(self, engine: Engine): - last_unit_schedule_size = self.num_microbatches % self._unit_schedule_size num_unit_schedules = self.num_microbatches // self._unit_schedule_size if last_unit_schedule_size > 0: @@ -2081,7 +2059,6 @@ def _forward_backward_step_gpipe(self, engine: Engine): for i in range(num_unit_schedules - 1): self._unit_schedule_step(engine, self._unit_schedule_size) self._unit_schedule_step(engine, last_unit_schedule_size) - def _forward_backward_step_1f1b(self, engine: Engine): # Compute number of warmup and remaining microbatches. @@ -2130,7 +2107,6 @@ def _forward_backward_step_1f1b(self, engine: Engine): # 3. Cooldown self._run_cooldown_loop(engine, num_microsteps, num_1f1b_micropairs=num_1f1b_micropairs) - @llm_timeout(func_name="zeropp_forward_backward_step") def forward_backward_step(self, engine, data_iter, forward_only=False, return_loss=True, return_output_label=True): """Run interleaved 1F1B schedule (model split into model chunks), with diff --git a/internlm/model/modules/linear.py b/internlm/model/modules/linear.py index 1e8c829c7..f3e75cf76 100644 --- a/internlm/model/modules/linear.py +++ b/internlm/model/modules/linear.py @@ -47,7 +47,8 @@ def decoupled_grad_func(input, grad_output, module, communicator): grad_bias, async_op=True, module=module, is_bias=True ) return [module.weight, module.bias] - return [module.weight,] + return [module.weight, ] + # adpated from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/fused_dense.py class SPFusedDenseFunc(torch.autograd.Function): @@ -259,14 +260,15 @@ def backward(ctx, grad_output, *args): if use_zeropp: grad_weight = None grad_bias = None - ZeroppManager.put(x.reshape(batch_dim, x.shape[-1]), grad_output, module, communicator, decoupled_grad_func) + ZeroppManager.put( + x.reshape(batch_dim, x.shape[-1]), grad_output, module, communicator, decoupled_grad_func + ) else: grad_weight, grad_bias = linear_backward_op( x.reshape(batch_dim, x.shape[-1]), grad_output, ctx.needs_input_grad[2], ) - grad_weight, grad_weight_sync = communicator.grad_hook( grad_weight, async_op=True, module=module, is_bias=False ) diff --git a/internlm/solver/optimizer/hybrid_zero_optim.py b/internlm/solver/optimizer/hybrid_zero_optim.py index 65e2f6b17..f58df02b4 100644 --- a/internlm/solver/optimizer/hybrid_zero_optim.py +++ b/internlm/solver/optimizer/hybrid_zero_optim.py @@ -417,7 +417,8 @@ def _accum_grads_store_in_bucket(self, bucket: BucketStore, reduce_rank: Optiona bucket.reset_by_rank(reduce_rank) - def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional[int] = None, skip_decoupled_grad_accum = True): + def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional[int] = None, + skip_decoupled_grad_accum=True): param_size = param.numel() group_id = getattr(param, "group_id") diff --git a/internlm/utils/zeropp_manager.py b/internlm/utils/zeropp_manager.py index 58f06fe4b..ec1c23bb4 100644 --- a/internlm/utils/zeropp_manager.py +++ b/internlm/utils/zeropp_manager.py @@ -1,5 +1,6 @@ import queue + class ZeroppManager: cache = [] cached_wp_params = {} @@ -15,11 +16,11 @@ def flush(cls): # Collect all stored computations during backward as a W pass. cls.grad_queue.put(cls.cache) cls.cache = [] - + @classmethod def empty(cls): return cls.grad_queue.empty() - + @classmethod def pop(cls): item = cls.grad_queue.get() @@ -27,21 +28,21 @@ def pop(cls): for total_input, grad_output, mod, communicator, func in item: params.extend(func(total_input, grad_output, mod, communicator)) return params - + @classmethod def cache_full_wp_parameters(cls, shard_param, full_param): cls.cached_wp_params[shard_param] = full_param - + @classmethod def retrieve_full_wp_parameters(cls, shard_param): return cls.cached_wp_params.get(shard_param, None) @classmethod def clear_cached_wp_parameters(cls): - for shard_param, full_param in cls.cached_wp_params.items(): + for shard_param, full_param in cls.cached_wp_params.items(): del full_param cls.cached_wp_params[shard_param] = None - + @classmethod def check_postpond_grad_accum(cls, param): return param in cls.cached_wp_params From 42288a3d4899604e26941896c561bbf331aea80c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Mon, 24 Jun 2024 15:30:24 +0800 Subject: [PATCH 4/6] fix lint --- internlm/initialize/initialize_trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internlm/initialize/initialize_trainer.py b/internlm/initialize/initialize_trainer.py index 3402e56d9..f33beddf8 100644 --- a/internlm/initialize/initialize_trainer.py +++ b/internlm/initialize/initialize_trainer.py @@ -97,8 +97,8 @@ def initialize_trainer( communication_overlap = gpc.config.parallel["pipeline"].get("interleaved_overlap", False) use_zeropp = hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp if use_zeropp: - unit_schedule_size = gpc.config.data["unit_schedule_size"] if(hasattr(gpc.config, "data") and hasattr(gpc.config.data, "unit_schedule_size")) \ - else None + unit_schedule_size = gpc.config.data["unit_schedule_size"] \ + if(hasattr(gpc.config, "data") and hasattr(gpc.config.data, "unit_schedule_size")) else None scheduler = ZeroPPScheduler( data_process_func=data_fn, num_microbatches=gpc.config.NUM_MICRO_BATCHES, From 579ec8523e99c12fde1dff717f379244d85713a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Mon, 24 Jun 2024 19:57:44 +0800 Subject: [PATCH 5/6] fix lint --- internlm/core/scheduler/pipeline_scheduler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py index 326fce451..b50c68035 100644 --- a/internlm/core/scheduler/pipeline_scheduler.py +++ b/internlm/core/scheduler/pipeline_scheduler.py @@ -1822,6 +1822,7 @@ def _run_1f1b_loop_with_overlap( forward_async_communicator = None backward_async_communicator = None + next_forward_chunk_id = next_backward_chunk_id = 0 # Run 1F1B in steady state. for k in range(num_1f1b_micropairs): From 85cd6f3093f1aa638e85314f3dba1905d4a24d6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=94=90=E4=B8=81?= Date: Tue, 25 Jun 2024 14:52:45 +0800 Subject: [PATCH 6/6] fix lint --- internlm/core/scheduler/__init__.py | 6 +++++- internlm/core/scheduler/pipeline_scheduler.py | 18 +++++++++++------- internlm/initialize/initialize_trainer.py | 9 ++++++--- internlm/model/modules/linear.py | 16 +++++++--------- internlm/solver/optimizer/hybrid_zero_optim.py | 5 +++-- internlm/utils/zeropp_manager.py | 4 ++++ 6 files changed, 36 insertions(+), 22 deletions(-) diff --git a/internlm/core/scheduler/__init__.py b/internlm/core/scheduler/__init__.py index a86dee753..3c25dc54a 100644 --- a/internlm/core/scheduler/__init__.py +++ b/internlm/core/scheduler/__init__.py @@ -1,6 +1,10 @@ from .base_scheduler import BaseScheduler from .no_pipeline_scheduler import NonPipelineScheduler -from .pipeline_scheduler import InterleavedPipelineScheduler, PipelineScheduler, ZeroPPScheduler +from .pipeline_scheduler import ( + InterleavedPipelineScheduler, + PipelineScheduler, + ZeroPPScheduler, +) __all__ = [ "BaseScheduler", diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py index b50c68035..d5a7088dd 100644 --- a/internlm/core/scheduler/pipeline_scheduler.py +++ b/internlm/core/scheduler/pipeline_scheduler.py @@ -1507,8 +1507,10 @@ def __init__( self._get_chunk_by_microbatch = self._get_chunk_by_microbatch_1f1b self._cur_unit_schedule_size = self._unit_schedule_size - if hasattr(gpc.config.parallel.pipeline, "decouple_grad") \ - and gpc.config.parallel.pipeline.decouple_grad is True: + if ( + hasattr(gpc.config.parallel.pipeline, "decouple_grad") + and gpc.config.parallel.pipeline.decouple_grad is True + ): self.decouple_grad = True else: self.decouple_grad = False @@ -1582,8 +1584,9 @@ def _forward_step(self, engine, chunk_id): gpc.set_virtual_pipeline_parallel_rank(chunk_id) if gpc.is_pipeline_first_stage(): - if (len(self._input_objs[chunk_id]) + len(self._input_objs_for_backward[chunk_id])) \ - == len(self._output_objs[chunk_id]): + if (len(self._input_objs[chunk_id]) + len(self._input_objs_for_backward[chunk_id])) == len( + self._output_objs[chunk_id] + ): self._input_objs[chunk_id].append(None) input_obj = self._input_objs[chunk_id].pop(0) @@ -1888,8 +1891,9 @@ def _run_1f1b_loop_with_overlap( input_obj_grad = None if gpc.is_pipeline_last_stage(ignore_virtual=True): - next_backward_chunk_id = \ - self._get_chunk_by_microbatch(backward_microstep_id - (self._pp_size - 1), backward=True) + next_backward_chunk_id = self._get_chunk_by_microbatch( + backward_microstep_id - (self._pp_size - 1), backward=True + ) next_backward_chunk_id -= 1 else: next_backward_chunk_id = self._get_chunk_by_microbatch(backward_microstep_id + 1, backward=True) @@ -2057,7 +2061,7 @@ def _forward_backward_step_gpipe(self, engine: Engine): else: last_unit_schedule_size = self._unit_schedule_size - for i in range(num_unit_schedules - 1): + for _ in range(num_unit_schedules - 1): self._unit_schedule_step(engine, self._unit_schedule_size) self._unit_schedule_step(engine, last_unit_schedule_size) diff --git a/internlm/initialize/initialize_trainer.py b/internlm/initialize/initialize_trainer.py index f33beddf8..7ba088f08 100644 --- a/internlm/initialize/initialize_trainer.py +++ b/internlm/initialize/initialize_trainer.py @@ -97,8 +97,11 @@ def initialize_trainer( communication_overlap = gpc.config.parallel["pipeline"].get("interleaved_overlap", False) use_zeropp = hasattr(gpc.config.parallel.pipeline, "use_zeropp") and gpc.config.parallel.pipeline.use_zeropp if use_zeropp: - unit_schedule_size = gpc.config.data["unit_schedule_size"] \ - if(hasattr(gpc.config, "data") and hasattr(gpc.config.data, "unit_schedule_size")) else None + unit_schedule_size = ( + gpc.config.data["unit_schedule_size"] + if (hasattr(gpc.config, "data") and hasattr(gpc.config.data, "unit_schedule_size")) + else None + ) scheduler = ZeroPPScheduler( data_process_func=data_fn, num_microbatches=gpc.config.NUM_MICRO_BATCHES, @@ -108,7 +111,7 @@ def initialize_trainer( scatter_gather_tensors=scatter_gather, scheduler_hooks=scheduler_hooks, communication_overlap=communication_overlap, - unit_schedule_size=unit_schedule_size + unit_schedule_size=unit_schedule_size, ) else: scheduler = InterleavedPipelineScheduler( diff --git a/internlm/model/modules/linear.py b/internlm/model/modules/linear.py index f3e75cf76..68b7eb526 100644 --- a/internlm/model/modules/linear.py +++ b/internlm/model/modules/linear.py @@ -33,21 +33,19 @@ custom_fwd = internlm_accelerator.return_custom_fwd() -def decoupled_grad_func(input, grad_output, module, communicator): +def decoupled_grad_func(x, grad_output, module, communicator): grad_weight, grad_bias = linear_backward_op( - input, + x, grad_output, module.bias is not None and module.bias.requires_grad, ) - communicator.grad_hook( - grad_weight, async_op=True, module=module, is_bias=False - ) + communicator.grad_hook(grad_weight, async_op=True, module=module, is_bias=False) if grad_bias is not None: - communicator.grad_hook( - grad_bias, async_op=True, module=module, is_bias=True - ) + communicator.grad_hook(grad_bias, async_op=True, module=module, is_bias=True) return [module.weight, module.bias] - return [module.weight, ] + return [ + module.weight, + ] # adpated from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/fused_dense.py diff --git a/internlm/solver/optimizer/hybrid_zero_optim.py b/internlm/solver/optimizer/hybrid_zero_optim.py index f58df02b4..65860e631 100644 --- a/internlm/solver/optimizer/hybrid_zero_optim.py +++ b/internlm/solver/optimizer/hybrid_zero_optim.py @@ -417,8 +417,9 @@ def _accum_grads_store_in_bucket(self, bucket: BucketStore, reduce_rank: Optiona bucket.reset_by_rank(reduce_rank) - def _wait_reduce_scatter_and_accumulate_grads(self, param, reduce_rank: Optional[int] = None, - skip_decoupled_grad_accum=True): + def _wait_reduce_scatter_and_accumulate_grads( + self, param, reduce_rank: Optional[int] = None, skip_decoupled_grad_accum=True + ): param_size = param.numel() group_id = getattr(param, "group_id") diff --git a/internlm/utils/zeropp_manager.py b/internlm/utils/zeropp_manager.py index ec1c23bb4..61615328f 100644 --- a/internlm/utils/zeropp_manager.py +++ b/internlm/utils/zeropp_manager.py @@ -2,6 +2,10 @@ class ZeroppManager: + """ + ZeroppManager is used to manage ZeroPP parameters gathering and gradient calculation. + """ + cache = [] cached_wp_params = {} grad_queue = queue.LifoQueue()