# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and from __future__ import annotations import os import queue import sys import time import warnings from collections import defaultdict, deque from dataclasses import dataclass from enum import Enum from functools import partial import paddle from paddle import framework from ..meta_optimizers.dygraph_optimizer import HybridParallelOptimizer from ..utils import timer_helper as timer from ..utils.hybrid_parallel_util import ( broadcast_dp_parameters, broadcast_moe_sharding_parameters, broadcast_mp_parameters, broadcast_sep_parameters, broadcast_sharding_parameters, ) from ..utils.log_util import get_sync_logger, logger from .meta_parallel_base import MetaParallelBase from .parallel_layers.pp_layers import PipelineLayer _use_four_directions = os.environ.get( 'PADDLE_USE_FOUR_DIRECTIONS_P2P', paddle.base.core.is_compiled_with_xpu() ) _use_four_directions = False # xpu use the same p2p method as gpu if _use_four_directions: from .pp_utils import four_directions_p2p_communication as p2p else: from .pp_utils import p2p_communication as p2p from typing import TYPE_CHECKING from paddle.distributed import fleet from paddle.distributed.fleet.utils.tensor_fusion_helper import ( HOOK_ACTION, FusedCommBuffer, assign_group_by_size, ) from .pipeline_hooks import ( PipelineHook, ) from .pp_utils.utils import dict_to_tuple_helper, tuple_to_dict_helper from .zero_bubble_utils import WeightGradStore if TYPE_CHECKING: from collections.abc import Callable g_profile_pipeline_details_steps = int( os.getenv("FLAGS_profile_pipeline_details_steps", "0") ) __all__ = [] def profile_pipeline_details(msg): GB = 1024.0 * 1024.0 * 1024.0 if paddle.base.core.is_compiled_with_cuda(): memory_allocated_size = paddle.device.cuda.memory_allocated() / GB memory_reserved_size = paddle.device.cuda.memory_reserved() / GB else: memory_allocated_size, memory_reserved_size = 0, 0 get_sync_logger().info( f"{msg}: memory_allocated_size={memory_allocated_size:.2f}, memory_reserved_size={memory_reserved_size:.2f}" ) def get_action(is_dp, shard_split_param=False): if is_dp: return HOOK_ACTION.ALL_REDUCE if shard_split_param: return HOOK_ACTION.REDUCE_SCATTER return HOOK_ACTION.REDUCE def _get_align_mode_scale(): hcg = fleet.get_hybrid_communicate_group() data_parallel_world_size = hcg.get_data_parallel_world_size() sharding_parallel_world_size = hcg.get_sharding_parallel_world_size() return max(data_parallel_world_size, 1) * max( sharding_parallel_world_size, 1 ) def _can_free(t): """ Check if a tensor can be freed. A tensor can be freed only if all of the following conditions are met: 1. Tensor is not None 2. Is a paddle.Tensor type 3. Has been initialized 4. inplace_version is 0 (not using in-place ops) or explicitly marked as freeable Args: t: The tensor to check Returns: bool: True if the tensor can be freed, False otherwise """ return ( t is not None and isinstance(t, paddle.Tensor) and t._is_initialized() and (t.inplace_version == 0 or getattr(t, "pp_can_free", False)) ) def _collect_all_tensors(obj, tensor_set): """ Recursively collect all tensors from a complex object. This function traverses nested data structures (tuple, list, dict) and finds all paddle.Tensor instances, adding them to the tensor_set. Used in Pipeline Parallel to identify all tensors that need to be managed. Args: obj: Any complex object that may contain nested tuple, list, dict and paddle.Tensor tensor_set: A set to store the collected tensors """ visited = set() stack = [obj] while stack: current = stack.pop() obj_id = id(current) if obj_id in visited: continue visited.add(obj_id) if isinstance(current, (tuple, list)): stack.extend(current) elif isinstance(current, dict): stack.extend(current.values()) elif isinstance(current, paddle.Tensor): # Check for duplicate addition if current in tensor_set: logger.debug(f"Duplicate tensor detected: {current}") tensor_set.add(current) def _release_output(output): """ Release the data pointer of output tensors. Collects all tensors from output and frees the data pointer of those that meet the release criteria. Used in Pipeline Parallel to release output tensor memory after forward propagation to avoid unnecessary memory usage. Args: output: The output object, which can be a tensor, tuple, list, or dict """ all_tensors = set() _collect_all_tensors(output, all_tensors) for t in all_tensors: if _can_free(t): t._clear_dataptr() def _release_input(input, output): """ Release the data pointer of input tensors. Only releases input tensors that do not appear in the output. This is because in Pipeline Parallel, if an input tensor is used in the output (e.g., residual connection), it cannot be freed early. This function ensures that input memory is released without affecting tensors needed for subsequent computation. Args: input: The input object, which can be a tensor, tuple, list, or dict output: The output object, used to determine which input tensors should not be freed """ output_tensors = set() _collect_all_tensors(output, output_tensors) def can_release(t): if not _can_free(t): return False return t not in output_tensors input_tensors = set() _collect_all_tensors(input, input_tensors) for t in input_tensors: if can_release(t): t._clear_dataptr() # assume only the first stage and last stage need data, and data consumption is ordered # to be replaced by real micro dataset from reader class FakeMicroDataset: def __init__( self, data, is_first_stage, is_last_stage, acc_steps, micro_batch_size, ): self._data = data self._index = 0 self._acc_steps = acc_steps self._is_first_stage = is_first_stage self._is_last_stage = is_last_stage self._micro_batch_size = micro_batch_size def __iter__(self): return self def __next__(self): if self._index >= self._acc_steps: raise StopIteration assert self._is_first_stage or self._is_last_stage micro_batch_data = self._load_micro_batch(self._index) self._index += 1 if self._index >= self._acc_steps: self._data = None # clearup return micro_batch_data def _load_micro_batch(self, micro_step): inputs = self._data data = None label = None if self._is_first_stage: assert len(inputs) == 2, "length of input should be 2" data = self._load_micro_batch_impl(inputs[0], micro_step) if self._is_last_stage: assert len(inputs) == 2, "length of input should be 2" label = self._load_micro_batch_impl(inputs[1], micro_step) return (data, label) def _load_micro_batch_impl(self, inputs, micro_step): begin = micro_step * self._micro_batch_size end = begin + self._micro_batch_size if isinstance(inputs, tuple): output = [] for data in inputs: if isinstance(data, list): assert len(data) == self._acc_steps, ( f"length of data should be {self._acc_steps}, but it is {len(data)}" ) output.append( data[micro_step].detach() if data[micro_step] is not None else None ) elif data is not None: self._check_data_valid(data) output.append(data[begin:end, :].detach()) else: output.append(None) return tuple(output) elif isinstance(inputs, dict): output_dict = {} for key, data in inputs.items(): if isinstance(data, list): assert len(data) == self._acc_steps, ( f"length of data should be {self._acc_steps}, but it is {len(data)}" ) output_dict[key] = ( data[micro_step].detach() if data[micro_step] is not None else None ) elif data is not None: self._check_data_valid(data) output_dict[key] = data[begin:end, :].detach() else: output_dict[key] = None return output_dict elif isinstance(inputs, list): assert len(inputs) == self._acc_steps, ( f"length of data should be {self._acc_steps}, but it is {len(inputs)}" ) if isinstance(inputs[micro_step], list): return [ tensor.detach() if tensor is not None else None for tensor in inputs[micro_step] ] return inputs[micro_step].detach() elif inputs is not None: self._check_data_valid(inputs) return inputs[begin:end, :].detach() else: return None def _check_data_valid(self, data): batch_size = data.shape[0] assert self._micro_batch_size * self._acc_steps == batch_size, ( "batch_size needs to be divisible by micro_batch_size. Currently, " f"batch_size = {batch_size}, micro_batch_size = {self._micro_batch_size}, accumulate_steps = {self._acc_steps}." ) # A wrapper for pipeline dataser, to avoid GPU memory leaks. class PipelineDatasetPreprocessor: def __init__(self, function): self.function = function def __call__(self): return self.function() # Enum for specifying the pipeline parallel micro-step locations. class PipelineParallelMicroStepLocations(Enum): FORWARD_BEGIN = 'forward_begin' FORWARD_END = 'forward_end' BACKWARD_BEGIN = 'backward_begin' BACKWARD_END = 'backward_end' # A callback class for managing hooks at different stages of a pipeline parallel process. class PipelineParallelMicroStepCallback: def __init__(self): # Initializes a dictionary to store hooks for each micro-step location in the pipeline. self.hooks: dict[PipelineParallelMicroStepLocations, list[Callable]] = { PipelineParallelMicroStepLocations.FORWARD_BEGIN: [], PipelineParallelMicroStepLocations.FORWARD_END: [], PipelineParallelMicroStepLocations.BACKWARD_BEGIN: [], PipelineParallelMicroStepLocations.BACKWARD_END: [], } def register_hook( self, location: PipelineParallelMicroStepLocations, hook: Callable ): """ Registers a hook function to be called at a specified pipeline parallel micro-step location. Args: location (PipelineParallelMicroStepLocations): The micro-step location where the hook should be registered. hook (Callable): The hook function to be registered. The function should accept the following optional keyword arguments: - input_tensor (paddle.Tensor): The input tensor to the current micro-step. - output_tensor (paddle.Tensor): The output tensor from the current micro-step. - input_tensor_grad (paddle.Tensor): The gradient of the input tensor. - output_tensor_grad (paddle.Tensor): The gradient of the output tensor. - step_id (paddle.Tensor): An identifier for the current step in the pipeline. Raises: AssertionError: If the specified location is not a valid micro-step location. """ assert location in self.hooks, ( f"Invalid location '{location}'. Valid locations are 'forward_begin', 'forward_end', 'backward_begin', or 'backward_end'." ) self.hooks[location].append(hook) def on_location( self, location: PipelineParallelMicroStepLocations, **kwargs ): """ Triggers all registered hooks at a specified pipeline parallel micro-step location. Args: location (PipelineParallelMicroStepLocations): The micro-step location where the hooks should be triggered. kwargs: Additional keyword arguments to be passed to the hook functions. Raises: AssertionError: If the specified location is not a valid micro-step location. """ assert location in self.hooks, ( f"Invalid location '{location}'. Valid locations are 'forward_begin', 'forward_end', 'backward_begin', or 'backward_end'." ) for hook in self.hooks[location]: hook(**kwargs) pipeline_parallel_callbacks_ = PipelineParallelMicroStepCallback() # It is typically very difficult for us to directly access the PipelineParallel object. # Users may use fleet.distributed_model to wrap a model into a pipeline parallel model (PP model). # We may not have access to the wrapped model when we want to register hooks, for example, when using PaddleNLP trainer to wrap around the PP model. # Additionally, we usually have only one `PipelineParallel` model, so the callbacks are registered globally. def register_global_pipeline_parallel_hook( location: PipelineParallelMicroStepLocations, hook: Callable ): """ Registering global hooks for pipeline parallelism. """ pipeline_parallel_callbacks_.register_hook(location, hook) class NoPipelineParallel(MetaParallelBase): def __init__(self, layers, strategy, hcg=None): assert isinstance(layers, PipelineLayer) super().__init__(layers, hcg, strategy) self._layers = layers self._strategy = strategy self._hcg = hcg self.micro_batch_size = self._strategy.pipeline_configs[ "micro_batch_size" ] self.accumulate_steps = self._strategy.pipeline_configs[ "accumulate_steps" ] self._dp_comm_overlap = False self._sharding_comm_overlap = False # store total loss of entire batch. It contains the loss of each micro batch in a list, then contains many loss_fn's list in total_loss. self.total_loss = None # default loss function index self.loss_fn_idx = 0 if self._hcg is not None: self.use_data_parallel = ( self._hcg.get_data_parallel_world_size() > 1 ) self.use_model_parallel = ( self._hcg.get_model_parallel_world_size() > 1 ) self.use_sep_parallel = self._hcg.get_sep_parallel_world_size() > 1 self.use_sharding_parallel = ( self._hcg.get_sharding_parallel_world_size() > 1 ) self.use_moe_sharding_parallel = ( self._hcg.get_moe_sharding_parallel_world_size() > 1 ) self.dp_group = self._hcg.get_data_parallel_group() # fused sep and dp if self.use_sep_parallel: self.dp_group = self._hcg.get_dp_sep_parallel_group() if self.use_model_parallel: logger.info("start broadcast mp parameters") broadcast_mp_parameters(self._layers, self._hcg) if self.use_sep_parallel: logger.info("start broadcast sep parameters") broadcast_sep_parameters(self._layers, self._hcg) if self.use_sharding_parallel: logger.info("start broadcast sharding parameters") broadcast_sharding_parameters(self._layers, self._hcg) if self.use_data_parallel: logger.info("start broadcast dp parameters") broadcast_dp_parameters(self._layers, self._hcg) if self.use_moe_sharding_parallel: logger.info("start broadcast moe_sharding parameters") broadcast_moe_sharding_parameters(self._layers, self._hcg) def is_pipeline_last_stage(self, ignore_virtual=False): return True def _check_micro_batch_data_valid(self, micro_batch_data): if isinstance(micro_batch_data, (tuple, list)): for data in micro_batch_data: self._check_micro_batch_data_valid(data) elif isinstance(micro_batch_data, dict): for value in micro_batch_data.values(): self._check_micro_batch_data_valid(value) elif micro_batch_data is not None: assert isinstance(micro_batch_data, paddle.Tensor) def _prepare_training(self, data, optimizer, lr_scheduler): assert framework._dygraph_tracer()._has_grad, ( "Please enable the generation of gradients." ) self.optimizer = optimizer self.lr_scheduler = lr_scheduler self._layers.train() return data def _optimizer_step(self): for p in self._layers.parameters(): if hasattr(p, "main_grad") and p.main_grad is not None: assert p.grad is None p.main_grad = p.main_grad.scale(1.0 / self.accumulate_steps) elif p.grad is not None: p.grad = p.grad.scale(1.0 / self.accumulate_steps) if self.scaler: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() self.optimizer.clear_grad() if self.lr_scheduler: self.lr_scheduler.step() def forward_backward_pipeline( self, data, scaler=None, return_micro_batch_loss=False, ): self.scaler = scaler self.total_loss = None if isinstance(data, PipelineDatasetPreprocessor): data = data() if (not isinstance(data, tuple)) and (not isinstance(data, list)): micro_dataset = data else: micro_dataset = FakeMicroDataset( data, True, True, self.accumulate_steps, self.micro_batch_size, ) loss_list = [] for _ in range(self.accumulate_steps): # data prepare data_iter = next(micro_dataset) input_tensor = data_iter[0] label = data_iter[1] self._check_micro_batch_data_valid(input_tensor) self._check_micro_batch_data_valid(label) # forward output_tensor = self._layers.forward(input_tensor) # loss is loss_fn[loss_fn_idx]'s result loss = None # cal loss for idx, loss_fn in enumerate(self._layers._loss_fn): loss_tensor = loss_fn(output_tensor, label) assert isinstance(loss_tensor, paddle.Tensor), ( "Currently, loss_fn should obtain Paddle.Tensor dtype" ) if self.total_loss is None: self.total_loss = [] # when self.total_loss length is less than idx, append a new tensor if len(self.total_loss) <= idx: self.total_loss.append([]) self.total_loss[idx].append(loss_tensor.detach()) if idx == self.loss_fn_idx: loss = loss_tensor # backward with paddle.amp.auto_cast(enable=False): if self.scaler: paddle.autograd.backward(self.scaler.scale(loss)) else: paddle.autograd.backward(loss) assert self.total_loss is not None, ( "train_batch() in last stage should obtain valid loss" ) losses = [] with paddle.amp.auto_cast(enable=False): for idx in range(len(self._layers._loss_fn)): self.total_loss[idx] = paddle.to_tensor(self.total_loss[idx]) if not return_micro_batch_loss: # TODO(shenliang03): it will use mean/sum to calculate loss tmp = paddle.zeros_like(self.total_loss[idx][0]) for loss in self.total_loss[idx]: tmp += loss.detach() losses.append(tmp / self.accumulate_steps) else: losses.append(self.total_loss[idx].detach()) return losses[0] if len(losses) == 1 else losses def train_batch( self, data, optimizer, lr_scheduler=None, scaler=None, loss_fn_idx=0, return_micro_batch_loss=False, ): data = self._prepare_training(data, optimizer, lr_scheduler) # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx # no pipeline parallel train_loss = self.forward_backward_pipeline( data, scaler, return_micro_batch_loss=return_micro_batch_loss ) # optimizer with paddle.amp.auto_cast(enable=False): self._optimizer_step() return train_loss def eval_batch( self, data, compute_loss=False, loss_fn_idx=0, return_host_tensor=False ): # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx self.total_loss = None if isinstance(data, PipelineDatasetPreprocessor): data = data() if (not isinstance(data, tuple)) and (not isinstance(data, list)): micro_dataset = data else: micro_dataset = FakeMicroDataset( data, True, True, self.accumulate_steps, self.micro_batch_size, ) loss_list = [] output_list = [] for _ in range(self.accumulate_steps): # data prepare data_iter = next(micro_dataset) input_tensor = data_iter[0] label = data_iter[1] self._check_micro_batch_data_valid(input_tensor) self._check_micro_batch_data_valid(label) # forward output_tensor = self._layers.forward(input_tensor) if compute_loss: # loss is loss_fn[loss_fn_idx]'s result loss = None # cal loss for idx, loss_fn in enumerate(self._layers._loss_fn): loss_tensor = loss_fn(output_tensor, label) assert isinstance(loss_tensor, paddle.Tensor), ( "Currently, loss_fn should obtain Paddle.Tensor dtype" ) if self.total_loss is None: self.total_loss = [] # when self.total_loss length is less than idx, append a new tensor if len(self.total_loss) <= idx: self.total_loss.append([]) self.total_loss[idx].append(loss_tensor.detach()) if idx == self.loss_fn_idx: loss = loss_tensor assert self.total_loss is not None, ( "train_batch() in last stage should obtain valid loss" ) else: if return_host_tensor: self._offload_tensors(output_tensor) output_list.append(output_tensor) if compute_loss: losses = [] return_micro_batch_loss = False for idx in range(len(self._layers._loss_fn)): self.total_loss[idx] = paddle.to_tensor(self.total_loss[idx]) # if not return_micro_batch_loss: # TODO(shenliang03): it will use mean/sum to calculate loss tmp = paddle.zeros_like(self.total_loss[idx][0]) for loss in self.total_loss[idx]: tmp += loss.detach() losses.append(tmp / self.accumulate_steps) # else: # losses.append(self.total_loss[idx].detach()) res = losses[0] if len(losses) == 1 else losses else: res = output_list return res def _offload_tensors(self, output_tensor): if isinstance(output_tensor, (tuple, list)): for t in output_tensor: if not isinstance(t, paddle.Tensor): continue host_tensor = ( t.pin_memory() if hasattr(t, "pin_memory") else t.cpu() ) host_tensor._share_buffer_to(t) else: if not isinstance(output_tensor, paddle.Tensor): return host_tensor = ( output_tensor.pin_memory() if hasattr(output_tensor, "pin_memory") else output_tensor.cpu() ) host_tensor._share_buffer_to(output_tensor) class PipelineParallel(MetaParallelBase): def __init__(self, layers, hcg, strategy): if not isinstance(layers, PipelineLayer): raise TypeError( "The Layer should be a derived class of PipelineLayer." ) super().__init__(layers, hcg, strategy) self.use_data_parallel = self._hcg.get_data_parallel_world_size() > 1 self.use_model_parallel = self._hcg.get_model_parallel_world_size() > 1 self.use_sep_parallel = self._hcg.get_sep_parallel_world_size() > 1 self.use_sharding_parallel = ( self._hcg.get_sharding_parallel_world_size() > 1 ) self.use_moe_sharding_parallel = ( self._hcg.get_moe_sharding_parallel_world_size() > 1 ) self.use_dict_in_pp = True self.total_loss = None self.micro_batch_size = self._strategy.pipeline_configs[ 'micro_batch_size' ] self.accumulate_steps = self._strategy.pipeline_configs[ 'accumulate_steps' ] # If sent tensor are not the same from different hosts, # they shouldn't been sent partially and then concatenated as a whole tensor. self._enable_partial_send_recv = self._strategy.pipeline_configs[ 'enable_partial_send_recv' ] self._using_cache = self._strategy.pipeline_configs['p2p_cache_shape'] self.num_stages = self._hcg.get_pipe_parallel_world_size() self.stage_id = self._hcg.get_stage_id() self.global_rank = self._hcg.get_global_rank() self.pp_group = self._hcg.get_pipe_parallel_group() self.dp_group = self._hcg.get_data_parallel_group() # fused sep and dp if self.use_sep_parallel: self.dp_group = self._hcg.get_dp_sep_parallel_group() self.sharding_group = self._hcg.get_sharding_parallel_group() self._virtual_pp_world_size = None self._virtual_pp_rank = None self._real_pp_world_size = self.num_stages self._real_pp_rank = self.stage_id # TODO(PP Dev): support dp_comm_overlap without use_main_grad training. # This combination will trigger inplace check error during `reshape_` in function `_split_tensors`. self._dp_comm_overlap = self._strategy.hybrid_configs[ "pp_configs" ].dp_comm_overlap self._sharding_comm_overlap = self._strategy.hybrid_configs[ "pp_configs" ].sharding_comm_overlap self._enable_timer = self._strategy.hybrid_configs[ "pp_configs" ].enable_timer self._release_gradients = self._strategy.hybrid_configs[ "pp_configs" ].release_gradients self._sharding_split_param = self._strategy.hybrid_configs[ "sharding_configs" ].split_param self._overlap_p2p_comm = self._strategy.hybrid_configs[ "pp_configs" ].overlap_p2p_comm self._clear_every_step_cache = self._strategy.hybrid_configs[ "pp_configs" ].clear_every_step_cache self._use_batch_p2p_comm = self._strategy.hybrid_configs[ "pp_configs" ].use_batch_p2p_comm self._dynamic_shape = self._strategy.hybrid_configs[ 'pp_configs' ].enable_dynamic_shape logger.info( f"Pipeline scheduler is in dynamic_shape mode={self._dynamic_shape}" ) if self._use_batch_p2p_comm and self._overlap_p2p_comm: warnings.warn( "non_batch_p2p_comm should be enabled when overlap_p2p_comm is activated, setting non_batch_p2p_comm=True." ) self._use_batch_p2p_comm = False logger.info( f"dp_comm_overlap {self._dp_comm_overlap}; \ sharding_comm_overlap {self._sharding_comm_overlap}; \ sharding_split_param {self._sharding_split_param};" ) self._profiling = self._strategy.hybrid_configs["pp_configs"].profiling self._records = [] self._record_format = ( '"name": "{}{}", "cat": "pipeline timeline", "ph": {}, "pid": 0, "tid": ' + str(self.stage_id + 1) + ', "ts": {}, "cname": "{}"' ) self._forward_color = "thread_state_running" # RGB: 126, 200, 148 self._backward_color = "rail_idle" # RGB: 238, 142, 0 if self._profiling: logger.info( "If enable pp profiling, the max training steps should be restricted " "to a reasonable value (such as 5) to avoid generating large profile files. " "The profiler will generate a profile file 'profile_record_tmp_file_for_rank_*' " "for each rank. Users should gather all profile files for one entire pipeline " "to one node (rank 0 is recommended) to get the full view of the pipeline profile. " "[DONT CHANGE THE NAME OF THE PROFILE FILES!]. " "Then get the profile parser from this url: " "https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/distributed/fleet/meta_parallel/pp_utils/profiler_helper.py " "and save the script to the same directory of all profile files." "Parse those files by this command: `python profiler_helper.py`. " "After parsing, a new file 'pipeline_profile.json' will be generated. " "Users can inspect this file by chrome://tracing website." ) if self._dp_comm_overlap: assert self.use_data_parallel and self.num_stages > 1 if self._sharding_comm_overlap: assert self.use_sharding_parallel and self.num_stages > 1 assert not (self._dp_comm_overlap and self._sharding_comm_overlap), ( "Cannot use dp pp overlap and sharding pp overlap at the same time." ) self._chunk_2_comm_buffers = defaultdict(list) self._comm_overlap = ( self._dp_comm_overlap or self._sharding_comm_overlap ) if self._enable_timer: if not timer.is_timer_initialized(): timer.set_timers() self.timers = timer.get_timers() p2p.initialize_p2p_groups( hcg, self._enable_partial_send_recv, self._enable_timer, ) # construct pipeline meta info self._p2p_helper = p2p.P2pHelper( self._using_cache, dynamic_shape=self._dynamic_shape ) self.global_rank = self._hcg.get_global_rank() self.micro_batch_id = 0 # default loss function index self.loss_fn_idx = 0 self._compute_loss = True self._return_host_tensor = False self.callbacks = pipeline_parallel_callbacks_ logger.info( f"Pipeline Info -- num_stages: {self.num_stages}, stage_id: {self.stage_id}" ) if self.use_model_parallel: logger.info("start broadcast mp parameters") broadcast_mp_parameters(self._layers, self._hcg) if self.use_sep_parallel: logger.info("start broadcast sep parameters") broadcast_sep_parameters(self._layers, self._hcg) if self.use_sharding_parallel: logger.info("start broadcast sharding parameters") broadcast_sharding_parameters(self._layers, self._hcg) if self.use_data_parallel: logger.info("start broadcast dp parameters") broadcast_dp_parameters(self._layers, self._hcg) if self.use_moe_sharding_parallel: logger.info("start broadcast moe_sharding parameters") broadcast_moe_sharding_parameters(self._layers, self._hcg) if self._dp_comm_overlap: self.register_allreduce_overlap_hook( self._layers, self.dp_group, self.accumulate_steps, True ) self.processed_steps = 0 self._init_user_hooks() # only support user hooks during training self.user_hooks_enabled = True def register_hook( self, location: PipelineParallelMicroStepLocations, hook: Callable ): self.callbacks.register_hook(location, hook) def _init_user_hooks(self): self._init_user_forward_backward_hooks() self._init_user_bubble_hooks() def _init_user_forward_backward_hooks(self): # initialize forward hooks self.forward_hooks = PipelineHook() self.forward_hooks.set_hooks_capacity( ( self._virtual_pp_world_size if self._virtual_pp_world_size is not None else 1 ) * self.accumulate_steps ) # initialize backward hooks self.backward_hooks = PipelineHook() self.backward_hooks.set_hooks_capacity( ( self._virtual_pp_world_size if self._virtual_pp_world_size is not None else 1 ) * self.accumulate_steps ) def _init_user_bubble_hooks(self): # (TODO:gexiao) support bubble hooks if needed self.bubble_hooks = None # self.bubble_hooks = PipelineHook() # self.bubble_hooks.set_hooks_capacity(2 * self.num_stages - 2) def _reset_user_hooks_status(self): if self.bubble_hooks: self.bubble_hooks.reset_current_id() if self.forward_hooks: self.forward_hooks.reset_current_id() if self.backward_hooks: self.backward_hooks.reset_current_id() def _check_user_hooks_status_at_step_end(self): if not self.user_hooks_enabled: return expected_bubble_step = 2 * self.num_stages - 2 expected_forward_step = ( self._virtual_pp_world_size if self._virtual_pp_world_size is not None else 1 ) * self.accumulate_steps expected_backward_step = ( self._virtual_pp_world_size if self._virtual_pp_world_size is not None else 1 ) * self.accumulate_steps if self.bubble_hooks: assert (self.bubble_hooks.current_id) == expected_bubble_step, ( f"bubble hooks status is not correct, current id is {self.bubble_hooks.current_id}, expected id is {expected_bubble_step}" ) if self.forward_hooks: assert (self.forward_hooks.current_id) == expected_forward_step, ( f"forward hooks status is not correct, current id is {self.forward_hooks.current_id}, expected id is {expected_forward_step}" ) if self.backward_hooks: assert (self.backward_hooks.current_id) == expected_backward_step, ( f"backward hooks status is not correct, current id is {self.backward_hooks.current_id}, expected id is {expected_backward_step}" ) def register_bubble_pipeline_parallel_hook( self, location: int, hook: Callable ): """ Registering bubble hooks for pipeline parallelism. """ if not self.bubble_hooks: raise ValueError("Bubble hooks are not supported yet.") self.bubble_hooks.register_hook(location, hook) def register_forward_pipeline_parallel_hook( self, location: int, hook: Callable ): """ Registering forward hooks for pipeline parallelism. """ if not self.forward_hooks: raise ValueError("Forward hooks are not supported yet.") self.forward_hooks.register_hook(location, hook) def register_backward_pipeline_parallel_hook( self, location: int, hook: Callable ): """ Registering backward hooks for pipeline parallelism. """ if not self.backward_hooks: raise ValueError("Backward hooks are not supported yet.") self.backward_hooks.register_hook(location, hook) @property def bubble_pipeline_parallel_hook_capacity(self): capacity = 0 if self.bubble_hooks: capacity = self.bubble_hooks.hooks_capacity return capacity @property def forward_pipeline_parallel_hook_capacity(self): capacity = 0 if self.forward_hooks: capacity = self.forward_hooks.hooks_capacity return capacity @property def backward_pipeline_parallel_hook_capacity(self): capacity = 0 if self.backward_hooks: capacity = self.backward_hooks.hooks_capacity return capacity def is_pipeline_first_stage(self, ignore_virtual=False): if not ignore_virtual: if self._virtual_pp_world_size is not None: assert self._virtual_pp_rank is not None if self._virtual_pp_rank != 0: return False assert self._real_pp_rank is not None return self._real_pp_rank == 0 def is_pipeline_last_stage(self, ignore_virtual=False): if not ignore_virtual: if self._virtual_pp_world_size is not None: assert self._virtual_pp_rank is not None if self._virtual_pp_rank != (self._virtual_pp_world_size - 1): return False assert self._real_pp_rank is not None assert self._real_pp_world_size is not None return self._real_pp_rank == (self._real_pp_world_size - 1) def set_virtual_pipeline_rank(self, rank): self._virtual_pp_rank = rank def fused_gradient( self, model, comm_group, acc_steps, dp, group_size=128 * 1024 * 1024 ): if model.get_num_virtual_stages() > 1: models = model.get_model_chunks() else: models = [model] act = get_action(dp, self._sharding_split_param) if act == HOOK_ACTION.REDUCE: assert hasattr(self, "optimizer") assert hasattr(self.optimizer, "_param2rank") _param2rank = self.optimizer._param2rank for chunk_idx, model in enumerate(models): # For virtual pipeline. Will separate parameters in different chunk into # different groups to get the best performance. fused_parameter_group = {} parameter_list = [ p for p in model.parameters() if not p.stop_gradient ] if len(parameter_list) < 1: return if act == HOOK_ACTION.REDUCE: # Sort parameters for sharding, since they have different dst rank for p in parameter_list: assert p.name in _param2rank dst_rank = _param2rank[p.name] if dst_rank in fused_parameter_group: fused_parameter_group[dst_rank].append(p) else: fused_parameter_group[dst_rank] = [p] else: fused_parameter_group[-1] = parameter_list for dst in fused_parameter_group: parameter_list = fused_parameter_group[dst] if act == HOOK_ACTION.REDUCE: # parse the relative dst rank to absolute dst rank for sharding dst = comm_group.ranks[dst] var_groups = assign_group_by_size(parameter_list, group_size) for group_idx, parameters in var_groups.items(): buffer = FusedCommBuffer( group_idx, parameters, comm_group, acc_steps, act, dst, release_grads=self._release_gradients, ) self._chunk_2_comm_buffers[chunk_idx].append(buffer) return self._chunk_2_comm_buffers def bw_hook_func(self, buffer, param): @paddle.autograd.no_grad() def fused_allreduce(*_): buffer.add_grad(param) return fused_allreduce def register_allreduce_overlap_hook( self, model, comm_group, acc_steps, dp, group_size=128 * 1024 * 1024 ): # register hook self.fused_gradient(model, comm_group, acc_steps, dp, group_size) for _, buffers in self._chunk_2_comm_buffers.items(): for buffer in buffers: for param in buffer._params: param._register_backward_hook( self.bw_hook_func(buffer, param) ) def timer_printer(self): if not self._enable_timer: return all_flag_names = self.timers.timers.keys() self.timers.log(all_flag_names) def _record_stamp(self, name, step, phase, color): if self._profiling: paddle.device.synchronize() self._records.append( '{' + self._record_format.format( name, step, phase, int(time.time() * 1000), color, ) + '}' ) def _flush_records(self): if self._profiling: with open( f'./profile_record_tmp_file_for_rank_{self.global_rank}', 'a+', ) as f: f.writelines(record + '\n' for record in self._records) self._records = [] def forward_backward_pipeline( self, data, scaler=None, static_scheduler=False, return_micro_batch_loss=False, ): # use the 1f1b scheduling strategy. # this strategy is inspired by: # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/schedules.py self._reset_user_hooks_status() # no _forward_only mode self.user_hooks_enabled = True if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] Start_forward_backward_pipeline" ) if static_scheduler: assert not self._profiling, ( "While _profiling, static scheduler is not available" ) if data is not None: warnings.warn( "Static scheduler run won't real run the model, but data has been provided" ) logger.info( "enable static_scheduler will return the pp schedule instead of the loss" ) schedule = "" self.scaler = scaler # store total loss of entire batch self.total_loss = None # store data id for micro_batch self.micro_batch_id = 0 startup_steps = self.num_stages - self.stage_id - 1 startup_steps = min(startup_steps, self.accumulate_steps) steady_steps = self.accumulate_steps - startup_steps input_buffers = [] output_buffers = [] micro_dataset = self._wrap_data(data) for step_id in range(startup_steps): if static_scheduler: schedule += f"f{step_id};" logger.info(f"forward step for micro step {step_id}") continue input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor) self._record_stamp("F", step_id, '"B"', self._forward_color) output_tensor, _, _ = self._forward_step( input_tensor=input_tensor_dict if use_dict else input_tensor, micro_dataset=micro_dataset, step_id=step_id, ) # convert dict to tuple whose tensor element has a key attribution output_tensor_tuple = dict_to_tuple_helper(output_tensor) self._record_stamp("F", step_id, '"E"', self._forward_color) # fwd output dict -> send tuple self._p2p_helper.send_forward( output_tensor=output_tensor_tuple, pp_last_stage=self.is_pipeline_last_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) input_buffers.append(input_tensor) output_buffers.append(output_tensor_tuple) if not self.is_pipeline_last_stage(): _release_output(output_tensor_tuple) if steady_steps > 0 and not static_scheduler: input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) for i in range(steady_steps): if static_scheduler: schedule += f"f{startup_steps + i};" schedule += f"b{i};" logger.info(f"forward step for micro step {startup_steps + i}") logger.info(f"backward step for micro step {i}") continue last_iter = i == (steady_steps - 1) input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor) self._record_stamp( "F", startup_steps + i, '"B"', self._forward_color ) output_tensor, _, _ = self._forward_step( input_tensor=input_tensor_dict if use_dict else input_tensor, micro_dataset=micro_dataset, step_id=startup_steps + i, ) self._record_stamp( "F", startup_steps + i, '"E"', self._forward_color ) output_tensor_tuple = dict_to_tuple_helper(output_tensor) # NOTE: `send_forward_recv_backward` is intentionally unused to # prevent hanging bugs in dynamic shape mode. self._p2p_helper.send_forward( output_tensor_tuple, self.is_pipeline_last_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) output_tensor_grad = self._p2p_helper.recv_backward( self.is_pipeline_last_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) input_buffers.append(input_tensor) output_buffers.append(output_tensor_tuple) if not self.is_pipeline_last_stage(): _release_output(output_tensor_tuple) input_tensor, output_tensor = ( input_buffers.pop(0), output_buffers.pop(0), ) self._record_stamp("B", i, '"B"', self._backward_color) input_tensor_grad = self._backward_step( input_tensor, output_tensor, output_tensor_grad, step_id=i ) self._record_stamp("B", i, '"E"', self._backward_color) if last_iter: input_tensor = None self._p2p_helper.send_backward( input_tensor_grad, self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) else: # NOTE: `send_backward_recv_forward` is intentionally unused to # prevent hanging bugs in dynamic shape mode. input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) self._p2p_helper.send_backward( input_tensor_grad, self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) for i in range(startup_steps): if static_scheduler: schedule += f"b{steady_steps + i};" logger.info(f"backward step for micro step {steady_steps + i}") continue input_tensor = input_buffers.pop(0) output_tensor = output_buffers.pop(0) output_tensor_grad = self._p2p_helper.recv_backward( self.is_pipeline_last_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) self._record_stamp( "B", steady_steps + i, '"B"', self._backward_color ) input_tensor_grad = self._backward_step( input_tensor, output_tensor, output_tensor_grad, step_id=steady_steps + i, ) self._record_stamp( "B", steady_steps + i, '"E"', self._backward_color ) self._p2p_helper.send_backward( input_tensor_grad, self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) if static_scheduler: return schedule self._flush_records() if self._comm_overlap: assert len(self._chunk_2_comm_buffers) > 0, ( "comm buffers should be created" ) for _, buffers in self._chunk_2_comm_buffers.items(): for buffer in buffers: buffer.scale_grads() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").start() self._layers.allreduce_shared_weight_gradients() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").stop() self.timers("broadcast_final_loss").start() with paddle.amp.auto_cast(enable=False): train_loss = self._broadcast_final_loss(return_micro_batch_loss) if self._enable_timer: self.timers("broadcast_final_loss").stop() if self._clear_every_step_cache: self._p2p_helper.clear_meta_cache() self.timer_printer() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] End_forward_backward_pipeline" ) self.processed_steps += 1 self._check_user_hooks_status_at_step_end() return train_loss def register_sharding_comm_overlap_hook(self, optimizer): """for delayed hook register until we get optimizer""" assert isinstance(optimizer, HybridParallelOptimizer), ( 'optimizer should be HybridParallelOptimizer subclass.' ) self.optimizer = optimizer if self._sharding_comm_overlap and len(self._chunk_2_comm_buffers) == 0: self.register_allreduce_overlap_hook( self._layers, self.sharding_group, self.accumulate_steps, False ) def _prepare_training(self, data, optimizer, lr_scheduler): # reset the virtual pp rank for each run self.set_virtual_pipeline_rank(0) assert isinstance(optimizer, HybridParallelOptimizer), ( 'optimizer should be HybridParallelOptimizer subclass.' ) assert framework._dygraph_tracer()._has_grad, ( 'Please enable the generation of gradients.' ) if self.is_pipeline_first_stage( ignore_virtual=True ) or self.is_pipeline_last_stage(ignore_virtual=True): assert data is not None, ( "For the first and the last stage, the data must be set." ) else: data = None self.optimizer = optimizer self.lr_scheduler = lr_scheduler self._layers.train() self.register_sharding_comm_overlap_hook(optimizer) return data def _wrap_data(self, data): """ for backward compatibility, wrap data to Fake FakeMicroDataset if it is of type list or tuple """ if isinstance(data, PipelineDatasetPreprocessor): data = data() if (not isinstance(data, tuple)) and (not isinstance(data, list)): return data micro_dataset = FakeMicroDataset( data, self.is_pipeline_first_stage(ignore_virtual=True), self.is_pipeline_last_stage(ignore_virtual=True), self.accumulate_steps, self.micro_batch_size, ) return micro_dataset def train_batch( self, data, optimizer, lr_scheduler=None, scaler=None, loss_fn_idx=0, return_micro_batch_loss=False, ): data = self._prepare_training(data, optimizer, lr_scheduler) # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx # 1f1b scheduler for pipeline parallel train_loss = self.forward_backward_pipeline( data, scaler, return_micro_batch_loss=return_micro_batch_loss ) # optimizer with paddle.amp.auto_cast(enable=False): self._optimizer_step() return train_loss def eval_batch( self, data, compute_loss=False, loss_fn_idx=0, return_host_tensor=False ): self.user_hooks_enabled = False # reset the virtual pp rank for each run self.set_virtual_pipeline_rank(0) self._layers.eval() origin_compute_loss = self._compute_loss self._compute_loss = compute_loss origin_return_host_tensor = self._return_host_tensor self._return_host_tensor = return_host_tensor # store data id for micro_batch self.micro_batch_id = 0 # store total loss of entire batch self.total_loss = None # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx startup_steps = self.num_stages - self.stage_id - 1 startup_steps = min(startup_steps, self.accumulate_steps) steady_steps = self.accumulate_steps - startup_steps output_buffers = [] # convert to micro dataset micro_dataset = self._wrap_data(data) for step_id in range(startup_steps): input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) # p2p data type: tuple # model input/return type: dict # here, convert p2p tuple -> dict input input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor) output_tensor, _, _ = self._forward_step( input_tensor_dict if use_dict else input_tensor, micro_dataset, step_id=None, ) # convert dict to tuple whose tensor element has a key attribution output_tensor_tuple = dict_to_tuple_helper(output_tensor) self._p2p_helper.send_forward( output_tensor_tuple, self.is_pipeline_last_stage(), skip_check_meta=True, batch_p2p_comm=self._use_batch_p2p_comm, ) if not self.is_pipeline_last_stage(): _release_output(output_tensor_tuple) else: self._offload_tensors(output_tensor_tuple) output_buffers.append(output_tensor_tuple) if steady_steps > 0: input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) for i in range(steady_steps): last_iter = i == (steady_steps - 1) input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor) output_tensor, _, _ = self._forward_step( input_tensor_dict if use_dict else input_tensor, micro_dataset, step_id=None, ) output_tensor_tuple = dict_to_tuple_helper(output_tensor) self._p2p_helper.send_forward( output_tensor_tuple, self.is_pipeline_last_stage(), skip_check_meta=True, batch_p2p_comm=self._use_batch_p2p_comm, ) if not self.is_pipeline_last_stage(): _release_output(output_tensor_tuple) else: self._offload_tensors(output_tensor_tuple) output_buffers.append(output_tensor_tuple) if not last_iter: input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) if self._compute_loss: train_loss = self._broadcast_final_loss() else: train_loss = output_buffers self._compute_loss = origin_compute_loss self._return_host_tensor = origin_return_host_tensor return train_loss def _maybe_loss_compute( self, output_tensor, micro_dataset, overlap_schedule_mode=False ): backward_loss_tensor = None backward_loss_fn_node = None loss_fn_node = None if self.is_pipeline_last_stage(): # train calculate loss for train if self._compute_loss: assert self._layers._loss_fn[self.loss_fn_idx] is not None, ( "loss function should exist to compute loss" ) labels = next(micro_dataset)[1] self._check_micro_batch_data_valid(labels) for idx, loss_fn in enumerate(self._layers._loss_fn): if overlap_schedule_mode: loss_fn_node = loss_fn.build_schedule_node() loss_fn_node.labels = labels loss_tensor = loss_fn_node.forward(output_tensor) else: loss_tensor = loss_fn(output_tensor, labels) assert isinstance(loss_tensor, paddle.Tensor), ( "Currently, loss_fn should obtain Paddle.Tensor dtype" ) if self.total_loss is None: self.total_loss = [] # when self.total_loss length is less than idx, append a new tensor if len(self.total_loss) <= idx: self.total_loss.append([]) self.total_loss[idx].append(loss_tensor.detach()) if idx == self.loss_fn_idx: backward_loss_tensor = loss_tensor backward_loss_fn_node = loss_fn_node return backward_loss_tensor, backward_loss_fn_node def _forward_step( self, input_tensor, micro_dataset, chunk_id=None, step_id=None, overlap_schedule_mode=False, ): if self.user_hooks_enabled: self.forward_hooks.run_hook() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( f"[Pipeline details] Before_forward_step_chunk_{chunk_id}_step_{step_id}" ) if self._enable_timer: self.timers("forward_step").start() if self.is_pipeline_first_stage(): input_tensor = next(micro_dataset)[0] self._check_micro_batch_data_valid(input_tensor) assert chunk_id is None or isinstance(chunk_id, int) self.callbacks.on_location( PipelineParallelMicroStepLocations.FORWARD_BEGIN, input_tensor=input_tensor, step_id=step_id, ) schedule_chunk = None if overlap_schedule_mode: schedule_chunk = self._layers.get_schedule_chunk(chunk_id=chunk_id) output_tensor = schedule_chunk.forward(input_tensor) else: output_tensor = self._layers.forward( input_tensor, chunk_id=chunk_id ) self.callbacks.on_location( PipelineParallelMicroStepLocations.FORWARD_END, input_tensor=input_tensor, output_tensor=output_tensor, step_id=step_id, ) backward_loss_tensor, backward_loss_fn_node = self._maybe_loss_compute( output_tensor, micro_dataset, overlap_schedule_mode ) if self.is_pipeline_first_stage() or self.is_pipeline_last_stage(): # Only increase micro batch id at virtual first/last pp stage. # The micro batch id is used to load data, therefore, only increase it when load data. self.micro_batch_id += 1 _release_input(input_tensor, output_tensor) if self._enable_timer: self.timers("forward_step").stop() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( f"[Pipeline details] After_forward_step_chunk_{chunk_id}_step_{step_id}" ) if self.is_pipeline_last_stage() and self._compute_loss: return backward_loss_tensor, schedule_chunk, backward_loss_fn_node return output_tensor, schedule_chunk, backward_loss_fn_node def _backward_step( self, input_tensor, output_tensor, output_tensor_grad, chunk_id=None, step_id=None, overlap_schedule_mode=False, schedule_chunk=None, loss_fn_node=None, ): if self.user_hooks_enabled: self.backward_hooks.run_hook() if self._enable_timer: self.timers("backward_step").start() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( f"[Pipeline details] Before_backward_step_chunk_{chunk_id}_step_{step_id}" ) with paddle.amp.auto_cast(enable=False): self.callbacks.on_location( PipelineParallelMicroStepLocations.BACKWARD_BEGIN, input_tensor=input_tensor, output_tensor=output_tensor, output_tensor_grad=output_tensor_grad, step_id=step_id, ) if self.is_pipeline_last_stage(): assert output_tensor_grad is None if overlap_schedule_mode: assert ( loss_fn_node is not None and schedule_chunk is not None ), ( "loss_fn_node and schedule_chunk should not be None in overlap_schedule_mode" ) input_tensor_grad = loss_fn_node.backward( scaler=self.scaler ) input_tensor_grad = schedule_chunk.backward( input_tensor_grad ) else: # In align mode, we scale the grad directly after forward if paddle.distributed.in_auto_parallel_align_mode(): output_tensor = output_tensor / _get_align_mode_scale() if self.scaler: paddle.autograd.backward( self.scaler.scale(output_tensor) ) else: paddle.autograd.backward(output_tensor) else: if isinstance(output_tensor, tuple): outputs = [t for t in output_tensor if not t.stop_gradient] assert len(outputs) == len(output_tensor_grad) grad_tensors = list(output_tensor_grad) else: outputs = [output_tensor] grad_tensors = [output_tensor_grad] if overlap_schedule_mode: assert schedule_chunk is not None, ( "schedule_chunk should not be None in overlap_schedule_mode" ) input_tensor_grad = schedule_chunk.backward(grad_tensors) else: paddle.autograd.backward( tensors=outputs, grad_tensors=grad_tensors, ) if not overlap_schedule_mode: # Extract input_tensor_grad from the input tensor. In overlap_schedule_mode, # the input_tensor_grad is extracted inside the schedule_chunk. input_tensor_grad = None if input_tensor is not None: if isinstance(input_tensor, tuple): input_tensor_grad = tuple( [ t.grad for t in input_tensor if not t.stop_gradient ] ) else: input_tensor_grad = input_tensor.grad if self._enable_timer: self.timers("backward_step").stop() self.callbacks.on_location( PipelineParallelMicroStepLocations.BACKWARD_END, input_tensor=input_tensor, output_tensor=output_tensor, input_tensor_grad=input_tensor_grad, output_tensor_grad=output_tensor_grad, step_id=step_id, ) if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( f"[Pipeline details] After_backward_step_chunk_{chunk_id}_step_{step_id}" ) return input_tensor_grad def _check_micro_batch_data_valid(self, micro_batch_data): if isinstance(micro_batch_data, (tuple, list)): for data in micro_batch_data: self._check_micro_batch_data_valid(data) elif isinstance(micro_batch_data, dict): for value in micro_batch_data.values(): self._check_micro_batch_data_valid(value) elif micro_batch_data is not None: assert isinstance(micro_batch_data, paddle.Tensor) def _broadcast_final_loss(self, return_micro_batch_loss=False): # Since the last backward run in interleave will set the virtual rank to 0, # here we need to check last stage ignoring virtual stage. if self.is_pipeline_last_stage(ignore_virtual=True): assert self.total_loss is not None, ( "train_batch() in last stage should obtain valid loss" ) losses = [] for idx in range(len(self._layers._loss_fn)): self.total_loss[idx] = paddle.to_tensor(self.total_loss[idx]) if not return_micro_batch_loss: # TODO(shenliang03): it will use mean/sum to calculate loss tmp = paddle.zeros_like(self.total_loss[idx][0]) for loss in self.total_loss[idx]: tmp += loss.detach() losses.append(tmp / self.accumulate_steps) else: losses.append(self.total_loss[idx].detach()) for idx in range(len(self._layers._loss_fn)): is_fp32 = ( paddle.full([], 1, 'int64') if losses[idx].dtype == paddle.float32 else paddle.full([], 0, 'int64') ) paddle.distributed.broadcast( is_fp32, src=self.global_rank, sync_op=True, group=self.pp_group, ) paddle.distributed.broadcast( losses[idx], src=self.global_rank, sync_op=True, group=self.pp_group, ) else: losses = [] for idx in range(len(self._layers._loss_fn)): is_fp32 = paddle.full([], 1, 'int64') paddle.distributed.broadcast( is_fp32, src=self._hcg.get_rank_from_stage(self.num_stages - 1), sync_op=True, group=self.pp_group, ) if return_micro_batch_loss: loss_shape = [self.accumulate_steps] else: loss_shape = [1] losses.append( paddle.zeros(shape=loss_shape, dtype="float32") if is_fp32.item() else paddle.zeros(shape=loss_shape, dtype="float16") ) paddle.distributed.broadcast( losses[idx], src=self._hcg.get_rank_from_stage(self.num_stages - 1), sync_op=True, group=self.pp_group, ) return losses[0] if len(losses) == 1 else losses def _optimizer_step(self): for p in self._layers.parameters(): if hasattr(p, "main_grad") and p.main_grad is not None: assert p.grad is None p.main_grad = p.main_grad.scale(1.0 / self.accumulate_steps) elif p.grad is not None: p.grad = p.grad.scale(1.0 / self.accumulate_steps) if self.scaler: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() if self._release_gradients: self.optimizer.clear_grad(set_to_zero=False) for _, buffers in self._chunk_2_comm_buffers.items(): for buffer in buffers: buffer._clear_grad_storage() else: self.optimizer.clear_grad() if self.lr_scheduler: self.lr_scheduler.step() def _offload_tensors(self, output_tensor): if not self._return_host_tensor: return if isinstance(output_tensor, (tuple, list)): for t in output_tensor: if not isinstance(t, paddle.Tensor) or isinstance( t, paddle.base.framework.EagerParamBase ): continue host_tensor = ( t.pin_memory() if hasattr(t, "pin_memory") else t.cpu() ) host_tensor._share_buffer_to(t) else: if not isinstance(output_tensor, paddle.Tensor): return host_tensor = ( output_tensor.pin_memory() if hasattr(output_tensor, "pin_memory") else output_tensor.cpu() ) host_tensor._share_buffer_to(output_tensor) def _release_output(self, output): def can_free(t): return ( t is not None and isinstance(t, paddle.Tensor) and t._is_initialized() and (t.inplace_version == 0 or getattr(t, "pp_can_free", False)) ) if isinstance(output, (tuple, list)): for t in output: if can_free(t): t._clear_dataptr() elif can_free(output): output._clear_dataptr() def get_static_scheduler(self): return self.forward_backward_pipeline(data=None, static_scheduler=True) @dataclass class P2PAsyncHandle: # funcs forward_handle_wait_fn: Callable forward_async_comm_fn: Callable backward_handle_wait_fn: Callable backward_async_comm_fn: Callable # outputs next_forward_virtual_pp_rank = None input_tensor = None out_fwd_wait_handles = None next_backward_virtual_pp_rank = None output_tensor_grad = None recv_next = None out_bwd_wait_handles = None def forward_handle_wait(self): self.forward_handle_wait_fn() def forward_async_comm(self, output_tensor): ( self.next_forward_virtual_pp_rank, self.input_tensor, self.out_fwd_wait_handles, ) = self.forward_async_comm_fn(output_tensor=output_tensor) def backward_handle_wait(self): self.backward_handle_wait_fn() def backward_async_comm(self, input_tensor_grad): ( self.next_backward_virtual_pp_rank, self.output_tensor_grad, self.recv_next, self.out_bwd_wait_handles, ) = self.backward_async_comm_fn(input_tensor_grad=input_tensor_grad) class PipelineParallelWithInterleave(PipelineParallel): # pipeline parallel with interleave scheduler def __init__(self, layers, hcg, strategy): super().__init__(layers=layers, hcg=hcg, strategy=strategy) self.overlap_schedule_mode = ( hasattr(type(self._layers), "overlapped_forward_backward") and self._strategy.hybrid_configs[ "pp_configs" ].forward_backward_overlap_scheduler ) if self.overlap_schedule_mode: assert not self._profiling, ( "Profiling is not compatible with overlap_schedule_mode." ) logger.info(f"Using {self._get_scheduler_name()}") self._record_format = ( '"name": "{}{}_VP{}", "cat": "virtual pipeline timeline", "ph": {}, "pid": 0, "tid": ' + str(self.stage_id + 1) + ', "ts": {}, "cname": "{}"' ) self._forward_colors = [ "thread_state_running", # RGB: 126, 200, 148 "thread_state_unknown", # RGB: 199, 155, 125 ] self._backward_colors = [ "rail_load", # RGB: 13, 168, 97 "rail_idle", # RGB: 238, 142, 0 ] # Structures to record the micro step for each layer chunk self._forward_micro_step_counter = {} self._backward_micro_step_counter = {} assert layers.get_num_virtual_stages() > 1 # setup for interleave scheduler self._check_sanity() self.num_model_chunks = layers.get_num_virtual_stages() self.model_chunks = layers.get_model_chunks() assert self.model_chunks is not None assert len(self.model_chunks) == self.num_model_chunks self._virtual_pp_world_size = self.num_model_chunks self._virtual_pp_rank = 0 self._reset_counter() self._best_unbalanced_scheduler = self._strategy.hybrid_configs[ "pp_configs" ].best_unbalanced_scheduler if self._best_unbalanced_scheduler: assert not self._comm_overlap, ( "pp best unbalaced scheduler can not run together with dp/sharding overlap" ) self._enable_offload_queue = self._strategy.hybrid_configs[ "pp_configs" ].enable_offload_queue # reinit user hook since now we have virtual stages self._init_user_hooks() def _get_scheduler_name(self): return f"PipelineParallelWithInterleave with overlapping forward backward={self.overlap_schedule_mode}, overlap p2p comm={self._overlap_p2p_comm}" def _init_user_bubble_hooks(self): # initialize bubble hooks self.bubble_hooks = PipelineHook() self.bubble_hooks.set_hooks_capacity(2 * self.num_stages - 2) def _check_sanity(self): assert framework.in_dynamic_mode(), ( "virtual pipeline stage with interleave only support eager dygraph mode" ) assert self.num_stages > 2, ( "virtual pipeline must run under pp degree > 2" ) assert self.accumulate_steps >= 2 * self.num_stages, ( f"accumulate_steps({self.accumulate_steps}) should be greater than or equal to 2 * num_stages({self.num_stages}) for pipeline with interleave" ) def _reset_counter(self): for i in range(self.num_model_chunks): self._forward_micro_step_counter[i] = 0 self._backward_micro_step_counter[i] = 0 def _record_stamp(self, name, step, phase, forward=True): if self._profiling: paddle.device.synchronize() virtual_pp_rank = self._get_virtual_pp_rank(step, forward=forward) color_idx = virtual_pp_rank % 2 # Get the profile color and micro step for current layer chunk if forward: color = self._forward_colors[color_idx] micro_step = self._forward_micro_step_counter[virtual_pp_rank] if phase == '"E"': self._forward_micro_step_counter[virtual_pp_rank] += 1 else: color = self._backward_colors[color_idx] micro_step = self._backward_micro_step_counter[virtual_pp_rank] if phase == '"E"': self._backward_micro_step_counter[virtual_pp_rank] += 1 self._records.append( '{' + self._record_format.format( name, micro_step, virtual_pp_rank, phase, int(time.time() * 1000), color, ) + '}' ) def _flush_records(self): if self._profiling: with open( f'./profile_record_tmp_file_for_rank_{self.global_rank}', 'a+', ) as f: f.writelines(record + '\n' for record in self._records) self._records = [] self._reset_counter() def _get_virtual_pp_rank(self, micro_step, forward): first_chunk_acc = ( self.accumulate_steps % self.num_stages + self.num_stages ) first_chunk_steps = first_chunk_acc * self.num_model_chunks if self._best_unbalanced_scheduler: num_group_last_chunk_forward = ( (micro_step - first_chunk_acc) // self.num_stages ) // self.num_model_chunks misplace_start = ( first_chunk_acc + self.num_model_chunks * self.num_stages * num_group_last_chunk_forward ) misplace_end = ( self.accumulate_steps % self.num_stages + num_group_last_chunk_forward * self.num_stages ) * self.num_model_chunks + self.num_stages forward_virtual_pp_stage = ( (micro_step - first_chunk_acc) // self.num_stages ) % self.num_model_chunks if micro_step < first_chunk_steps: virtual_pp_stage = micro_step // first_chunk_acc if not forward and self._best_unbalanced_scheduler: if ( micro_step >= first_chunk_acc + (self.num_model_chunks - 1) * self.num_stages ): if forward_virtual_pp_stage == self.num_model_chunks - 1: virtual_pp_stage = 0 elif ( micro_step >= misplace_start and micro_step < misplace_end ): virtual_pp_stage = ( micro_step - self.num_stages ) // first_chunk_acc else: origin_micro_step = micro_step micro_step -= first_chunk_steps virtual_pp_stage = micro_step % ( self.num_stages * self.num_model_chunks ) virtual_pp_stage = virtual_pp_stage // self.num_stages if not forward and self._best_unbalanced_scheduler: total_num_forward_step_from_steady = ( first_chunk_acc + (self.accumulate_steps - first_chunk_acc) * self.num_model_chunks ) if ( origin_micro_step <= total_num_forward_step_from_steady and forward_virtual_pp_stage == self.num_model_chunks - 1 ): virtual_pp_stage = 0 elif ( misplace_start <= total_num_forward_step_from_steady and origin_micro_step >= misplace_start and origin_micro_step < misplace_end ): if origin_micro_step < first_chunk_steps + self.num_stages: virtual_pp_stage = ( origin_micro_step - self.num_stages ) // first_chunk_acc else: virtual_pp_stage = (micro_step - self.num_stages) % ( self.num_stages * self.num_model_chunks ) virtual_pp_stage = virtual_pp_stage // self.num_stages if not forward: virtual_pp_stage = self.num_model_chunks - virtual_pp_stage - 1 return virtual_pp_stage def _get_forward_input(self, virtual_pp_rank): # some checkers assert hasattr(self, 'input_tensors') assert hasattr(self, 'output_tensors') if not self._forward_only: assert hasattr(self, 'output_tensor_grads') assert len(self.input_tensors[virtual_pp_rank]) == ( len(self.output_tensors[virtual_pp_rank]) + 1 ) input_tensor = self.input_tensors[virtual_pp_rank][-1] else: input_tensor = self.input_tensors[virtual_pp_rank].pop() return input_tensor def _store_forward_outputs( self, virtual_pp_rank, output_tensor, schedule_chunk=None, loss_fn_node=None, ): self.output_tensors[virtual_pp_rank].append(output_tensor) # If overlap_schedule_mode eq False, the schedule chunk is a None self.schedule_chunks[virtual_pp_rank].append(schedule_chunk) if self.is_pipeline_last_stage(): self.loss_fn_chunks.append(loss_fn_node) if self._forward_only: # no need to store tensor for backward if self._compute_loss: self.output_tensors[virtual_pp_rank].pop() # save output_tensors for return value of eval batch else: self._offload_tensors(output_tensor) else: # no need to store tensor for backward if self._forward_only: self.output_tensors[virtual_pp_rank].pop() def _forward_step_helper( self, micro_dataset, micro_step, overlap_schedule_mode=False, check_is_last_chunk=False, ): virtual_pp_rank = self._get_virtual_pp_rank(micro_step, forward=True) if check_is_last_chunk and virtual_pp_rank == self.num_model_chunks - 1: os.environ["FLAGS_last_vpp_chunk_forward"] = "1" self.set_virtual_pipeline_rank(virtual_pp_rank) input_tensor = self._get_forward_input(virtual_pp_rank) input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor) output_tensor, schedule_chunk, loss_fn_node = self._forward_step( input_tensor_dict if use_dict else input_tensor, micro_dataset, virtual_pp_rank, # chunk_id step_id=micro_step, overlap_schedule_mode=overlap_schedule_mode, ) output_tensor_tuple = dict_to_tuple_helper(output_tensor) self._store_forward_outputs( virtual_pp_rank, output_tensor_tuple, schedule_chunk, loss_fn_node ) return output_tensor_tuple def _overlap_comm_grads(self): if self._comm_overlap: self._backward_step_count += 1 sync_step = self._backward_step_count - self.stage_id if sync_step > 0 and sync_step % self.num_stages == 0: chunk_idx = self._virtual_pp_world_size - ( sync_step // self.num_stages ) for buffer in self._chunk_2_comm_buffers[chunk_idx]: buffer.comm_grads() if self.stage_id != 0: if ( self._backward_step_count == self.num_stages * self.num_model_chunks ): for buffer in self._chunk_2_comm_buffers[0]: buffer.comm_grads() def _sync_overlap_grads(self): if self._comm_overlap: assert ( self._backward_step_count == self.num_stages * self.num_model_chunks ), ( "backward step count should be equal to accumulate steps * virtual pp world size," f" but get {self._backward_step_count}, excepted result is {self.num_stages * self.num_model_chunks}" ) for _, buffers in self._chunk_2_comm_buffers.items(): for buffer in buffers: buffer.scale_grads() def _get_backward_input(self, virtual_pp_rank): # some checkers assert hasattr(self, 'input_tensors') assert hasattr(self, 'output_tensors') assert hasattr(self, 'output_tensor_grads') assert len(self.output_tensor_grads[virtual_pp_rank]) > 0, ( f"output_tensor_grads is empty for virtual_pp_rank {virtual_pp_rank}" ) assert len(self.input_tensors[virtual_pp_rank]) > 0 assert len(self.output_tensors[virtual_pp_rank]) > 0 input_tensor = self.input_tensors[virtual_pp_rank].pop(0) output_tensor = self.output_tensors[virtual_pp_rank].pop(0) output_tensor_grad = self.output_tensor_grads[virtual_pp_rank].pop(0) schedule_chunk = self.schedule_chunks[virtual_pp_rank].pop(0) if self.is_pipeline_last_stage(): loss_fn_node = self.loss_fn_chunks.pop(0) else: loss_fn_node = None return ( input_tensor, output_tensor, output_tensor_grad, schedule_chunk, loss_fn_node, ) def _backward_step_helper(self, micro_step, overlap_schedule_mode=False): virtual_pp_rank = self._get_virtual_pp_rank(micro_step, forward=False) self.set_virtual_pipeline_rank(virtual_pp_rank) ( input_tensor, output_tensor, output_tensor_grad, schedule_chunk, loss_fn_node, ) = self._get_backward_input(virtual_pp_rank) input_tensor_grad = self._backward_step( input_tensor, output_tensor, output_tensor_grad, chunk_id=virtual_pp_rank, step_id=micro_step, overlap_schedule_mode=overlap_schedule_mode, schedule_chunk=schedule_chunk, loss_fn_node=loss_fn_node, ) self._overlap_comm_grads() return input_tensor_grad def _forward_backward_helper( self, micro_dataset, forward_micro_step_id, backward_micro_step_id, p2p_async_handle=None, ): if not self.overlap_schedule_mode: if p2p_async_handle is not None: p2p_async_handle.forward_handle_wait() self._record_stamp("F", forward_micro_step_id, '"B"', forward=True) output_tensor = self._forward_step_helper( micro_dataset, forward_micro_step_id, ) self._record_stamp("F", forward_micro_step_id, '"E"', forward=True) if p2p_async_handle is not None: p2p_async_handle.forward_async_comm(output_tensor) p2p_async_handle.backward_handle_wait() # backward self._record_stamp( "B", backward_micro_step_id, '"B"', forward=False ) input_tensor_grad = self._backward_step_helper( backward_micro_step_id, ) self._record_stamp( "B", backward_micro_step_id, '"E"', forward=False ) if p2p_async_handle is not None: p2p_async_handle.backward_async_comm(input_tensor_grad) return else: return output_tensor, input_tensor_grad else: # 1. prepare forward inputs forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id, forward=True ) self.set_virtual_pipeline_rank(forward_virtual_pp_rank) if self.user_hooks_enabled: self.forward_hooks.run_hook() forward_inputs = self._get_forward_input(forward_virtual_pp_rank) input_tensor_dict, use_dict = tuple_to_dict_helper(forward_inputs) if self.is_pipeline_first_stage(): forward_inputs = next(micro_dataset)[0] self._check_micro_batch_data_valid(forward_inputs) if self.is_pipeline_last_stage(): labels = next(micro_dataset)[1] # 2. get forward chunks forward_chunk = self._layers.get_schedule_chunk( chunk_id=forward_virtual_pp_rank ) if self.is_pipeline_last_stage(): assert len(self._layers._loss_fn) == 1 forward_loss_fn_node = self._layers._loss_fn[ 0 ].build_schedule_node() forward_loss_fn_node.labels = labels else: forward_loss_fn_node = None # 3. prepare backward inputs & get backward chunks backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) self.set_virtual_pipeline_rank(backward_virtual_pp_rank) if self.user_hooks_enabled: self.backward_hooks.run_hook() ( _, _, backward_grads, backward_chunk, backward_loss_fn_node, ) = self._get_backward_input(backward_virtual_pp_rank) # 4. forward & backward if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] Start_forward_backward_step" ) if self._enable_timer: self.timers("forward_backward_step").start() output_tensor, forward_loss, input_tensor_grad = ( self._layers.overlapped_forward_backward( forward_chunk, input_tensor_dict if use_dict else forward_inputs, forward_loss_fn_node, backward_chunk, backward_loss_fn_node, backward_grads, self.scaler, p2p_async_handle=p2p_async_handle, ) ) output_tensor_tuple = dict_to_tuple_helper(output_tensor) if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] After_forward_backward_step" ) if self._enable_timer: self.timers("forward_backward_step").stop() # 5. process forward outputs forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id, forward=True ) self.set_virtual_pipeline_rank(forward_virtual_pp_rank) self._store_forward_outputs( forward_virtual_pp_rank, output_tensor_tuple, forward_chunk, forward_loss_fn_node, ) if self.is_pipeline_first_stage() or self.is_pipeline_last_stage(): # Only increase micro batch id at virtual first/last pp stage. # The micro batch id is used to load data, therefore, only increase it when load data. self.micro_batch_id += 1 if self.is_pipeline_last_stage(): # In overlap mode, only one loss_fn is supported. if self.total_loss is None: self.total_loss = [[]] self.total_loss[0].append(forward_loss.detach()) # 6. process backward outputs backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) self.set_virtual_pipeline_rank(backward_virtual_pp_rank) self._overlap_comm_grads() return output_tensor_tuple, input_tensor_grad def bw_hook_func(self, buffer, param): # For pipeline with interleave, we need to add grad to buffer without communication. # Use communication where appropriate to avoid dp communication and pp scheduling conflicts. # all reduce hook @paddle.autograd.no_grad() def fused_allreduce(*_): buffer.add_grad(param, use_comm=False) return fused_allreduce def register_allreduce_overlap_hook(self, model, comm_group, acc_steps, dp): super().register_allreduce_overlap_hook( model, comm_group, acc_steps, dp, group_size=sys.maxsize ) def _init_buffers(self): # init some data buffers for interleave scheduler self.input_tensors = [[] for _ in range(self.num_model_chunks)] self.output_tensors = [[] for _ in range(self.num_model_chunks)] self.output_tensor_grads = [[] for _ in range(self.num_model_chunks)] self.schedule_chunks = [[] for _ in range(self.num_model_chunks)] self.loss_fn_chunks = [] def forward_backward_pipeline( self, data, scaler, forward_only=False, compute_loss=True, static_scheduler=False, return_micro_batch_loss=False, ): """ Executes forward and backward passes for pipeline parallel training with interleaved scheduling. This method implements pipeline parallel training using interleaved scheduling strategy, inspired by Megatron-LM's implementation. It handles forward pass, backward pass, and gradient computation while managing communication and synchronization between stages. Args: data: Input data that will be wrapped into micro-batches scaler: Gradient scaler for mixed precision training forward_only: Whether to only perform forward pass (default: False) compute_loss: Whether to compute loss (default: True) return_micro_batch_loss: Whether to return micro-batch level loss (default: False) Returns: Training loss or logits if compute_loss is True; Otherwise returns output logits from the last stage Raises: AssertionError: - When compute_loss=False but forward_only=False - When cache is disabled but using interleaved pipeline - When buffers are not empty after execution Note: - Uses interleaved scheduling strategy (requires cache to be enabled) - Supports overlapping communication and computation for optimization - Handles startup phase, steady phase, and cooldown phase - Supports best unbalanced scheduler (_best_unbalanced_scheduler) """ self._reset_user_hooks_status() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] Start_forward_backward_step" ) # use interleave scheduling strategy. # this strategy is inspired by: # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/schedules.py if not compute_loss: assert forward_only, ( "compute_loss can only be set to False when forward_only is set to True" ) if static_scheduler: assert not forward_only, ( "static_scheduler only for training not for eval" ) assert not self._profiling, ( "While _profiling, static scheduler is not available" ) if data is not None: warnings.warn( "Static scheduler run won't real run the model, but data has been provided" ) logger.info( "enable static_scheduler will return the pp schedule instead of the loss" ) schedule = "" # NOTE(shenliang03): Due to ring_exchange for pipeline with interleave, cache should be enabled assert self._using_cache, ( "cache should be enabled for pipeline with interleave" ) self.overlap_schedule_mode = ( hasattr(type(self._layers), "overlapped_forward_backward") and self._strategy.hybrid_configs[ "pp_configs" ].forward_backward_overlap_scheduler ) if forward_only: self.overlap_schedule_mode = False # init some attributes for this batch run self.scaler = scaler self.total_loss = None self.micro_batch_id = 0 self._forward_only = forward_only self.user_hooks_enabled = not self._forward_only first_chunk_acc = ( self.accumulate_steps % self.num_stages + self.num_stages ) first_chunk_steps = first_chunk_acc * self.num_model_chunks fwd_buffer_queue = queue.Queue() bwd_buffer_queue = queue.Queue() skip_steps = self.accumulate_steps % self.num_stages last_stage_recv_queue = deque() left_id = skip_steps right_id = left_id + first_chunk_acc * (self.num_model_chunks - 1) def _process_fwd_buffer(step_id, tensor): if step_id < first_chunk_steps: if not self.is_pipeline_last_stage(): fwd_buffer_queue.put(tensor) if left_id <= step_id < right_id: tensor = fwd_buffer_queue.get() else: tensor = None else: if self.is_pipeline_last_stage(): tensor = None return tensor def _last_stage_need_recv_next(micro_step): if micro_step >= first_chunk_acc: if len(last_stage_recv_queue) == 0: return False else: res = last_stage_recv_queue[0] if micro_step - res[0] < self.num_stages: return False else: return True else: return False def _last_stage_recv_pp_rank(micro_step): if micro_step >= first_chunk_acc: assert len(last_stage_recv_queue) != 0, ( "last_stage_recv_queue can't be empty" ) virtual_pp_stage = (last_stage_recv_queue.popleft())[1] return virtual_pp_stage - 1 else: return self.num_model_chunks - 1 def _process_bwd_buffer(step_id, tensor): if self._best_unbalanced_scheduler: if not self.is_pipeline_first_stage(): bwd_buffer_queue.put(tensor) if step_id >= left_id and not bwd_buffer_queue.empty(): tensor = bwd_buffer_queue.get() else: tensor = None else: if step_id < first_chunk_steps: if not self.is_pipeline_first_stage(): bwd_buffer_queue.put(tensor) if left_id <= step_id < right_id: tensor = bwd_buffer_queue.get() else: tensor = None else: if self.is_pipeline_first_stage(): tensor = None return tensor per_stage_accumulate_steps = self.accumulate_steps // self.num_stages self._backward_step_count = -( first_chunk_steps + (per_stage_accumulate_steps - 2) * self.num_stages * self.num_model_chunks ) self._init_buffers() micro_dataset = self._wrap_data(data) num_steps = self.accumulate_steps * self.num_model_chunks if forward_only: # If only forward, since there is no backward during running, all steps are startup steps startup_steps = num_steps else: # actually startup_steps is calculated from two number: # first_forward_cross_to_end = (self.num_stages - self.stage_id - 1) + (self.num_model_chunks - 1) * self.num_stages # end_to_first_backward_cross = (self.num_stages - self.stage_id - 1) # startup_steps = first_forward_cross_to_end + end_to_first_backward_cross startup_steps = (self.num_stages - self.stage_id - 1) * 2 startup_steps += (self.num_model_chunks - 1) * first_chunk_acc startup_steps = min(startup_steps, num_steps) # An additional micro step is needed for overplapping schedule if self.overlap_schedule_mode: startup_steps += 1 steady_steps = num_steps - startup_steps for location in range(self.stage_id): if self.user_hooks_enabled: self.bubble_hooks.run_hook() rest_bubble_times = self.num_stages - 1 - self.stage_id self.set_virtual_pipeline_rank(0) if not static_scheduler: self.input_tensors[0].append( self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), sync_recv=False, batch_p2p_comm=self._use_batch_p2p_comm, ) ) fwd_wait_handles = None bwd_wait_handles = None # run startup steps for micro_step in range(startup_steps): if fwd_wait_handles is not None: for req in fwd_wait_handles: req.wait() if static_scheduler: virtual_pp_rank = self._get_virtual_pp_rank( micro_step, forward=True ) real_micro_step = self._forward_micro_step_counter[ virtual_pp_rank ] self._forward_micro_step_counter[virtual_pp_rank] += 1 schedule += f"f{real_micro_step}_vp{virtual_pp_rank};" logger.info( f"forward step for {real_micro_step} with virtual pp rank {virtual_pp_rank}" ) continue self._record_stamp("F", micro_step, '"B"', forward=True) output_tensor = self._forward_step_helper( micro_dataset, micro_step, overlap_schedule_mode=self.overlap_schedule_mode, ) self._record_stamp("F", micro_step, '"E"', forward=True) if micro_step >= startup_steps - rest_bubble_times: if self.user_hooks_enabled: self.bubble_hooks.run_hook() # determine whether recv forward tensor or not next_virtual_pp_rank = self._get_virtual_pp_rank( micro_step + 1, forward=True ) recv_prev = True if self.is_pipeline_first_stage(ignore_virtual=True): if next_virtual_pp_rank == 0: # next chunk is the first chunk, not need to pre recv an input tensor recv_prev = False # last micro step, no next run if micro_step == (num_steps - 1): recv_prev = False # last stage shouldn't send tensor to downstream if self.is_pipeline_last_stage(ignore_virtual=True): output_tensor = _process_fwd_buffer(micro_step, output_tensor) if not self._overlap_p2p_comm: # prepare for the first steady step if ( micro_step == (startup_steps - 1) and (not forward_only) and steady_steps ): input_tensor_grad = None recv_next = True if self.is_pipeline_last_stage(ignore_virtual=True): recv_next = False # the last startup step needs on four direction comm to set up for steady 1f1b ( input_tensor, output_tensor_grad, ) = self._p2p_helper.send_forward_backward_recv_forward_backward( output_tensor, input_tensor_grad, recv_prev=recv_prev, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, skip_check_meta=not self.training, ) # output_tensor_grad is not none if recv_next # append output_tensor_grad no matter none or not self.output_tensor_grads[self.num_model_chunks - 1].append( output_tensor_grad ) else: input_tensor = self._p2p_helper.send_forward_recv_forward( output_tensor, recv_prev=recv_prev, batch_p2p_comm=self._use_batch_p2p_comm, skip_check_meta=not self.training, ) # append input_tensor no matter none or not self.input_tensors[next_virtual_pp_rank].append(input_tensor) else: ( input_tensor, fwd_wait_handles, ) = self._p2p_helper.send_forward_recv_forward( output_tensor, recv_prev=recv_prev, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, skip_check_meta=not self.training, ) if ( micro_step == (startup_steps - 1) and (not forward_only) and steady_steps ): input_tensor_grad = None recv_next = True if self.is_pipeline_last_stage(ignore_virtual=True): recv_next = False ( output_tensor_grad, bwd_wait_handles, ) = self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) self.output_tensor_grads[self.num_model_chunks - 1].append( output_tensor_grad ) # append input_tensor no matter none or not self.input_tensors[next_virtual_pp_rank].append(input_tensor) _release_output(output_tensor) # run 1f1b steady steps for micro_step in range(steady_steps): if static_scheduler: forward_micro_step_id = micro_step + startup_steps forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id, forward=True ) backward_micro_step_id = micro_step backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) real_forward_micro_step = self._forward_micro_step_counter[ forward_virtual_pp_rank ] self._forward_micro_step_counter[forward_virtual_pp_rank] += 1 real_backward_micro_step = self._backward_micro_step_counter[ backward_virtual_pp_rank ] self._backward_micro_step_counter[backward_virtual_pp_rank] += 1 schedule += ( f"f{real_forward_micro_step}_vp{forward_virtual_pp_rank};" ) schedule += ( f"b{real_backward_micro_step}_vp{backward_virtual_pp_rank};" ) logger.info( f"forward step for {real_forward_micro_step} with virtual pp rank {forward_virtual_pp_rank}" ) logger.info( f"backward step for {real_backward_micro_step} with virtual pp rank {backward_virtual_pp_rank}" ) continue # forward forward_micro_step_id = micro_step + startup_steps if self._overlap_p2p_comm: backward_micro_step_id = micro_step def forward_handle_wait(fwd_wait_handles): if fwd_wait_handles is not None: for req in fwd_wait_handles: req.wait() def forward_async_comm(forward_micro_step_id, output_tensor): forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id, forward=True ) self.set_virtual_pipeline_rank(forward_virtual_pp_rank) # determine whether to recv input tensor from upstream recv_prev = True if self.is_pipeline_first_stage(ignore_virtual=True): next_forward_virtual_pp_rank = ( self._get_virtual_pp_rank( forward_micro_step_id + 1, forward=True ) ) if next_forward_virtual_pp_rank == 0: # next chunk is the first chunk, not need to pre recv an input tensor recv_prev = False else: next_forward_virtual_pp_rank = ( self._get_virtual_pp_rank( forward_micro_step_id + 1, forward=True ) ) # last iteration doesn't need recv from upstream if micro_step == (steady_steps - 1): recv_prev = False if self.is_pipeline_last_stage(ignore_virtual=True): output_tensor = _process_fwd_buffer( forward_micro_step_id, output_tensor ) # Send activation tensor to the next stage and receive activation tensor from the # previous stage ( input_tensor, fwd_wait_handles, ) = self._p2p_helper.send_forward_recv_forward( output_tensor, recv_prev=recv_prev, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, skip_check_meta=not self.training, ) _release_output(output_tensor) return ( next_forward_virtual_pp_rank, input_tensor, fwd_wait_handles, ) def backward_handle_wait(bwd_wait_handles): if bwd_wait_handles is not None: for req in bwd_wait_handles: req.wait() def backward_async_comm( backward_micro_step_id, input_tensor_grad ): if ( self._best_unbalanced_scheduler and self.is_pipeline_last_stage(ignore_virtual=True) ): cur_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) if cur_pp_rank != 0: last_stage_recv_queue.append( (backward_micro_step_id, cur_pp_rank) ) # first stage doesn't send grad to upstream backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) self.set_virtual_pipeline_rank(backward_virtual_pp_rank) if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor_grad = _process_bwd_buffer( backward_micro_step_id, input_tensor_grad ) recv_next = True if self.is_pipeline_last_stage(ignore_virtual=True): if self._best_unbalanced_scheduler: next_backward_virtual_pp_rank = ( self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False, ) ) if self.is_pipeline_last_stage(ignore_virtual=True): recv_next = _last_stage_need_recv_next( backward_micro_step_id + 1 ) else: next_backward_virtual_pp_rank = ( self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False, ) ) if next_backward_virtual_pp_rank == ( self.num_model_chunks - 1 ): # next chunk is the last chunk, not need to pre recv an output tensor grad recv_next = False else: next_backward_virtual_pp_rank = ( self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False, ) ) ( output_tensor_grad, bwd_wait_handles, ) = self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) return ( next_backward_virtual_pp_rank, output_tensor_grad, recv_next, bwd_wait_handles, ) # Package some closure functions and parameters into `P2PAsyncHandle` # structure to simplify function parameter passing p2p_async_handle = P2PAsyncHandle( partial( forward_handle_wait, fwd_wait_handles=fwd_wait_handles ), partial( forward_async_comm, forward_micro_step_id=forward_micro_step_id, ), partial( backward_handle_wait, bwd_wait_handles=bwd_wait_handles ), partial( backward_async_comm, backward_micro_step_id=backward_micro_step_id, ), ) self._forward_backward_helper( micro_dataset, forward_micro_step_id, backward_micro_step_id, p2p_async_handle, ) # Information that needs to be updated next_forward_virtual_pp_rank = ( p2p_async_handle.next_forward_virtual_pp_rank ) input_tensor = p2p_async_handle.input_tensor fwd_wait_handles = p2p_async_handle.out_fwd_wait_handles next_backward_virtual_pp_rank = ( p2p_async_handle.next_backward_virtual_pp_rank ) output_tensor_grad = p2p_async_handle.output_tensor_grad recv_next = p2p_async_handle.recv_next bwd_wait_handles = p2p_async_handle.out_bwd_wait_handles else: backward_micro_step_id = micro_step output_tensor, input_tensor_grad = ( self._forward_backward_helper( micro_dataset, forward_micro_step_id, backward_micro_step_id, ) ) if ( self._best_unbalanced_scheduler and self.is_pipeline_last_stage(ignore_virtual=True) ): cur_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) if cur_pp_rank != 0: last_stage_recv_queue.append( (backward_micro_step_id, cur_pp_rank) ) # four directions comm # send output tensor to downstream # send input tensor grad to upstream # recv input tensor from upstream # recv output tensor grad from downstream # last stage doesn't send rst to downstream forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id, forward=True ) self.set_virtual_pipeline_rank(forward_virtual_pp_rank) if self.is_pipeline_last_stage(ignore_virtual=True): output_tensor = _process_fwd_buffer( forward_micro_step_id, output_tensor ) # first stage doesn't send grad to upstream backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) self.set_virtual_pipeline_rank(backward_virtual_pp_rank) if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor_grad = _process_bwd_buffer( backward_micro_step_id, input_tensor_grad ) # determine whether to recv input tensor from upstream recv_prev = True next_forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id + 1, forward=True ) if self.is_pipeline_first_stage(ignore_virtual=True) and ( next_forward_virtual_pp_rank == 0 ): # first pp stage and first virtual stage recv_prev = False # last iteration doesn't need recv from upstream if micro_step == (steady_steps - 1): recv_prev = False # determine whether to recv grad from downstream recv_next = True if self._best_unbalanced_scheduler: next_backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False, ) if self.is_pipeline_last_stage(ignore_virtual=True): recv_next = _last_stage_need_recv_next( backward_micro_step_id + 1 ) else: next_backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False ) if self.is_pipeline_last_stage(ignore_virtual=True) and ( next_backward_virtual_pp_rank == (self.num_model_chunks - 1) ): # last pp stage and last virtual stage recv_next = False ( input_tensor, output_tensor_grad, ) = self._p2p_helper.send_forward_backward_recv_forward_backward( output_tensor, input_tensor_grad, recv_prev=recv_prev, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, skip_check_meta=not self.training, ) # append input_tensor no matter none or not self.input_tensors[next_forward_virtual_pp_rank].append( input_tensor ) # append output_tensor_grad no matter none or not if self._best_unbalanced_scheduler: if self.is_pipeline_last_stage(ignore_virtual=True): if recv_next: recv_next_virtual_pp_rank = _last_stage_recv_pp_rank( backward_micro_step_id + 1 ) self.output_tensor_grads[ recv_next_virtual_pp_rank ].append(output_tensor_grad) if ( next_backward_virtual_pp_rank == self.num_model_chunks - 1 and recv_next_virtual_pp_rank != next_backward_virtual_pp_rank ): self.output_tensor_grads[ self.num_model_chunks - 1 ].append(None) elif ( next_backward_virtual_pp_rank == self.num_model_chunks - 1 ): self.output_tensor_grads[ self.num_model_chunks - 1 ].append(None) else: self.output_tensor_grads[ next_backward_virtual_pp_rank ].append(output_tensor_grad) else: self.output_tensor_grads[next_backward_virtual_pp_rank].append( output_tensor_grad ) _release_output(output_tensor) assert fwd_buffer_queue.empty(), "forward buffer should be empty" if not static_scheduler: _release_output(output_tensor) # remaining backward steps if not forward_only: if self._overlap_p2p_comm and bwd_wait_handles is not None: for wait_handles in bwd_wait_handles: wait_handles.wait() # no steady steps, which only occurs when accumulate_step == num_stage if not steady_steps: output_tensor_grad = self._p2p_helper.recv_backward( self.is_pipeline_last_stage(), batch_p2p_comm=self._use_batch_p2p_comm, ) self.output_tensor_grads[self.num_model_chunks - 1].append( output_tensor_grad ) for micro_step in range(steady_steps, num_steps): if static_scheduler: virtual_pp_rank = self._get_virtual_pp_rank( micro_step, forward=False ) real_micro_step = self._backward_micro_step_counter[ virtual_pp_rank ] self._backward_micro_step_counter[virtual_pp_rank] += 1 schedule += f"b{real_micro_step}_vp{virtual_pp_rank};" logger.info( f"backward step for {real_micro_step} with virtual pp rank {virtual_pp_rank}" ) continue if ( micro_step < steady_steps + self.num_stages - 1 - self.stage_id ) and self.user_hooks_enabled: self.bubble_hooks.run_hook() # cooldown loop self._record_stamp("B", micro_step, '"B"', forward=False) input_tensor_grad = self._backward_step_helper( micro_step, overlap_schedule_mode=self.overlap_schedule_mode ) self._record_stamp("B", micro_step, '"E"', forward=False) next_backward_virtual_pp_rank = self._get_virtual_pp_rank( micro_step + 1, forward=False, ) if ( self._best_unbalanced_scheduler and self.is_pipeline_last_stage(ignore_virtual=True) ): cur_pp_rank = self._get_virtual_pp_rank( micro_step, forward=False ) if cur_pp_rank != 0: last_stage_recv_queue.append((micro_step, cur_pp_rank)) recv_next = True if self.is_pipeline_last_stage(ignore_virtual=True): if self._best_unbalanced_scheduler: recv_next = _last_stage_need_recv_next(micro_step + 1) else: if next_backward_virtual_pp_rank == ( self.num_model_chunks - 1 ): recv_next = False if micro_step == (num_steps - 1): recv_next = False if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor_grad = _process_bwd_buffer( micro_step, input_tensor_grad ) # append output_tensor_grad no matter none or not if self._best_unbalanced_scheduler: if self.is_pipeline_last_stage(ignore_virtual=True): output_tensor_grad = ( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, ) ) if recv_next: recv_next_virtual_pp_rank = ( _last_stage_recv_pp_rank(micro_step + 1) ) self.output_tensor_grads[ recv_next_virtual_pp_rank ].append(output_tensor_grad) else: self.output_tensor_grads[ next_backward_virtual_pp_rank ].append(output_tensor_grad) else: self.output_tensor_grads[ next_backward_virtual_pp_rank ].append( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, ) ) else: self.output_tensor_grads[ next_backward_virtual_pp_rank ].append( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, ) ) self._sync_overlap_grads() for _ in range(self.stage_id): self.bubble_hooks.run_hook() if static_scheduler: self._reset_counter() return schedule if self._enable_timer: self.timers("allreduce_shared_weight_gradients").start() self._layers.allreduce_shared_weight_gradients() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").stop() self._flush_records() assert bwd_buffer_queue.empty(), "backward buffer should be empty" if compute_loss: # return loss if compute loss if self._enable_timer: self.timers("broadcast_final_loss").start() with paddle.amp.auto_cast(enable=False): train_loss_or_logits = self._broadcast_final_loss( return_micro_batch_loss ) if self._enable_timer: self.timers("broadcast_final_loss").stop() else: # else just return logits without loss func calc train_loss_or_logits = self.output_tensors.pop() if self._clear_every_step_cache: self._p2p_helper.clear_meta_cache() self.timer_printer() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] End_forward_backward_step" ) self.processed_steps += 1 self._check_user_hooks_status_at_step_end() # reset dynamic meta counter if self._dynamic_shape: assert self._p2p_helper._dynamic_cnt == len( self._p2p_helper._send_recv_meta_list ), "p2p dynamic_cnt should equal to send_recv_meta_list" self._p2p_helper._dynamic_cnt = 0 return train_loss_or_logits def train_batch( self, data, optimizer, lr_scheduler=None, scaler=None, loss_fn_idx=0, return_micro_batch_loss=False, ): """ Execute one training batch with pipeline parallel interleaving schedule. Performs forward/backward passes and optimizer update for a batch of data using pipeline parallel with interleaved scheduling. Args: data: Input data for the batch optimizer: Optimizer instance for parameter updates lr_scheduler: Learning rate scheduler (optional) scaler: Gradient scaler for mixed precision training (optional) loss_fn_idx: Index of loss function to use (default: 0) return_micro_batch_loss: Whether to return per-micro-batch losses (default: False) Returns: The computed training loss. If return_micro_batch_loss is True, returns a tuple of (total_loss, micro_batch_losses). Note: - Handles both FP16/FP32 mixed precision training when scaler is provided - Supports multiple loss functions through loss_fn_idx - Uses interleaved pipeline parallel schedule for efficient training """ data = self._prepare_training(data, optimizer, lr_scheduler) # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx # interleave scheduler for pipeline parallel train_loss = self.forward_backward_pipeline( data, scaler, return_micro_batch_loss=return_micro_batch_loss ) # optimizer with paddle.amp.auto_cast(enable=False): self._optimizer_step() return train_loss def eval_batch( self, data, compute_loss=False, loss_fn_idx=0, return_host_tensor=False ): self.user_hooks_enabled = False # reset the virtual pp rank for each run self.set_virtual_pipeline_rank(0) self._layers.eval() origin_compute_loss = self._compute_loss self._compute_loss = compute_loss origin_return_host_tensor = self._return_host_tensor self._return_host_tensor = return_host_tensor # check loss_fn_idx is valid and loss_fn exists assert ( loss_fn_idx in range(len(self._layers._loss_fn)) and self._layers._loss_fn[loss_fn_idx] is not None ), f"loss function {loss_fn_idx} should exist to compute loss" self.loss_fn_idx = loss_fn_idx train_loss_or_logits = self.forward_backward_pipeline( data, None, forward_only=True, compute_loss=compute_loss ) self._init_buffers() self._compute_loss = origin_compute_loss self._return_host_tensor = origin_return_host_tensor return train_loss_or_logits def get_static_scheduler(self): return self.forward_backward_pipeline( data=None, scaler=None, static_scheduler=True ) class PipelineParallelWithInterleaveFthenB(PipelineParallelWithInterleave): def __init__(self, layers, hcg, strategy): # Initialize the basic parameters of the parent class PipelineParallel super().__init__(layers=layers, hcg=hcg, strategy=strategy) # Whether to enable overlapped scheduling mode (disabled by default) self.overlap_schedule_mode = False def _get_scheduler_name(self): return "PipelineParallelWithInterleaveFthenB" def _init_user_bubble_hooks(self): # (TODO:gexiao) support bubble hooks if needed self.bubble_hooks = None # self.bubble_hooks = PipelineHook() # self.bubble_hooks.set_hooks_capacity(2 * self.num_stages - 2) def _check_sanity(self): assert framework.in_dynamic_mode(), ( "virtual pipeline stage with interleave only support eager dygraph mode" ) assert self.num_stages > 2, ( "virtual pipeline must run under pp degree > 2" ) def _get_virtual_pp_rank(self, micro_step, forward): virtual_pp_stage = micro_step % ( self.accumulate_steps * self.num_model_chunks ) virtual_pp_stage = virtual_pp_stage // self.accumulate_steps if not forward: virtual_pp_stage = self.num_model_chunks - virtual_pp_stage - 1 return virtual_pp_stage def _overlap_comm_grads(self): if not self._comm_overlap: return self._backward_step_count += 1 sync_step = self._backward_step_count - self.stage_id if sync_step > 0 and sync_step % self.accumulate_steps == 0: chunk_idx = self._virtual_pp_world_size - ( sync_step // self.accumulate_steps ) for buffer in self._chunk_2_comm_buffers[chunk_idx]: buffer.comm_grads() if self.stage_id == 0: return if ( self._backward_step_count == self.accumulate_steps * self._virtual_pp_world_size ): for buffer in self._chunk_2_comm_buffers[0]: buffer.comm_grads() def _sync_overlap_grads(self): if not self._comm_overlap: return expected_count = self.accumulate_steps * self._virtual_pp_world_size assert self._backward_step_count == expected_count, ( f"backward step count should be equal to accumulate steps * virtual pp world size, " f"but got {self._backward_step_count}, expected result is {expected_count}" ) for buffers in self._chunk_2_comm_buffers.values(): for buffer in buffers: buffer.scale_grads() def forward_backward_pipeline( self, data, scaler, forward_only=False, compute_loss=True, return_micro_batch_loss=False, ): self._reset_user_hooks_status() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] Start_forward_backward_step" ) if not compute_loss: assert forward_only, ( "compute_loss can only be set to False when forward_only is set to True" ) # NOTE(shenliang03): Due to ring_exchange for pipeline with interleave, cache should be enabled assert self._using_cache, ( "cache should be enabled for pipeline with interleave" ) # init some attributes for this batch run self.scaler = scaler self.total_loss = None self.micro_batch_id = 0 self._forward_only = forward_only self.user_hooks_enabled = not self._forward_only assert ( self.accumulate_steps == self.num_stages or self.accumulate_steps % self.num_stages == 0 ), ( f"accumulate_steps({self.accumulate_steps}) and num_stages({self.num_stages}) should be a multiple or accumulate_steps % num_stages == 0" ) self._backward_step_count = 0 skip_steps = self.accumulate_steps - self.num_stages send_recv_buffer_queue = queue.Queue() self._init_buffers() micro_dataset = self._wrap_data(data) num_steps = self.accumulate_steps * self.num_model_chunks self.set_virtual_pipeline_rank(0) self.input_tensors[0].append( self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), sync_recv=False, batch_p2p_comm=self._use_batch_p2p_comm, ) ) for micro_step in range(num_steps): output_tensor = self._forward_step_helper(micro_dataset, micro_step) # determine whether recv forward tensor or not next_virtual_pp_rank = self._get_virtual_pp_rank( micro_step + 1, forward=True ) recv_prev = True if self.is_pipeline_first_stage(ignore_virtual=True): if next_virtual_pp_rank == 0: # next chunk is the first chunk, not need to pre recv an input tensor recv_prev = False # last micro step, no next run if micro_step == (num_steps - 1): recv_prev = False if self.is_pipeline_last_stage(ignore_virtual=True): # last stage skip send/recv if not self.is_pipeline_last_stage(): send_recv_buffer_queue.put(output_tensor) if micro_step < skip_steps or ( self.is_pipeline_last_stage() and micro_step % self.accumulate_steps >= skip_steps ): output_tensor = None else: output_tensor = send_recv_buffer_queue.get() input_tensor = self._p2p_helper.send_forward_recv_forward( output_tensor, recv_prev=recv_prev, batch_p2p_comm=self._use_batch_p2p_comm, skip_check_meta=not self.training, ) self.input_tensors[next_virtual_pp_rank].append(input_tensor) _release_output(output_tensor) assert send_recv_buffer_queue.empty(), ( "send_recv buffer should be empty" ) # remaining backward steps if not forward_only: self.output_tensor_grads[self.num_model_chunks - 1].append( self._p2p_helper.recv_backward( self.is_pipeline_last_stage(), sync_recv=False, batch_p2p_comm=self._use_batch_p2p_comm, ) ) for micro_step in range(num_steps): # cooldown loop input_tensor_grad = self._backward_step_helper(micro_step) next_backward_virtual_pp_rank = self._get_virtual_pp_rank( micro_step + 1, forward=False ) recv_next = True if self.is_pipeline_last_stage(ignore_virtual=True): if next_backward_virtual_pp_rank == ( self.num_model_chunks - 1 ): recv_next = False if micro_step == (num_steps - 1): recv_next = False if self.is_pipeline_first_stage(ignore_virtual=True): if not self.is_pipeline_first_stage(): send_recv_buffer_queue.put(input_tensor_grad) if micro_step < skip_steps or ( self.is_pipeline_first_stage() and micro_step % self.accumulate_steps >= skip_steps ): input_tensor_grad = None else: input_tensor_grad = send_recv_buffer_queue.get() self.output_tensor_grads[next_backward_virtual_pp_rank].append( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, ) ) assert send_recv_buffer_queue.empty(), ( "send_recv buffer should be empty" ) self._sync_overlap_grads() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").start() self._layers.allreduce_shared_weight_gradients() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").stop() if compute_loss: # return loss if compute loss if self._enable_timer: self.timers("broadcast_final_loss").start() with paddle.amp.auto_cast(enable=False): train_loss_or_logits = self._broadcast_final_loss( return_micro_batch_loss ) if self._enable_timer: self.timers("broadcast_final_loss").stop() else: # else just return logits without loss func calc train_loss_or_logits = self.output_tensors.pop() if self._clear_every_step_cache: self._p2p_helper.clear_meta_cache() self.timer_printer() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] End_forward_backward_step" ) self.processed_steps += 1 self._check_user_hooks_status_at_step_end() return train_loss_or_logits class OffloadQueue(queue.Queue): def __init__(self, offload=False, *args, **kwargs): super().__init__(*args, **kwargs) self.offload = offload def put(self, tensor, *args, **kwargs): if self.offload and isinstance( tensor, (paddle.Tensor, paddle.base.framework.core.eager.Tensor) ): tensor_cpu = tensor.pin_memory() tensor_cpu._share_buffer_to(tensor) elif self.offload and isinstance(tensor, tuple): for t in tensor: if isinstance( t, (paddle.Tensor, paddle.base.framework.core.eager.Tensor) ): t_cpu = t.pin_memory() t_cpu._share_buffer_to(t) super().put(tensor, *args, **kwargs) def get(self, *args, **kwargs): tensor = super().get(*args, **kwargs) if self.offload and isinstance( tensor, (paddle.Tensor, paddle.base.framework.core.eager.Tensor) ): tensor = tensor.to(paddle.base.framework._current_expected_place()) elif self.offload and isinstance(tensor, tuple): for t in tensor: if isinstance( t, (paddle.Tensor, paddle.base.framework.core.eager.Tensor) ): t_dev = t.to( paddle.base.framework._current_expected_place() ) t_dev._share_buffer_to(t) return tensor class VPPFhenBInBalancedMemory(PipelineParallelWithInterleaveFthenB): def __init__(self, layers, hcg, strategy): super().__init__(layers=layers, hcg=hcg, strategy=strategy) self.overlap_schedule_mode = False def _get_scheduler_name(self): return "VPPFhenBInBalancedMemory" def _init_user_bubble_hooks(self): self.bubble_hooks = PipelineHook() self.bubble_hooks.set_hooks_capacity(2 * self.num_stages - 2) def forward_backward_pipeline( self, data, scaler, forward_only=False, compute_loss=True, return_micro_batch_loss=False, ): self._reset_user_hooks_status() if not compute_loss: assert forward_only, ( "compute_loss can only be set to False when forward_only is set to True" ) assert self._using_cache, ( "cache should be enabled for pipeline with interleave" ) self.user_hooks_enabled = not forward_only if forward_only: return super().forward_backward_pipeline( data, scaler, forward_only, compute_loss, return_micro_batch_loss, ) if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] Start_forward_backward_step" ) # init some attributes for this batch run self.scaler = scaler self.total_loss = None self.micro_batch_id = 0 self._forward_only = forward_only self._init_buffers() backward_send_recv_buffer_queue = OffloadQueue( offload=self._enable_offload_queue ) forward_send_recv_buffer_queue = OffloadQueue( offload=self._enable_offload_queue ) skip_steps = self.accumulate_steps - self.num_stages micro_dataset = self._wrap_data(data) num_steps = self.accumulate_steps * self.num_model_chunks # the whole pipeline is splited into 3 parse: # startup_steps, steady_1f1b_steps, cooldown_steps startup_steps = ( self.accumulate_steps * (self.num_model_chunks - 1) + self.num_stages - self.stage_id - 1 ) steady_1f1b_steps = self.accumulate_steps - ( self.num_stages - self.stage_id - 1 ) cooldown_steps = startup_steps # Bubbles before startup_steps for _ in range(self.stage_id): if self.user_hooks_enabled: self.bubble_hooks.run_hook() self.set_virtual_pipeline_rank(0) self.input_tensors[0].append( self._p2p_helper.recv_forward( self.is_pipeline_first_stage(), sync_recv=False, batch_p2p_comm=self._use_batch_p2p_comm, ) ) # In startup_steps, we send every output_tensor of last stage, # to simplify the code logic of stage 1F1B. for micro_step in range(startup_steps): self._record_stamp("F", micro_step, '"B"', forward=True) output_tensor = self._forward_step_helper(micro_dataset, micro_step) self._record_stamp("F", micro_step, '"E"', forward=True) next_forward_virtual_pp_rank = self._get_virtual_pp_rank( micro_step + 1, forward=True ) recv_prev = True if self.is_pipeline_first_stage(ignore_virtual=True) and ( micro_step < self.num_stages - 1 ): recv_prev = False input_tensor = self._p2p_helper.send_forward_recv_forward( output_tensor, recv_prev=recv_prev, batch_p2p_comm=self._use_batch_p2p_comm, skip_check_meta=not self.training, ) if self.is_pipeline_first_stage(ignore_virtual=True): if input_tensor is not None: # stash the input_tensor and it will be used in the next chunk later forward_send_recv_buffer_queue.put(input_tensor) if next_forward_virtual_pp_rank == 0: input_tensor = None else: # when a input_tensor is needed, get one from the queue input_tensor = forward_send_recv_buffer_queue.get() self.input_tensors[next_forward_virtual_pp_rank].append( input_tensor ) _release_output(output_tensor) if self.is_pipeline_first_stage(ignore_virtual=True): assert ( forward_send_recv_buffer_queue.qsize() == num_steps - startup_steps - 1 ), forward_send_recv_buffer_queue.qsize() input_tensor_grad = None for micro_step in range(steady_1f1b_steps): first_iter = micro_step == 0 last_iter = micro_step == (steady_1f1b_steps - 1) forward_micro_step_id = micro_step + startup_steps backward_micro_step_id = micro_step self._record_stamp("F", forward_micro_step_id, '"B"', forward=True) output_tensor = self._forward_step_helper( micro_dataset, forward_micro_step_id, check_is_last_chunk=True, ) self._record_stamp("F", forward_micro_step_id, '"E"', forward=True) if first_iter: for _ in range(self.num_stages - self.stage_id - 1): if self.user_hooks_enabled: self.bubble_hooks.run_hook() # NOTE: `send_forward_recv_backward` is intentionally unused to # prevent hanging bugs in dynamic shape mode. self._p2p_helper.send_forward( output_tensor, self.is_pipeline_last_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) output_tensor_grad = self._p2p_helper.recv_backward( self.is_pipeline_last_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) # Unlike normal FthenB, in 1F1B steps, we recv output_tensor_grad # for the current step, but not for the next step cur_backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id, forward=False ) self.output_tensor_grads[cur_backward_virtual_pp_rank].append( output_tensor_grad ) self._record_stamp( "B", backward_micro_step_id, '"B"', forward=False ) input_tensor_grad = self._backward_step_helper( backward_micro_step_id ) self._record_stamp( "B", backward_micro_step_id, '"E"', forward=False ) WeightGradStore.flush() # stash the input_tensor_grad and it will be sent to ths last stage later if self.is_pipeline_first_stage(ignore_virtual=True): backward_send_recv_buffer_queue.put(input_tensor_grad) if not last_iter: if not WeightGradStore.funcs_queue.empty(): # NOTE: `send_backward_recv_forward` is intentionally unused to # prevent hanging bugs in dynamic shape mode. input_tensor, fw_wait_handles = ( self._p2p_helper.recv_forward( self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) ) bw_wait_handles = self._p2p_helper.send_backward( input_tensor_grad, self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) # Execute weight grad computation while P2P communication is in progress WeightGradStore.pop() # Wait for P2P communication to complete if fw_wait_handles is not None: for fw_wait_handle in fw_wait_handles: fw_wait_handle.wait() if bw_wait_handles is not None: for bw_wait_handle in bw_wait_handles: bw_wait_handle.wait() next_forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id + 1, forward=True ) if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor = forward_send_recv_buffer_queue.get() self.input_tensors[next_forward_virtual_pp_rank].append( input_tensor ) else: # NOTE: `send_backward_recv_forward` is intentionally unused to # prevent hanging bugs in dynamic shape mode. input_tensor = self._p2p_helper.recv_forward( self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) self._p2p_helper.send_backward( input_tensor_grad, self.is_pipeline_first_stage(ignore_virtual=True), batch_p2p_comm=self._use_batch_p2p_comm, ) next_forward_virtual_pp_rank = self._get_virtual_pp_rank( forward_micro_step_id + 1, forward=True ) if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor = forward_send_recv_buffer_queue.get() self.input_tensors[next_forward_virtual_pp_rank].append( input_tensor ) else: for _ in range(self.num_stages - self.stage_id - 1): if self.user_hooks_enabled: self.bubble_hooks.run_hook() assert forward_send_recv_buffer_queue.qsize() == 0, ( forward_send_recv_buffer_queue.qsize() ) next_backward_virtual_pp_rank = self._get_virtual_pp_rank( steady_1f1b_steps, forward=False ) # no more fwd, but we need to send the input_tensor_grad. if self.is_pipeline_first_stage(ignore_virtual=True): input_tensor_grad = backward_send_recv_buffer_queue.get() if not WeightGradStore.funcs_queue.empty(): output_tensor_grad, wait_handles = ( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=True, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) ) # Execute weight grad computation while P2P communication is in progress WeightGradStore.pop() if wait_handles is not None: for handle in wait_handles: handle.wait() self.output_tensor_grads[next_backward_virtual_pp_rank].append( output_tensor_grad ) else: self.output_tensor_grads[next_backward_virtual_pp_rank].append( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=True, batch_p2p_comm=self._use_batch_p2p_comm, ) ) # Flush any remaining deferred weight gradient computations if not WeightGradStore.funcs_queue.empty(): raise AssertionError("WeightGradStore.funcs_queue should be empty") WeightGradStore.clear() # run cooldown for micro_step in range(cooldown_steps): backward_micro_step_id = micro_step + steady_1f1b_steps self._record_stamp( "B", backward_micro_step_id, '"B"', forward=False ) input_tensor_grad = self._backward_step_helper( backward_micro_step_id ) self._record_stamp( "B", backward_micro_step_id, '"E"', forward=False ) next_backward_virtual_pp_rank = self._get_virtual_pp_rank( backward_micro_step_id + 1, forward=False ) # Flush deferred weight gradient computations to queue WeightGradStore.flush() recv_next = True if backward_micro_step_id == (num_steps - 1): recv_next = False if self.is_pipeline_first_stage(ignore_virtual=True): if not self.is_pipeline_first_stage(): backward_send_recv_buffer_queue.put(input_tensor_grad) if ( self.is_pipeline_first_stage() and backward_micro_step_id % self.accumulate_steps >= skip_steps ): # no need to send the input_tensor_grad anymore input_tensor_grad = None else: input_tensor_grad = backward_send_recv_buffer_queue.get() if not WeightGradStore.funcs_queue.empty(): output_tensor_grad, wait_handles = ( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, overlap_p2p_comm=True, ) ) # Execute weight grad computation while P2P communication is in progress WeightGradStore.pop() if wait_handles is not None: for handle in wait_handles: handle.wait() self.output_tensor_grads[next_backward_virtual_pp_rank].append( output_tensor_grad ) else: self.output_tensor_grads[next_backward_virtual_pp_rank].append( self._p2p_helper.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, batch_p2p_comm=self._use_batch_p2p_comm, ) ) # Flush any remaining deferred weight gradient computations if not WeightGradStore.funcs_queue.empty(): raise AssertionError( "WeightGradStore.funcs_queue should be empty" ) WeightGradStore.clear() assert backward_send_recv_buffer_queue.empty(), ( "send_recv buffer should be empty" ) # Bubbles after cooldown for _ in range(self.stage_id): if self.user_hooks_enabled: self.bubble_hooks.run_hook() # reset dynamic meta counter if self._dynamic_shape: assert self._p2p_helper._dynamic_cnt == len( self._p2p_helper._send_recv_meta_list ), "p2p dynamic_cnt should equal to send_recv_meta_list" self._p2p_helper._dynamic_cnt = 0 self._flush_records() self._sync_overlap_grads() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").start() self._layers.allreduce_shared_weight_gradients() if self._enable_timer: self.timers("allreduce_shared_weight_gradients").stop() if compute_loss: # return loss if compute loss if self._enable_timer: self.timers("broadcast_final_loss").start() with paddle.amp.auto_cast(enable=False): train_loss_or_logits = self._broadcast_final_loss( return_micro_batch_loss ) if self._enable_timer: self.timers("broadcast_final_loss").stop() else: # else just return logits without loss func calc train_loss_or_logits = self.output_tensors.pop() if self._clear_every_step_cache: self._p2p_helper.clear_meta_cache() self.timer_printer() if self.processed_steps < g_profile_pipeline_details_steps: profile_pipeline_details( "[Pipeline details] End_forward_backward_step" ) self.processed_steps += 1 self._check_user_hooks_status_at_step_end() return train_loss_or_logits