# 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 import time import copy import os from types import MethodType from numpy import prod import paddle import paddle.fluid as fluid from .meta_parallel_base import MetaParallelBase from .pp_utils.utils import get_tensor_bytes, is_float_tensor from .pp_utils import utils from .parallel_layers.pp_layers import PipelineLayer from ..utils.hybrid_parallel_util import * from ..utils.log_util import logger class PipelineParallel(MetaParallelBase): def __init__(self, layers, hcg, strategy): super(PipelineParallel, self).__init__(layers, hcg, strategy) self.use_pipe_parallel = self._hcg.get_pipe_parallel_world_size() > 1 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.num_caches = 0 self.caches = { 'inputs': [], 'labels': [], 'outputs': [], } self.recv_cache = None self.grad_tensors = None self.send_meta = True self.current_loss = paddle.to_tensor(0.0) self.total_loss = None self.use_amp = self._strategy.amp self.init_loss_scaling = self._strategy.amp_configs['init_loss_scaling'] self.micro_batch_size = self._strategy.pipeline_configs[ 'micro_batch_size'] self.accumulate_steps = self._strategy.pipeline_configs[ 'accumulate_steps'] self.num_stages = self._hcg.get_pipe_parallel_world_size() self.stage_id = self._hcg.get_stage_id() self.prev_stage_id = self.stage_id - 1 self.next_stage_id = self.stage_id + 1 self.pp_group = self._hcg.get_pipe_parallel_group() logger.info("Pipeline Info -- num_stages: {}, stage_id: {}".format( self.num_stages, self.stage_id)) if self.use_model_parallel: logger.info("start broadcast mp parameters") broadcast_mp_parameters(self._layers, self._hcg) if self.use_data_parallel: logger.info("start broadcast mp parameters") broadcast_dp_parameters(self._layers, self._hcg) def _allocate_caches(self, num_caches): if self.num_caches >= num_caches: return num = num_caches - self.num_caches self.num_caches = num_caches for key in self.caches: self.caches[key].extend([None] * num) def train_batch(self, data, optimizer): self.optimizer = optimizer assert fluid.framework._dygraph_tracer()._has_grad, ( 'Please enable the generation of gradients.') if self.stage_id == 0 or self.stage_id == self.num_stages - 1: assert data, ( "For the first and the last stage, the data_iter must be set.") else: assert data is None, ( "For pipe stages other than the first and the last one, " "the data_iter must be None.") self.data = data self._layers.train() self.total_loss = None minibatch_cmds = utils.TrainGenerator(self.accumulate_steps, self.num_stages, self.stage_id) self._train(minibatch_cmds) return self.total_loss def _train(self, minibatch_cmds): self._allocate_caches(self.accumulate_steps) for micro_cmds in minibatch_cmds: for cmd in micro_cmds: assert type(cmd) in self._COMMAND_MAP, "unknow cmd: {}".format( type(cmd)) self._apply_cmd = MethodType(self._COMMAND_MAP[type(cmd)], self) self._apply_cmd(**cmd.kwargs) def _allreduce_grads(self): if not self.use_data_parallel: return fused_allreduce_gradients(list(self._layers.parameters()), self._hcg) def _forward(self, cache_id): # load data self._load_micro_batch(cache_id) if self.stage_id != 0: self._recv_activations(cache_id) if isinstance(self.caches['inputs'][cache_id], tuple): inputs = tuple(t for t in self.caches['inputs'][cache_id]) else: inputs = self.caches['inputs'][cache_id] self._clear_grads(inputs) outputs = self._layers.forward(inputs) self.caches['outputs'][cache_id] = outputs if self.stage_id == self.num_stages - 1: if self._layers._loss_fn is not None: labels = self.caches['labels'][cache_id] outputs = self._layers._loss_fn(outputs, labels) if self.stage_id == self.num_stages - 1: self.current_loss = outputs if isinstance(self.current_loss, paddle.Tensor): if self.total_loss is None: self.total_loss = paddle.zeros_like(self.current_loss) self.total_loss += self.current_loss.detach() else: if self.total_loss is None: self.total_loss = [ paddle.zeros_like(v) for v in self.current_loss ] for idx, v in enumerate(self.current_loss): self.total_loss[idx] += v.detach() if self.use_data_parallel: self.current_loss = self.current_loss / self._hcg.get_data_parallel_world_size( ) if self.accumulate_steps > 1: self.current_loss = self.current_loss / self.accumulate_steps self.caches['outputs'][cache_id] = self.current_loss.clone() else: self._send_activations(cache_id) def _backward(self, cache_id): assert self.optimizer is not None if self.stage_id == self.num_stages - 1: paddle.autograd.backward(self.caches['outputs'][cache_id]) self._send_gradients(cache_id) return self._recv_gradients(cache_id) outputs = self.caches['outputs'][cache_id] grad_tensors = self.grad_tensors if isinstance(outputs, tuple): out_tensors = [t for t in outputs if is_float_tensor(t)] assert len(out_tensors) == len(grad_tensors) paddle.autograd.backward( tensors=out_tensors, grad_tensors=grad_tensors) else: paddle.autograd.backward( tensors=[outputs], grad_tensors=[grad_tensors]) grad_tensors = None if self.stage_id != 0: self._send_gradients(cache_id) self.caches['outputs'][cache_id] = None #self.caches['backward_tensors'][cache_id] = None def _get_data(self): if self.use_model_parallel: mp_rank = self._hcg.get_model_parallel_rank() else: mp_rank = 0 # mp rank 0 loads the data and broadcat it to others. data = self.data if self.use_model_parallel and (self.stage_id == 0 or self.stage_id == self.num_stages - 1): assert isinstance(data, (tuple, paddle.Tensor)) if isinstance(data, paddle.Tensor): paddle.distributed.broadcast( data, src=self._hcg.get_model_parallel_group_src_rank(), group=self._hcg.get_model_parallel_group()) else: data = [] for d in self.data: assert isinstance(d, paddle.Tensor) paddle.distributed.broadcast( d, src=self._hcg.get_model_parallel_group_src_rank(), group=self._hcg.get_model_parallel_group()) data.append(d) data = tuple(data) return data def _load_micro_batch(self, cache_id): inputs = self._get_data() if self.stage_id == 0: data = None #if isinstance(inputs[0], paddle.Tensor): if len(inputs) == 1: assert isinstance(inputs[0], paddle.Tensor) data = inputs[0].clone().detach() #data.stop_gradient = not is_float_tensor(data) data.stop_gradient = True else: assert isinstance(inputs, tuple) data = [] for d in inputs: assert isinstance(d, paddle.Tensor) i = d.clone().detach() #i.stop_gradient = not is_float_tensor(i) i.stop_gradient = True data.append(i) data = tuple(data) self.caches['inputs'][cache_id] = data if self.stage_id == self.num_stages - 1: labels = None #if isinstance(inputs[1], paddle.Tensor): if len(inputs) == 1: assert isinstance(inputs[0], paddle.Tensor) labels = inputs[0] elif isinstance(inputs, tuple): labels = [] for label in inputs: assert isinstance(label, paddle.Tensor) label = label.detach() labels.append(label) labels = tuple(labels) self.caches['labels'][cache_id] = labels def _send_meta(self, data, peer): """ % type (0: tensor, 1: tuple) % num_tensors if type=tuple foreach tensor: % ndims % shape """ if isinstance(data, paddle.Tensor): tensor_type = paddle.to_tensor([0]) paddle.distributed.send( tensor_type, peer, use_calc_stream=True, group=self.pp_group) dims = paddle.to_tensor(len(data.shape)) paddle.distributed.send( dims, peer, use_calc_stream=True, group=self.pp_group) shape = paddle.to_tensor(data.shape) paddle.distributed.send( shape, peer, use_calc_stream=True, group=self.pp_group) elif isinstance(data, tuple): tensor_type = paddle.to_tensor([1]) paddle.distributed.send( tensor_type, peer, use_calc_stream=True, group=self.pp_group) nums = paddle.to_tensor(len(data)) paddle.distributed.send( nums, peer, use_calc_stream=True, group=self.pp_group) for idx, d in enumerate(data): assert isinstance(d, paddle.Tensor) dims = paddle.to_tensor(len(d.shape)) paddle.distributed.send( dims, peer, use_calc_stream=True, group=self.pp_group) shape = paddle.to_tensor(d.shape) paddle.distributed.send( shape, peer, use_calc_stream=True, group=self.pp_group) def _recv_meta(self, peer): tensor_type = paddle.to_tensor([0]) paddle.distributed.recv( tensor_type, peer, use_calc_stream=True, group=self.pp_group) tensor_type = tensor_type.numpy()[0] if tensor_type == 0: dims = paddle.to_tensor([0]) paddle.distributed.recv( dims, peer, use_calc_stream=True, group=self.pp_group) dims = dims.numpy()[0] shape = paddle.to_tensor([0] * dims) paddle.distributed.recv( shape, peer, use_calc_stream=True, group=self.pp_group) shape = shape.numpy().tolist() return self._allocate_buffer( shape, dtype="float32", num_caches=1)[0] elif tensor_type == 1: num = paddle.to_tensor([0]) paddle.distributed.recv( num, peer, use_calc_stream=True, group=self.pp_group) num = num.numpy()[0] shapes = [] for i in range(num): dims = paddle.to_tensor([0]) paddle.distributed.recv( dims, peer, use_calc_stream=True, group=self.pp_group) dims = dims.numpy()[0] shape = paddle.to_tensor([0] * dims) paddle.distributed.recv( shape, peer, use_calc_stream=True, group=self.pp_group) shapes.append(shape.numpy().tolist()) dtypes = ["float32"] * len(shapes) caches = self._allocate_buffers(shapes, dtypes, num_caches=1)[0] caches = tuple(caches) return caches def _send_activations(self, cache_id): outputs = self.caches['outputs'][cache_id] if self.send_meta: self.send_meta = False self._send_meta(outputs, self.next_stage_id) if isinstance(outputs, paddle.Tensor): paddle.distributed.send( outputs, self.next_stage_id, use_calc_stream=True, group=self.pp_group) elif isinstance(outputs, tuple): for output in outputs: paddle.distributed.send( output, self.next_stage_id, use_calc_stream=True, group=self.pp_group) def _send_gradients(self, cache_id): inputs = self.caches['inputs'][cache_id] if isinstance(inputs, paddle.Tensor): assert inputs.grad is not None paddle.distributed.send( paddle.to_tensor(inputs.grad), self.prev_stage_id, use_calc_stream=True, group=self.pp_group) else: for idx, d in enumerate(inputs): # Skip tensors that will not produce a grad if not is_float_tensor(d): assert d.grad is None continue assert d.grad is not None paddle.distributed.send( d.grad, self.prev_stage_id, use_calc_stream=True, group=self.pp_group) self.caches['inputs'][cache_id] = None def _recv_activations(self, cache_id): inputs = None # Allocate the buffer if necessary if self.recv_cache is None: self.recv_cache = self._recv_meta(self.prev_stage_id) if isinstance(self.recv_cache, paddle.Tensor): paddle.distributed.recv( self.recv_cache, self.prev_stage_id, use_calc_stream=True, group=self.pp_group) inputs = self.recv_cache.clone().detach() inputs.stop_gradient = not is_float_tensor(inputs) else: assert isinstance(self.recv_cache, tuple) inputs = [None] * len(self.recv_cache) for idx, d in enumerate(self.recv_cache): assert isinstance(d, paddle.Tensor) paddle.distributed.recv( d, self.prev_stage_id, use_calc_stream=True, group=self.pp_group) inputs[idx] = d.clone().detach() inputs = tuple(inputs) for d in inputs: d.stop_gradient = not is_float_tensor(d) self.caches['inputs'][cache_id] = inputs def _recv_gradients(self, cache_id): outputs = self.caches['outputs'][cache_id] if self.grad_tensors is None: if isinstance(outputs, paddle.Tensor): s = list(outputs.shape) dtype = 'float16' if self.use_amp else "float32" self.grad_tensors = self._allocate_buffer( s, dtype, num_buffers=1)[0] else: sizes = [list(d.shape) for d in outputs if is_float_tensor(d)] dtypes = ['float16'] * len( sizes) if self.use_amp else ['float32'] * len(sizes) self.grad_tensors = self._allocate_buffers( sizes, dtypes, num_caches=1)[0] if isinstance(self.grad_tensors, paddle.Tensor): paddle.distributed.recv( self.grad_tensors, self.next_stage_id, use_calc_stream=True, group=self.pp_group) else: assert isinstance(outputs, tuple) for d in self.grad_tensors: paddle.distributed.recv( d, self.next_stage_id, use_calc_stream=True, group=self.pp_group) def _step(self): self._allreduce_grads() self.optimizer.step() self.optimizer.clear_gradients() def _clear_grads(self, inputs): if isinstance(inputs, paddle.Tensor): if inputs.grad is not None: inputs.clear_gradient() else: for d in inputs: if d.grad is not None: d.clear_gradient() def _allocate_zeros(self, shape, dtype): return paddle.zeros(shape, dtype) def _allocate_buffer(self, shape, dtype, num_caches=-1): caches = [] if num_caches == -1: num_caches = self.num_caches for count in range(num_caches): caches.append(self._allocate_zeros(shape, dtype)) return caches def _allocate_buffers(self, shapes, dtypes, num_caches=-1): caches = [] if num_caches == -1: num_caches = self.num_caches for count in range(num_caches): cache = [] for shape, dtype in zip(shapes, dtypes): cache.append(self._allocate_zeros(shape, dtype)) caches.append(cache) return caches def save_state_dict(self, model_path): state_dict = self._layers.state_dict() paddle.save(state_dict, model_path) def load_state_dict(self, model_path): state_dict = paddle.load(self.model_path) self._layers.set_state_dict(state_dict) _COMMAND_MAP = { utils.Optimize: _step, utils.Forward: _forward, utils.Backward: _backward, } def forward(self, *inputs, **kwargs): raise RuntimeError("Call train_batch for pipeline instead of forward.")