chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,851 @@
|
||||
# Copyright (c) 2025 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
|
||||
# limitations under the License.
|
||||
|
||||
# The file has been adapted from DeepSeek DualPipe project
|
||||
# Copyright (c) 2025 DeepSeek
|
||||
# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import paddle
|
||||
from paddle import framework
|
||||
from paddle.distributed.communication.batch_isend_irecv import (
|
||||
P2POp,
|
||||
batch_isend_irecv,
|
||||
)
|
||||
|
||||
try:
|
||||
from paddle.distributed.communication import deep_ep
|
||||
except ImportError:
|
||||
deep_ep = None
|
||||
|
||||
from ..utils.log_util import logger
|
||||
from .pipeline_parallel import (
|
||||
FakeMicroDataset,
|
||||
HybridParallelOptimizer,
|
||||
PipelineDatasetPreprocessor,
|
||||
PipelineParallel,
|
||||
)
|
||||
from .pp_utils.batch_comm_helper import BatchCommHelper
|
||||
from .pp_utils.forward_backward_overlap_utils import ScheduleChunk
|
||||
from .zero_bubble_utils import EventStore, WeightGradStore
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def detach_and_requires_grad(x):
|
||||
o = x.detach()
|
||||
o.stop_gradient = False
|
||||
return o
|
||||
|
||||
|
||||
class DualPipeVParallel(PipelineParallel):
|
||||
"""
|
||||
An implementation of the DualPipeV, based on
|
||||
https://github.com/deepseek-ai/DualPipe/blob/main/dualpipe/dualpipe.py.
|
||||
"""
|
||||
|
||||
def __init__(self, layers, hcg, strategy):
|
||||
super().__init__(layers=layers, hcg=hcg, strategy=strategy)
|
||||
self.overlapped_forward_backward = hasattr(
|
||||
type(self._layers), "overlapped_forward_backward"
|
||||
)
|
||||
logger.info(
|
||||
f"Using DualPipeVParallel with overlapping forward backward={self.overlapped_forward_backward}"
|
||||
)
|
||||
|
||||
self.num_ranks = self.num_stages
|
||||
self.group_rank = self.pp_group.rank
|
||||
self.prev_rank = self.pp_group.ranks[
|
||||
(self.group_rank - 1) % self.pp_group.world_size
|
||||
]
|
||||
self.next_rank = self.pp_group.ranks[
|
||||
(self.group_rank + 1) % self.pp_group.world_size
|
||||
]
|
||||
|
||||
# NOTE(zhangyuqin1998): The first rank has to broadcast the meta information
|
||||
# of the P2P communication after the first forward.
|
||||
self.need_broadcast_meta = self.is_pipeline_first_stage()
|
||||
self.need_recv_meta = not self.is_pipeline_first_stage()
|
||||
self._p2p_helper = BatchCommHelper(self._using_cache)
|
||||
|
||||
def is_pipeline_first_stage(self):
|
||||
return self.group_rank == 0
|
||||
|
||||
def is_pipeline_last_stage(self):
|
||||
return self.group_rank == self.num_ranks - 1
|
||||
|
||||
def _reset_states(self):
|
||||
self.input_tensors = ([], [])
|
||||
self.output_tensors = ([], [])
|
||||
self.input_grad_tensors = ([], [])
|
||||
self.output_grad_tensors = ([], [])
|
||||
self.loss_tensors: list[paddle.Tensor] = []
|
||||
self.schedule_chunks = ([], [])
|
||||
self.loss_fn_chunks = []
|
||||
|
||||
# The first value in the list corresponds to phase 0, and the second value corresponds to phase 1.
|
||||
self.current_f_acc_id = [0, 0]
|
||||
self.current_b_acc_id = [0, 0]
|
||||
self.current_send_f_acc_id = [0, 0]
|
||||
self.current_send_b_acc_id = [0, 0]
|
||||
self.current_recv_f_acc_id = [0, 0]
|
||||
self.current_recv_b_acc_id = [0, 0]
|
||||
self.comm_forward_ops: list[P2POp] = []
|
||||
self.comm_backward_ops: list[P2POp] = []
|
||||
self.to_free: list[paddle.Tensor] = []
|
||||
|
||||
def _get_forward_inputs(self, micro_datasets, phase, acc_id):
|
||||
is_first_stage = self.is_pipeline_first_stage() and phase == 0
|
||||
if is_first_stage:
|
||||
assert micro_datasets is not None
|
||||
self.input_tensors[phase].append(next(micro_datasets[phase])[0])
|
||||
if self.forward_only:
|
||||
self.input_tensors[phase][acc_id] = None
|
||||
return self.input_tensors[phase][acc_id]
|
||||
|
||||
def _get_forward_labels(self, micro_datasets, phase, acc_id):
|
||||
is_last_stage = self.is_pipeline_first_stage() and phase == 1
|
||||
if is_last_stage and self._compute_loss:
|
||||
assert micro_datasets is not None
|
||||
labels = next(micro_datasets[phase])[1]
|
||||
self._check_micro_batch_data_valid(labels)
|
||||
return labels
|
||||
else:
|
||||
return None
|
||||
|
||||
def _loss_compute(self, micro_datasets, phase, acc_id, logits):
|
||||
labels = self._get_forward_labels(micro_datasets, phase, acc_id)
|
||||
loss_fn_node = None
|
||||
if not self.overlapped_forward_backward:
|
||||
loss_tensor = self._layers._loss_fn[0](logits, labels)
|
||||
else:
|
||||
loss_fn_node = self._layers._loss_fn[0].build_schedule_node()
|
||||
loss_fn_node.labels = labels
|
||||
loss_tensor = loss_fn_node.forward(logits)
|
||||
self._store_forward_loss(phase, loss_tensor, loss_fn_node)
|
||||
|
||||
def _store_forward_tensors(self, phase, outputs, schedule_chunk):
|
||||
self.schedule_chunks[phase].append(schedule_chunk)
|
||||
if self.is_pipeline_last_stage() and phase == 0:
|
||||
self.input_tensors[1].append(
|
||||
[detach_and_requires_grad(output) for output in outputs]
|
||||
)
|
||||
is_last_stage = self.is_pipeline_first_stage() and phase == 1
|
||||
if not is_last_stage:
|
||||
self.output_tensors[phase].append(outputs)
|
||||
|
||||
def _forward_compute(self, phase: int, micro_datasets=None) -> None:
|
||||
acc_id = self.current_f_acc_id[phase]
|
||||
self.current_f_acc_id[phase] += 1
|
||||
|
||||
inputs = self._get_forward_inputs(micro_datasets, phase, acc_id)
|
||||
|
||||
if self.overlapped_forward_backward:
|
||||
schedule_chunk = self._layers.get_schedule_chunk(chunk_id=phase)
|
||||
outputs = schedule_chunk.forward(inputs)
|
||||
else:
|
||||
schedule_chunk = None
|
||||
outputs = self._layers.forward(inputs, chunk_id=phase)
|
||||
outputs = [outputs] if isinstance(outputs, paddle.Tensor) else outputs
|
||||
|
||||
is_last_stage = self.is_pipeline_first_stage() and phase == 1
|
||||
if is_last_stage and self._compute_loss:
|
||||
self._loss_compute(micro_datasets, phase, acc_id, outputs)
|
||||
self._store_forward_tensors(phase, outputs, schedule_chunk)
|
||||
|
||||
def _get_backward_inputs(self, phase, acc_id):
|
||||
outputs = self.output_tensors[phase][acc_id]
|
||||
self.output_tensors[phase][acc_id] = None
|
||||
output_grads = self.output_grad_tensors[phase][acc_id]
|
||||
self.output_grad_tensors[phase][acc_id] = None
|
||||
non_empty = [
|
||||
(t, g) for t, g in zip(outputs, output_grads) if g is not None
|
||||
]
|
||||
outputs, output_grads = list(zip(*non_empty))
|
||||
return outputs, output_grads
|
||||
|
||||
def _store_backward_tensors(self, phase, acc_id, input_grads=None):
|
||||
if input_grads is None:
|
||||
inputs = self.input_tensors[phase][acc_id]
|
||||
input_grads = [
|
||||
t.grad
|
||||
for t in inputs
|
||||
if (t is not None and not t.stop_gradient)
|
||||
]
|
||||
self.input_tensors[phase][acc_id] = None
|
||||
|
||||
if isinstance(input_grads, paddle.Tensor):
|
||||
input_grads = (input_grads,)
|
||||
if self.is_pipeline_last_stage() and phase == 1:
|
||||
self.output_grad_tensors[0].append(input_grads)
|
||||
else:
|
||||
self.input_grad_tensors[phase].append(input_grads)
|
||||
|
||||
def _store_forward_loss(self, phase, loss_tensor, loss_fn_node=None):
|
||||
is_last_stage = self.is_pipeline_first_stage() and phase == 1
|
||||
if is_last_stage and self._compute_loss:
|
||||
if isinstance(loss_tensor, (tuple, list)):
|
||||
assert len(loss_tensor) == 1
|
||||
loss_tensor = loss_tensor[0]
|
||||
assert isinstance(loss_tensor, paddle.Tensor), (
|
||||
"Currently, loss_fn should obtain Paddle.Tensor dtype"
|
||||
)
|
||||
|
||||
self.loss_tensors.append(loss_tensor)
|
||||
self.loss_fn_chunks.append(loss_fn_node)
|
||||
|
||||
def _backward_compute(self, phase: int, enable_zb: bool = False) -> None:
|
||||
if self.forward_only:
|
||||
return
|
||||
|
||||
acc_id = self.current_b_acc_id[phase]
|
||||
self.current_b_acc_id[phase] += 1
|
||||
|
||||
is_last_stage = self.is_pipeline_first_stage() and phase == 1
|
||||
|
||||
WeightGradStore.enabled = enable_zb
|
||||
input_grads = None
|
||||
with paddle.amp.auto_cast(enable=False):
|
||||
if is_last_stage:
|
||||
loss = self.loss_tensors[acc_id]
|
||||
if self.overlapped_forward_backward:
|
||||
loss_fn_node = self.loss_fn_chunks[acc_id]
|
||||
backward_chunk = self.schedule_chunks[phase][acc_id]
|
||||
_, _, input_grads = (
|
||||
self._layers.overlapped_forward_backward(
|
||||
ScheduleChunk([]), # forward_chunk
|
||||
None, # forward_inputs
|
||||
None, # forward_loss_fn_node
|
||||
backward_chunk,
|
||||
loss_fn_node,
|
||||
None, # input_grads
|
||||
self.scaler,
|
||||
combine_bw_event_to_wait=None,
|
||||
pp_stream=None,
|
||||
)
|
||||
)
|
||||
self.loss_fn_chunks[acc_id] = None
|
||||
self.schedule_chunks[phase][acc_id] = None
|
||||
else:
|
||||
if self.scaler:
|
||||
paddle.autograd.backward(self.scaler.scale(loss))
|
||||
else:
|
||||
paddle.autograd.backward(loss)
|
||||
else:
|
||||
outputs, output_grads = self._get_backward_inputs(phase, acc_id)
|
||||
if self.overlapped_forward_backward:
|
||||
backward_chunk = self.schedule_chunks[phase][acc_id]
|
||||
_, _, input_grads = (
|
||||
self._layers.overlapped_forward_backward(
|
||||
ScheduleChunk([]), # forward_chunk
|
||||
None, # forward_inputs
|
||||
None, # forward_loss_fn_node
|
||||
backward_chunk,
|
||||
None, # backward_loss_fn_node
|
||||
output_grads,
|
||||
None, # scaler
|
||||
combine_bw_event_to_wait=None,
|
||||
pp_stream=None,
|
||||
)
|
||||
)
|
||||
self.schedule_chunks[phase][acc_id] = None
|
||||
else:
|
||||
if len(outputs) > 0:
|
||||
outputs = [t for t in outputs if not t.stop_gradient]
|
||||
paddle.autograd.backward(
|
||||
tensors=outputs,
|
||||
grad_tensors=output_grads,
|
||||
)
|
||||
WeightGradStore.enabled = False
|
||||
if enable_zb:
|
||||
WeightGradStore.flush()
|
||||
|
||||
self._store_backward_tensors(phase, acc_id, input_grads=input_grads)
|
||||
|
||||
def _forward_backward_compute(
|
||||
self,
|
||||
forward_phase: int,
|
||||
backward_phase: int,
|
||||
micro_datasets=None,
|
||||
combine_backward_event_to_wait=None,
|
||||
pass_pp_stream=False,
|
||||
) -> None:
|
||||
if self.forward_only:
|
||||
self._forward_compute(forward_phase, micro_datasets)
|
||||
return
|
||||
|
||||
if not self.overlapped_forward_backward:
|
||||
self._forward_compute(forward_phase, micro_datasets)
|
||||
self._backward_compute(backward_phase)
|
||||
return
|
||||
|
||||
# pre-forward
|
||||
forward_acc_id = self.current_f_acc_id[forward_phase]
|
||||
self.current_f_acc_id[forward_phase] += 1
|
||||
|
||||
forward_inputs = self._get_forward_inputs(
|
||||
micro_datasets, forward_phase, forward_acc_id
|
||||
)
|
||||
forward_labels = self._get_forward_labels(
|
||||
micro_datasets, forward_phase, forward_acc_id
|
||||
)
|
||||
if forward_labels is not None:
|
||||
forward_loss_fn_node = self._layers._loss_fn[
|
||||
0
|
||||
].build_schedule_node()
|
||||
forward_loss_fn_node.labels = forward_labels
|
||||
else:
|
||||
forward_loss_fn_node = None
|
||||
|
||||
# pre-backward
|
||||
backward_acc_id = self.current_b_acc_id[backward_phase]
|
||||
self.current_b_acc_id[backward_phase] += 1
|
||||
|
||||
is_last_stage1 = self.is_pipeline_first_stage() and backward_phase == 1
|
||||
if is_last_stage1:
|
||||
backward_loss_fn_node = self.loss_fn_chunks[backward_acc_id]
|
||||
backward_grads = None
|
||||
else:
|
||||
backward_loss_fn_node = None
|
||||
_, backward_grads = self._get_backward_inputs(
|
||||
backward_phase, backward_acc_id
|
||||
)
|
||||
|
||||
# event_to_wait = deep_ep.get_event_from_custom_stream(paddle.device.current_stream().stream_base)
|
||||
|
||||
# forward & backward
|
||||
forward_chunk = self._layers.get_schedule_chunk(chunk_id=forward_phase)
|
||||
backward_chunk = self.schedule_chunks[backward_phase][backward_acc_id]
|
||||
forward_outputs, forward_loss, backward_input_grads = (
|
||||
self._layers.overlapped_forward_backward(
|
||||
forward_chunk,
|
||||
forward_inputs,
|
||||
forward_loss_fn_node,
|
||||
backward_chunk,
|
||||
backward_loss_fn_node,
|
||||
backward_grads,
|
||||
self.scaler,
|
||||
combine_bw_event_to_wait=combine_backward_event_to_wait,
|
||||
pp_stream=(
|
||||
self.pp_group.process_group.get_stream(
|
||||
paddle.framework._current_expected_place_()
|
||||
)
|
||||
if pass_pp_stream
|
||||
else None
|
||||
),
|
||||
)
|
||||
)
|
||||
self.schedule_chunks[backward_phase][backward_acc_id] = None
|
||||
|
||||
# post-forward
|
||||
self._store_forward_tensors(
|
||||
forward_phase, forward_outputs, forward_chunk
|
||||
)
|
||||
self._store_forward_loss(
|
||||
forward_phase, forward_loss, forward_loss_fn_node
|
||||
)
|
||||
|
||||
# post-backward
|
||||
self._store_backward_tensors(
|
||||
backward_phase, backward_acc_id, input_grads=backward_input_grads
|
||||
)
|
||||
|
||||
def _commit_and_wait_comm(
|
||||
self, p2p_overlap=False, use_outer_event_wait=False
|
||||
) -> None:
|
||||
common_forward_ops_num = (
|
||||
len(self.comm_forward_ops)
|
||||
if self.comm_forward_ops is not None
|
||||
else 0
|
||||
)
|
||||
common_backward_ops_num = (
|
||||
len(self.comm_backward_ops)
|
||||
if self.comm_backward_ops is not None
|
||||
else 0
|
||||
)
|
||||
if common_forward_ops_num == 0 and common_backward_ops_num == 0:
|
||||
if EventStore.event is not None:
|
||||
e_t = EventStore.event
|
||||
EventStore.event = None
|
||||
return e_t
|
||||
return deep_ep.get_event_from_custom_stream(
|
||||
paddle.device.current_stream().stream_base
|
||||
)
|
||||
|
||||
use_stream_wait_event = (
|
||||
p2p_overlap and self._overlap_p2p_comm and deep_ep is not None
|
||||
)
|
||||
|
||||
pp_raw_stream = self.pp_group.process_group.get_stream(
|
||||
paddle.framework._current_expected_place_()
|
||||
)
|
||||
if use_outer_event_wait:
|
||||
self.pp_group.process_group.set_outer_wait(True)
|
||||
|
||||
if common_forward_ops_num > 0:
|
||||
fwd_reqs = batch_isend_irecv(self.comm_forward_ops)
|
||||
|
||||
if not use_stream_wait_event:
|
||||
for req in fwd_reqs:
|
||||
req.wait()
|
||||
|
||||
if use_outer_event_wait:
|
||||
self.pp_group.process_group.set_outer_wait(False)
|
||||
|
||||
if use_stream_wait_event:
|
||||
forward_event_to_wait = deep_ep.get_event_from_custom_stream(
|
||||
pp_raw_stream
|
||||
)
|
||||
|
||||
backward_outer_event_wait = False
|
||||
if EventStore.event is not None:
|
||||
with paddle.device.stream_guard(
|
||||
paddle.device.Stream(stream_base=pp_raw_stream)
|
||||
):
|
||||
EventStore.event.current_stream_wait()
|
||||
|
||||
EventStore.set(None)
|
||||
self.pp_group.process_group.set_outer_wait(True)
|
||||
|
||||
backward_outer_event_wait = True
|
||||
|
||||
if common_backward_ops_num > 0:
|
||||
bwd_reqs = batch_isend_irecv(self.comm_backward_ops)
|
||||
|
||||
if not use_stream_wait_event:
|
||||
for req in bwd_reqs:
|
||||
req.wait()
|
||||
|
||||
if backward_outer_event_wait:
|
||||
self.pp_group.process_group.set_outer_wait(False)
|
||||
|
||||
if use_stream_wait_event:
|
||||
forward_event_to_wait.current_stream_wait()
|
||||
|
||||
combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
|
||||
pp_raw_stream
|
||||
)
|
||||
else:
|
||||
combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
|
||||
paddle.device.current_stream().stream_base
|
||||
)
|
||||
|
||||
self.comm_forward_ops = []
|
||||
self.comm_backward_ops = []
|
||||
|
||||
self._free_tensors()
|
||||
|
||||
return combine_bw_event_to_wait
|
||||
|
||||
def _weight_pass(self) -> None:
|
||||
if self.forward_only:
|
||||
return
|
||||
|
||||
self._commit_and_wait_comm()
|
||||
|
||||
# Assume FIFO
|
||||
WeightGradStore.pop()
|
||||
|
||||
def _free_tensors(self) -> None:
|
||||
self._release_output(self.to_free)
|
||||
self.to_free = []
|
||||
|
||||
def _recv_forward(self, phase: int) -> None:
|
||||
if (self.is_pipeline_first_stage() and phase == 0) or (
|
||||
self.is_pipeline_last_stage() and phase == 1
|
||||
):
|
||||
return
|
||||
|
||||
self.current_recv_f_acc_id[phase] += 1
|
||||
|
||||
tensors = self._p2p_helper.append_irecv(
|
||||
self.comm_forward_ops,
|
||||
self.prev_rank if phase == 0 else self.next_rank,
|
||||
self.pp_group,
|
||||
alloc_on_comm_stream=self._overlap_p2p_comm,
|
||||
)
|
||||
self.input_tensors[phase].append(tensors)
|
||||
|
||||
def _send_forward(self, phase: int) -> None:
|
||||
if (self.is_pipeline_first_stage() and phase == 1) or (
|
||||
self.is_pipeline_last_stage() and phase == 0
|
||||
):
|
||||
return
|
||||
|
||||
acc_id = self.current_send_f_acc_id[phase]
|
||||
self.current_send_f_acc_id[phase] += 1
|
||||
tensors = self.output_tensors[phase][acc_id]
|
||||
|
||||
self._p2p_helper.append_isend(
|
||||
self.comm_forward_ops,
|
||||
tensors,
|
||||
self.next_rank if phase == 0 else self.prev_rank,
|
||||
self.pp_group,
|
||||
self.need_broadcast_meta,
|
||||
)
|
||||
self.need_broadcast_meta = False
|
||||
|
||||
self.to_free.extend(tensors)
|
||||
|
||||
def _recv_backward(self, phase: int) -> None:
|
||||
if self.forward_only:
|
||||
return
|
||||
|
||||
if (self.is_pipeline_first_stage() and phase == 1) or (
|
||||
self.is_pipeline_last_stage() and phase == 0
|
||||
):
|
||||
return
|
||||
|
||||
self.current_recv_b_acc_id[phase] += 1
|
||||
tensors = self._p2p_helper.append_irecv(
|
||||
self.comm_backward_ops,
|
||||
self.next_rank if phase == 0 else self.prev_rank,
|
||||
self.pp_group,
|
||||
alloc_on_comm_stream=self._overlap_p2p_comm,
|
||||
)
|
||||
self.output_grad_tensors[phase].append(tensors)
|
||||
|
||||
def _send_backward(self, phase: int) -> None:
|
||||
if self.forward_only:
|
||||
return
|
||||
|
||||
if (self.is_pipeline_first_stage() and phase == 0) or (
|
||||
self.is_pipeline_last_stage() and phase == 1
|
||||
):
|
||||
return
|
||||
|
||||
acc_id = self.current_send_b_acc_id[phase]
|
||||
self.current_send_b_acc_id[phase] += 1
|
||||
tensors = self.input_grad_tensors[phase][acc_id]
|
||||
self.input_grad_tensors[phase][acc_id] = None
|
||||
|
||||
self._p2p_helper.append_isend(
|
||||
self.comm_backward_ops,
|
||||
tensors,
|
||||
self.prev_rank if phase == 0 else self.next_rank,
|
||||
self.pp_group,
|
||||
)
|
||||
|
||||
def _forward_pass(
|
||||
self,
|
||||
phase: int,
|
||||
micro_datasets=None,
|
||||
recv: bool = True,
|
||||
send: bool = True,
|
||||
) -> None:
|
||||
if recv:
|
||||
self._recv_forward(phase)
|
||||
self._commit_and_wait_comm()
|
||||
|
||||
self._forward_compute(phase, micro_datasets)
|
||||
|
||||
if send:
|
||||
self._send_forward(phase)
|
||||
|
||||
def _backward_pass(
|
||||
self,
|
||||
phase: int,
|
||||
enable_zb: bool = False,
|
||||
recv: bool = True,
|
||||
send: bool = True,
|
||||
) -> None:
|
||||
if recv:
|
||||
self._recv_backward(phase)
|
||||
self._commit_and_wait_comm()
|
||||
|
||||
self._backward_compute(phase, enable_zb)
|
||||
|
||||
if send:
|
||||
self._send_backward(phase)
|
||||
|
||||
def _forward_backward_pass(
|
||||
self,
|
||||
forward_phase: int,
|
||||
backward_phase: int,
|
||||
micro_datasets=None,
|
||||
recv0: bool = True,
|
||||
first_chunk=False,
|
||||
last_chunk=False,
|
||||
main_stage=False,
|
||||
last_stage_and_first_chunk=False,
|
||||
) -> None:
|
||||
if recv0:
|
||||
self._recv_forward(forward_phase)
|
||||
self._recv_backward(backward_phase)
|
||||
|
||||
need_send_forward = not (
|
||||
self.is_pipeline_first_stage() and forward_phase == 1
|
||||
) or (self.is_pipeline_last_stage() and forward_phase == 0)
|
||||
need_send_backward = not (
|
||||
self.is_pipeline_first_stage() and backward_phase == 0
|
||||
) or (self.is_pipeline_last_stage() and backward_phase == 1)
|
||||
|
||||
use_outer_event_wait = (
|
||||
main_stage
|
||||
and not first_chunk
|
||||
and self._overlap_p2p_comm
|
||||
and deep_ep is not None
|
||||
and (need_send_forward and need_send_backward)
|
||||
)
|
||||
|
||||
pass_pp_stream = (
|
||||
main_stage
|
||||
and not last_chunk
|
||||
and self._overlap_p2p_comm
|
||||
and deep_ep is not None
|
||||
and (need_send_forward and need_send_backward)
|
||||
and (not last_stage_and_first_chunk)
|
||||
)
|
||||
|
||||
combine_bw_wait_event = self._commit_and_wait_comm(
|
||||
not last_chunk, use_outer_event_wait
|
||||
)
|
||||
|
||||
self._forward_backward_compute(
|
||||
forward_phase,
|
||||
backward_phase,
|
||||
micro_datasets,
|
||||
combine_backward_event_to_wait=combine_bw_wait_event,
|
||||
pass_pp_stream=pass_pp_stream,
|
||||
)
|
||||
|
||||
self._send_forward(forward_phase)
|
||||
self._send_backward(backward_phase)
|
||||
|
||||
def _wrap_data(self, data, phase):
|
||||
"""
|
||||
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() and phase == 0,
|
||||
self.is_pipeline_first_stage() and phase == 1,
|
||||
self.accumulate_steps,
|
||||
self.micro_batch_size,
|
||||
)
|
||||
return micro_dataset
|
||||
|
||||
def _prepare_training(self, data, optimizer, lr_scheduler):
|
||||
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():
|
||||
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 _broadcast_final_loss(self):
|
||||
loss_sum_tensor = paddle.zeros([1], "float32")
|
||||
if self.is_pipeline_first_stage():
|
||||
assert len(self.loss_tensors) > 0, (
|
||||
"train_batch() in last stage should obtain valid loss"
|
||||
)
|
||||
for loss in self.loss_tensors:
|
||||
loss_sum_tensor += loss.detach().astype("float32")
|
||||
loss_sum_tensor /= self.accumulate_steps
|
||||
|
||||
paddle.distributed.all_reduce(
|
||||
loss_sum_tensor, group=self.pp_group, sync_op=True
|
||||
)
|
||||
return loss_sum_tensor
|
||||
|
||||
def forward_backward_pipeline(
|
||||
self,
|
||||
data,
|
||||
scaler,
|
||||
forward_only=False,
|
||||
compute_loss=True,
|
||||
):
|
||||
self.scaler = scaler
|
||||
|
||||
rank = self.group_rank
|
||||
num_ranks = self.num_ranks
|
||||
assert (
|
||||
self.accumulate_steps > 0 and self.accumulate_steps >= num_ranks * 2
|
||||
), f"{self.accumulate_steps=}, {num_ranks=}"
|
||||
self.forward_only = forward_only
|
||||
|
||||
self._reset_states()
|
||||
|
||||
# NOTE(zhangyuqin1998): Tensors to be sent or received must have a
|
||||
# consistent shape and data type throughout the entire pipeline. We
|
||||
# broadcast the meta info in the first forward of the first rank.
|
||||
self._p2p_helper.recv_meta_from_head(self.pp_group, self.need_recv_meta)
|
||||
self.need_recv_meta = False
|
||||
|
||||
micro_dataset_phase0 = self._wrap_data(data, 0)
|
||||
micro_dataset_phase1 = self._wrap_data(data, 1)
|
||||
micro_datasets = [micro_dataset_phase0, micro_dataset_phase1]
|
||||
|
||||
# Step 1: nF0
|
||||
step_1 = (num_ranks - rank - 1) * 2
|
||||
for i in range(step_1):
|
||||
self._forward_pass(0, micro_datasets)
|
||||
|
||||
# Step 2: nF0F1
|
||||
step_2 = rank + 1
|
||||
self._recv_forward(0)
|
||||
for i in range(step_2):
|
||||
self._forward_pass(0, micro_datasets, recv=False, send=False)
|
||||
self._recv_forward(0)
|
||||
self._forward_pass(
|
||||
1,
|
||||
micro_datasets,
|
||||
send=(not self.is_pipeline_last_stage()) or (i < step_2 - 1),
|
||||
)
|
||||
self._send_forward(0)
|
||||
|
||||
# Step 3: nB1W1F1 (Use zero bubble)
|
||||
step_3 = num_ranks - rank - 1
|
||||
for i in range(step_3):
|
||||
self._backward_pass(1, enable_zb=True)
|
||||
self._recv_forward(1)
|
||||
self._weight_pass()
|
||||
self._forward_pass(1, micro_datasets, recv=False)
|
||||
|
||||
# Step 4 (Main step): nF0B1F1B0
|
||||
step_4 = self.accumulate_steps - num_ranks * 2 + rank + 1
|
||||
have_step5 = num_ranks - rank - 1 > 0
|
||||
# Update code to support send/recv overlap
|
||||
# Only support send/recv overlap in MainStep
|
||||
for i in range(step_4):
|
||||
is_last_chunk = i + 1 == step_4
|
||||
if i == 0:
|
||||
if self.is_pipeline_last_stage():
|
||||
# NOTE: We don't overlap these two passes to further reduce bubble size.
|
||||
self._forward_pass(
|
||||
0, micro_datasets, recv=False, send=False
|
||||
)
|
||||
self._send_forward(1)
|
||||
self._backward_pass(1, send=False)
|
||||
self._send_forward(0)
|
||||
self._send_backward(1)
|
||||
|
||||
self._forward_backward_pass(
|
||||
1,
|
||||
0,
|
||||
micro_datasets,
|
||||
first_chunk=True,
|
||||
last_chunk=is_last_chunk,
|
||||
main_stage=True,
|
||||
)
|
||||
else:
|
||||
self._forward_backward_pass(
|
||||
0,
|
||||
1,
|
||||
micro_datasets,
|
||||
recv0=False,
|
||||
first_chunk=True,
|
||||
main_stage=True,
|
||||
)
|
||||
|
||||
self._forward_backward_pass(
|
||||
1,
|
||||
0,
|
||||
micro_datasets,
|
||||
last_chunk=is_last_chunk,
|
||||
main_stage=True,
|
||||
)
|
||||
else:
|
||||
self._forward_backward_pass(
|
||||
0,
|
||||
1,
|
||||
micro_datasets,
|
||||
main_stage=True,
|
||||
last_stage_and_first_chunk=self.is_pipeline_last_stage(),
|
||||
)
|
||||
self._forward_backward_pass(
|
||||
1,
|
||||
0,
|
||||
micro_datasets,
|
||||
last_chunk=is_last_chunk,
|
||||
main_stage=True,
|
||||
)
|
||||
|
||||
# Step 5: nB1F1B0
|
||||
step_5 = num_ranks - rank - 1
|
||||
for i in range(step_5):
|
||||
self._backward_pass(1)
|
||||
self._forward_backward_pass(1, 0, micro_datasets)
|
||||
|
||||
# Step 6: nB1B0 (The second half of the passes use zero bubble)
|
||||
step_6 = rank + 1
|
||||
enable_zb = False
|
||||
for i in range(step_6):
|
||||
if i == step_6 // 2 and rank % 2 == 1:
|
||||
enable_zb = True
|
||||
self._backward_pass(1, enable_zb=enable_zb)
|
||||
if i == step_6 // 2 and rank % 2 == 0:
|
||||
enable_zb = True
|
||||
self._backward_pass(0, enable_zb=enable_zb)
|
||||
|
||||
# Step 7: nWB0 (Use zero bubble)
|
||||
step_7 = num_ranks - rank - 1
|
||||
for i in range(step_7):
|
||||
self._weight_pass()
|
||||
self._backward_pass(0, enable_zb=True)
|
||||
|
||||
# Step 8: nW
|
||||
step_8 = rank + 1
|
||||
for i in range(step_8):
|
||||
self._weight_pass()
|
||||
assert WeightGradStore.funcs_queue.empty()
|
||||
|
||||
self._commit_and_wait_comm()
|
||||
|
||||
self._layers.allreduce_shared_weight_gradients()
|
||||
|
||||
with paddle.amp.auto_cast(enable=False):
|
||||
train_loss = self._broadcast_final_loss()
|
||||
|
||||
self._reset_states()
|
||||
return train_loss
|
||||
|
||||
def train_batch(
|
||||
self,
|
||||
data,
|
||||
optimizer,
|
||||
lr_scheduler=None,
|
||||
scaler=None,
|
||||
):
|
||||
data = self._prepare_training(data, optimizer, lr_scheduler)
|
||||
|
||||
train_loss = self.forward_backward_pipeline(data, scaler)
|
||||
|
||||
# optimizer
|
||||
with paddle.amp.auto_cast(enable=False):
|
||||
self._optimizer_step()
|
||||
|
||||
return train_loss
|
||||
Reference in New Issue
Block a user