wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
paddlepaddle--paddle/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

4193 lines
161 KiB
Python
Raw Permalink Blame History

# 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
Reference in New Issue View Git Blame Copy Permalink